module Day6

import Prelude
import Node
import Aoc
import SortedMap

Point : U
Point = Int × Int

instance Eq Point where
  (a,b) == (c,d) = a == c && b == d

instance Ord Point where
  (a,b) < (c,d) = a < c || a == c && b < d

Grid : U
Grid = SortedMap Point Char

loadData : String → Grid
loadData text = go (unpack text) 0 0 EmptyMap
  where
    go : List Char → Int → Int → SortedMap Point Char → SortedMap Point Char
    go Nil r c map = map
    go ('\n' :: cs) r c map = go cs (r + 1) 0 map
    go (x :: xs) r c map = go xs r (c + 1) $ updateMap (r,c) x map

data Dir : U where North East South West : Dir

instance Show Dir where
  show North = "N"
  show East = "E"
  show South = "S"
  show West = "W"

instance Ord Dir where
  a < b = show a < show b

instance Eq (Point × Dir) where
  (a,b) == (c,d) = a == c && show b == show d

instance Ord (Point × Dir) where
  (a,b) < (c,d) =
    if a < c then True
    else if a /= c then False
    else b < d

Done : U
Done = SortedMap (Point × Dir) Unit

turn : Dir → Dir
turn North = East
turn East = South
turn South = West
turn West = North

instance Cast Dir Char where
  cast North = '^'
  cast East = '>'
  cast South = 'v'
  cast West = '<'

step : Dir → Point → Point
step North (r, c) = (r - 1, c)
step East (r, c) = (r, c + 1)
step South (r, c) = (r + 1, c)
step West (r, c) = (r, c - 1)

bad : Point → Bool
bad (x,y) = x < 0 || y < 0

-- third is
walk : Dir → Point → Grid → Grid
walk dir pos grid =
  let grid = updateMap pos 'X' grid in
  let pos' = step dir pos in
  case lookupMap pos' grid of
    Just (_, '#') => walk (turn dir) pos grid
    Nothing => grid
    _   => walk dir pos' grid

checkLoop : Grid → Done → Dir → Point → Bool
checkLoop grid done dir pos =
  let (Nothing) = lookupMap (pos,dir) done | _ => True in
  let done = updateMap (pos, dir) MkUnit done
      pos' = step dir pos
    in case lookupMap pos' grid of
        Nothing => False
        Just (_, '#') => checkLoop grid done (turn dir) pos
        Just _ => checkLoop grid done dir pos'

part2 : Dir → Point → Grid → Done → List Point → List Point
part2 dir pos grid done sol =
  let done = updateMap (pos, dir) MkUnit done
      grid = updateMap pos 'X' grid
      turnDir = turn dir
      turnPos = step turnDir pos
      pos' = step dir pos in
  case lookupMap pos' grid of
    Nothing => sol
    Just (_, '#') => part2 (turn dir) pos grid done sol
    Just (_, 'X') => part2 dir pos' grid done sol
    Just (_, '.') => if checkLoop (updateMap pos' '#' grid) done turnDir pos
      then part2 dir pos' grid done (pos' :: sol)
      else part2 dir pos' grid done sol
    Just x => part2 (trace ("WAT " ++ debugStr x) dir) pos' grid done sol

lookupV : ∀ a. Char → List (a × Char) → Maybe a
lookupV _ Nil = Nothing
lookupV needle ((k,v) :: rest) =
  if v == needle then Just k else lookupV needle rest

run : String -> IO Unit
run fn = do
  putStrLn fn
  text <- readFile fn
  let grid = loadData text
  let (Just pos) = lookupV '^' (toList grid) | _ => putStrLn "no guard"
  let grid' = walk North pos grid
  let xs = filter (\ x => 'X' == snd x) $ toList grid'
  let part1 = length xs
  putStrLn $ "part1 " ++ show part1

  let cands = part2 North pos grid EmptyMap Nil
  -- debugLog $ length cands -- turns out nub isn't needed for these cases, but we'll leave it in
  putStrLn $ "part2 " ++ show (length $ ordNub cands)
  printLn $ length $ toList grid

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