{-# LANGUAGE LambdaCase, MultiWayIf #-}
module Main where

import Data.Functor
import Data.List
import Data.Matrix

data Dir where
  N,E,W,S :: Dir
  deriving Eq

data Pos where
  Marked   :: Pos
  Unmarked :: Pos
  Guard    :: Dir -> Pos
  Crate    :: Pos
  deriving Eq

instance Show Pos where
  show = \case
              Unmarked -> "."
              Marked   -> "X"
              Guard N  -> "^"
              Guard E  -> ">"
              Guard W  -> "<"
              Guard S  -> "v"
              Crate    -> "#"

parse :: String -> Matrix Pos
parse = fromLists . (((\case
              '.' -> Unmarked
              'X' -> Marked
              '^' -> Guard N
              '>' -> Guard E
              '<' -> Guard W
              'v' -> Guard S
              '#' -> Crate
             ) <$>) <$>) . lines

guardPos :: Matrix Pos -> (Int, Int, Dir)
guardPos s = let
        s' = toLists s
        Just y = findIndex (any (\case
                                Guard _ -> True
                                _ -> False)) s'
        Just x = findIndex (\case
                        Guard _ -> True
                        _ -> False) (s' !! y)
        d = (\case
                Guard a -> a
                _ -> error ""
            ) (getElem (y+1) (x+1) s)
             in
               (y+1,x+1,d)

move :: Matrix Pos -> (Int, Int, Dir) -> Matrix Pos
move s (y,x,d) = let (y',x',d') = case d of
                                 N -> (y-1,x,E)
                                 E -> (y,x+1,S)
                                 W -> (y,x-1,N)
                                 S -> (y+1,x,W)
                     t = getElem y' x' s
                 in
                   case t of
                     Crate -> setElem (Guard d') (y,x) s
                     _ -> setElem (Guard d) (y',x') $ setElem (Marked) (y,x) s


traversal :: Matrix Pos -> Matrix Pos
traversal s = let i@(y,x,d) = guardPos s in
  if | y == 1       && d == N -> s
     | y == nrows s && d == S -> s
     | x == 1       && d == W -> s
     | x == ncols s && d == E -> s
     | otherwise -> traversal $ move s i

solve1 :: Matrix Pos -> Int
solve1 = sum . (length . filter (\case
                                  Marked -> True
                                  Guard _ -> True
                                  _ -> False
                              ) <$>) . toLists . traversal

solve2 :: Matrix Pos -> Int
solve2 = undefined

main :: IO ()
main = readFile "inputs/6" <&> parse >>= \i ->
       print (solve1 i)
       -- >> print (solve2 i)