module Main (main) where

import Control.Parallel.Strategies (parList, rdeepseq, using)
import Data.List (nub)
import Prelude hiding (Left, Right)

-- | Part 1
data Tile = Empty | Blockage | Guard deriving (Show, Eq)

newtype Grid = Grid [[Tile]] deriving (Show)

data Direction = Up | Down | Left | Right deriving (Show, Eq)

data WalkResult = Forward (Int, Int) | BlockEncountered Direction | OutOfBounds

tileFromString :: Char -> Tile
tileFromString '.' = Empty
tileFromString '#' = Blockage
tileFromString '^' = Guard

guardRow :: [Char] -> [Tile]
guardRow = map tileFromString

guardGrid :: String -> Grid
guardGrid = Grid . map guardRow . lines

findStart :: Grid -> (Int, Int)
findStart grid = findStart' grid 0 0
  where
    findStart' :: Grid -> Int -> Int -> (Int, Int)
    findStart' (Grid grid') row col
      | col >= length (grid' !! row) = findStart' (Grid grid') (row + 1) 0
      | (grid' !! row) !! col == Guard = (row, col)
      | otherwise = findStart' (Grid grid') row (col + 1)

move :: Direction -> (Int, Int) -> (Int, Int)
move Up (row, col) = (row - 1, col)
move Down (row, col) = (row + 1, col)
move Left (row, col) = (row, col - 1)
move Right (row, col) = (row, col + 1)

turnRight :: Direction -> Direction
turnRight Up = Right
turnRight Right = Down
turnRight Down = Left
turnRight Left = Up

outOfBounds :: Grid -> (Int, Int) -> Bool
outOfBounds (Grid g) (row, col) = row < 0 || row >= length g || col < 0 || col >= length (g !! row)

stepGridWith :: Grid -> Direction -> (Int, Int) -> WalkResult
stepGridWith (Grid grid) direction pos
  | outOfBounds (Grid grid) pos = OutOfBounds
  | outOfBounds (Grid grid) (move direction pos) = Forward (move direction pos)
  | otherwise =
      let (r', c') = move direction pos
          newDirection = turnRight direction
       in case (grid !! r') !! c' of
            Blockage -> BlockEncountered newDirection
            _ -> Forward (r', c')

walkGrid :: Grid -> Direction -> (Int, Int) -> [(Int, Int)]
walkGrid grid direction pos =
  case stepGridWith grid direction pos of
    OutOfBounds -> []
    Forward newPos -> pos : walkGrid grid direction newPos
    BlockEncountered newDirection -> walkGrid grid newDirection pos

part1 :: String -> IO ()
part1 s =
  let grid = guardGrid s
      start = findStart grid
   in print $ length $ nub $ walkGrid grid Up start

-- | Part 2
isLoop :: Direction -> (Int, Int) -> Grid -> Bool
isLoop direction pos grid = isLoop' grid direction pos []
  where
    isLoop' :: Grid -> Direction -> (Int, Int) -> [((Int, Int), Direction)] -> Bool
    isLoop' grid' direction' pos' visited
      | (pos', direction') `elem` visited = True
      | otherwise =
          case stepGridWith grid' direction' pos' of
            OutOfBounds -> False
            Forward newPos -> isLoop' grid' direction' newPos ((pos', direction') : visited)
            BlockEncountered newDirection -> isLoop' grid' newDirection pos' visited

blockPos :: Grid -> (Int, Int) -> Grid
blockPos (Grid g) (row, col) =
  let (g', gs) = splitAt row g
   in Grid $ g' ++ [blockPosInCol' (g !! row) col] ++ drop 1 gs
  where
    blockPosInCol' :: [Tile] -> Int -> [Tile]
    blockPosInCol' row' col' =
      let (r, rs) = splitAt col' row'
       in r ++ [Blockage] ++ drop 1 rs

countLoopers :: Grid -> (Int, Int) -> Int
countLoopers (Grid grid) startPos =
  let guardPath = nub $ walkGrid (Grid grid) Up startPos
      emptyPath = filter (\(row, col) -> (grid !! row) !! col == Empty) guardPath
      loops = map countLooper emptyPath `using` parList rdeepseq
   in sum loops
  where
    countLooper :: (Int, Int) -> Int
    countLooper (row, col)
      | isLoop Up startPos (blockPos (Grid grid) (row, col)) = 1
      | otherwise = 0

part2 :: String -> IO ()
part2 s =
  let grid = guardGrid s
      start = findStart grid
   in print $ countLoopers grid start

-- | Main
main :: IO ()
main = readFile "resources/day6.txt" >>= part2