module Day5

import Prelude
import Node
import Aoc
import SortedMap

data Prob : U where
  MkProb : List (Int × Int) -> List (List Int) → Prob

parseRule : String → Maybe (Int × Int)
parseRule txt =
  let (a :: b :: Nil) = nums' "|" txt  | _ => Nothing
  in Just (a,b)

parse : String → Maybe Prob
parse text = do
  let (a :: b :: Nil) = split (trim text) "\n\n" | pts => Nothing
  rules <- traverse parseRule $ split a "\n"
  let updates = map (nums' ",") $ split b "\n"
  Just $ MkProb rules updates

RuleMap : U
RuleMap = SortedMap Int (List Int)

getDisallowed : Int → RuleMap → List Int
getDisallowed key rmap = fromMaybe Nil (map snd $ lookupMap key rmap)

mkRuleMap : List (Int × Int) -> RuleMap
mkRuleMap rules = foldl go EmptyMap rules
  where
    go : RuleMap → Int × Int → RuleMap
    go rmap (b,a) = updateMap a (b :: getDisallowed a rmap) rmap

scan : RuleMap → List Int -> List Int -> Bool
scan rmap interdit Nil = True
scan rmap interdit (x :: xs) =
  if elem x interdit then False
  else scan rmap (getDisallowed x rmap ++ interdit) xs

fix : RuleMap → List Int → List Int
fix rmap Nil = Nil
fix rmap (x :: xs) =
  let interdit = getDisallowed x rmap in
  let (prefix,rest) = partition (flip elem interdit) xs
  in case prefix of
    Nil => x :: fix rmap rest
    ys => fix rmap (ys ++ x :: rest)

middle : List Int -> Int
middle xs = go xs xs
  where
    go : List Int → List Int → Int
    go (x :: xs) (_ :: _ :: ys) = go xs ys
    go (x :: xs) (_ :: ys) = x
    go _ _ = 0

run : String -> IO Unit
run fn = do
  putStrLn fn
  text <- readFile fn
  let (Just prob) = parse text | _ => putStrLn "Parse Error"
  let (MkProb rules things) = prob
  let rmap = mkRuleMap rules
  let good = filter (scan rmap Nil) things
  let part1 = foldl _+_ 0 $ map middle good
  let bad = filter (not ∘ scan rmap Nil) things
  putStrLn $ "part1 " ++ show part1
  let fixed = map (fix rmap) bad
  printLn $ length bad
  let part2 = foldl _+_ 0 $ map middle fixed
  putStrLn $ "part2 " ++ show part2

main : IO Unit
main = do
  run "aoc2024/day5/eg.txt"
  run "aoc2024/day5/input.txt"