module Day6.Part1

import Data.List
import Data.Maybe
import Data.Vect
import Data.Vect.Quantifiers

import Utils

%default total

--- TYPES

data Direction = U | R | D | L

-- True if there is a wall
Map : Nat -> Nat -> Type
Map h w = Grid h w Bool

--- PARSING

parseDir : Char -> Maybe Direction
parseDir '^' = Just U
parseDir '>' = Just R
parseDir 'v' = Just D
parseDir '<' = Just L
parseDir _ = Nothing

parseInput : String -> Maybe (h ** w ** (Map h w, Pos h w, Direction))
parseInput inp = do
  (h ** w ** grid) <- parseGrid inp
  g <- findGuard grid
  pure (h ** w ** ((map . map) (=='#') grid, g))
  where
    findGuard : Vect h (Vect w Char) -> Maybe (Pos h w, Direction)
    findGuard [] = Nothing
    findGuard (row :: vs) =
      case any (\c => isItJust (parseDir c)) row of
        Yes p => Just ((FZ, cast $ anyToFin p), fromJust _ @{anyToFinCorrect p})
        No _ => map (mapFst (mapFst FS)) $ findGuard vs

--- DATA

predFin : Fin n -> Maybe (Fin n)
predFin FZ = Nothing
predFin (FS x) = Just $ weaken x

moveDir : {h, w : _} -> Direction -> Pos h w -> Maybe (Pos h w)
moveDir U (i, j) = (,j) <$> predFin i
moveDir R (i, j) = (i,) <$> strengthen (FS j)
moveDir D (i, j) = (,j) <$> strengthen (FS i)
moveDir L (i, j) = (i,) <$> predFin j

rotate : Direction -> Direction
rotate U = R
rotate R = D
rotate D = L
rotate L = U


covering
getRoute : {h, w : _} -> Direction -> Pos h w -> Map h w -> List (Pos h w)
getRoute = go []
  where
    covering
    go : List (Pos h w) -> Direction -> Pos h w -> Map h w -> List (Pos h w)
    go route dir pos map =
      let Just next = moveDir dir pos
            | Nothing => pos :: route
      in if index next map
         then go (pos :: route) (rotate dir) pos map
         else go (pos :: route) dir next map

--- SOLUTION

export covering
solution : String -> Maybe Nat
solution inp = do
  (h ** w ** (mp, pos, dir)) <- parseInput inp
  pure $ length $ nub $ getRoute dir pos mp