#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [pretty-simple attoparsec arithmoi])"
#! nix-shell -i "ghcid -c 'ghci' -T main"

{-# OPTIONS_GHC -Wall -Wincomplete-uni-patterns #-}
{-# OPTIONS_GHC -Wno-unused-top-binds -Wno-unused-imports #-}
{-# OPTIONS_GHC -Wno-unused-matches -Wno-type-defaults #-}
{-# OPTIONS_GHC -Wno-unused-local-binds #-}
{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative
import Data.Attoparsec.Text (Parser)
import Data.Text (Text)
import Debug.Trace (trace)
import Text.Pretty.Simple
import Data.Maybe (fromMaybe,catMaybes)
import Data.List (find,sortOn)
import qualified Data.Attoparsec.Text as A
import qualified Data.Text as T
import Math.NumberTheory.Primes

exampleData :: String
exampleData = "939\n7,13,x,x,59,x,31,19"

numOrXParser :: Parser (Maybe Int)
numOrXParser = (Just <$> A.decimal) <|> ("x" *> pure Nothing)

inputParser :: Parser (Int,[Int])
inputParser = do
  n <- A.decimal
  A.skipSpace
  xs <- numOrXParser `A.sepBy` ","
  pure (n,catMaybes $ xs)

parseInput :: String -> Either String (Int,[Int])
parseInput = (A.parseOnly inputParser) . T.pack

solvePart1 :: String -> Either String Int
solvePart1 str = do
  (n,xs) <- parseInput str
  let (bus, time) = head
                  $ sortOn (snd)
                  $ map (\x -> (x,fromMaybe (-1) $ find (> n) [0,x..])) xs
  pure $ bus * (time - n)

inputParser2 :: Parser (Int,[Maybe Int])
inputParser2 = do
  n <- A.decimal
  A.skipSpace
  xs <- numOrXParser `A.sepBy` ","
  pure (n,xs)

parseInput2 :: String -> Either String (Int,[Maybe Int])
parseInput2 = (A.parseOnly inputParser2) . T.pack

solvePart2 :: String -> Either String [(Int,Int)]
solvePart2 str = do
  (_,xs) <- parseInput2 str
  let startAndIds = catMaybes $ sequence <$> zip [0..] xs
  pure startAndIds

main :: IO ()
main = do
  putStrLn ":: Test"
  pPrint $ A.parseOnly inputParser $ T.pack exampleData
  pPrint $ take 3 ((\n -> [1,n..]) 59)
  putStrLn ":: Day 13 - Part 1"
  input <- readFile "day13/input"
  pPrint $ solvePart1 exampleData
  pPrint $ solvePart1 input
  putStrLn ":: Test"
  print $ unPrime . fst <$> factorise (1068781::Integer)
  putStrLn ":: Day 13 - Part 2"
  print $ solvePart2 exampleData