{-# LANGUAGE LambdaCase #-}

module Day8 where

import Control.Monad (replicateM)
import Data.Char
import Data.Map (Map)
import Data.Map.Lazy qualified as Map
import Text.ParserCombinators.ReadP

run :: IO ()
run = do
  (instructions, network) <- readInput
  putStrLn $ "Part 1 Answer: " <> either id show (part1 instructions network)
  putStrLn $ "Part 2 Answer: " <> either id show (part2 instructions network)

part1 :: [Instruction] -> Network -> Either String Int
part1 instructions network =
  followInstructionsUntil (== "ZZZ") network instructions "AAA"

part2 :: [Instruction] -> Network -> Either String Int
part2 instructions network =
  let isStartingNode = \case
        [_, _, 'A'] -> True
        _ -> False
      startingNodes = filter isStartingNode $ Map.keys network
      isDestinationNode = \case
        [_, _, 'Z'] -> True
        _ -> False
      steps = traverse (followInstructionsUntil isDestinationNode network instructions) startingNodes
   in foldr lcm 1 <$> steps

nextDest :: Network -> Instruction -> Node -> Either String Node
nextDest network i currentPosition =
  case Map.lookup currentPosition network of
    Nothing -> Left $ "We're lost, current position: " <> currentPosition
    Just (leftDest, rightDest) ->
      case i of
        R -> Right rightDest
        L -> Right leftDest

-- |
-- >>> network = Map.fromList [("AAA", ("BBB", "BBB")), ("BBB", ("AAA", "ZZZ")), ("ZZZ", ("ZZZ", "ZZZ"))]
-- >>> followInstructionsUntil (== "ZZZ") network [L,L,R] "AAA"
-- 6
followInstructionsUntil :: (Node -> Bool) -> Network -> [Instruction] -> Node -> Either String Int
followInstructionsUntil isDestination network initInstructions initPosition =
  go (cycle initInstructions) initPosition 0
  where
    go :: [Instruction] -> Node -> Int -> Either String Int
    go [] _ _ = Left "No instructions"
    go (i : is) currentPosition steps
      | isDestination currentPosition = Right steps
      | otherwise =
        case nextDest network i currentPosition of
          Left e -> Left e
          Right next -> go is next (steps + 1)

type Node = String

type Network = Map Node (Node, Node)

data Instruction = R | L
  deriving (Show, Read)

readInput :: IO ([Instruction], Network)
readInput = do
  input <- getContents
  case readP_to_S inputParser input of
    [] -> error "Parser failed"
    [(x, "")] -> pure x
    xs -> error $ "Parser failed: " <> show xs

instructionParser :: ReadP Instruction
instructionParser =
  read . (: []) <$> satisfy (\c -> c == 'R' || c == 'L')

nodeParser :: ReadP Node
nodeParser = do
  replicateM 3 $ satisfy isAlphaNum

nodeConnectionParser :: ReadP (Node, (Node, Node))
nodeConnectionParser = do
  src <- nodeParser
  _ <- string " = ("
  leftDest <- nodeParser
  _ <- string ", "
  rightDest <- nodeParser
  _ <- string ")"
  pure (src, (leftDest, rightDest))

networkParser :: ReadP Network
networkParser =
  Map.fromList <$> sepBy1 nodeConnectionParser (char '\n')

inputParser :: ReadP ([Instruction], Network)
inputParser = do
  instructions <- many1 instructionParser
  _ <- string "\n\n"
  network <- networkParser
  eof
  pure (instructions, network)