{-# LANGUAGE LambdaCase #-}

import Control.Arrow
import Data.Char
import Data.List.Split
import qualified Data.Map as M

type LR a = (a, a) -> a

main :: IO ()
main = do
  (directions, input) <- parse . lines <$> readFile "input"
  print $ day08a directions input "AAA"
  let starts = filter ((== 'A') . last) (M.keys input)
  print $ day08b directions input starts

parse :: [String] -> ([LR String], M.Map String (String, String))
parse =
  cycle .
  map
    (\case
       'L' -> fst
       'R' -> snd
       _ -> error "Malformed directions input") .
  head &&&
  M.fromList . map parseLine . tail . tail

parseLine :: String -> (String, (String, String))
parseLine s = (a, (b, c))
  where
    [a, b, c] = wordsBy (not . isAlphaNum) s

day08a :: [LR String] -> M.Map String (String, String) -> String -> Int
day08a _ _ [_, _, 'Z'] = 0
day08a (lr:lrs) m pos = 1 + day08a lrs m (lr (m M.! pos))

day08b :: [LR String] -> M.Map String (String, String) -> [String] -> Int
day08b lrs m starts = foldr1 lcm $ map (day08a lrs m) starts