module Main where

import Data.List (intercalate)
import Data.Maybe (fromMaybe)
import Data.Set (Set)
import qualified Data.Set as S

data Direction = East | South deriving (Show, Eq, Ord)

type Coord = (Int, Int)

type Board = Set (Coord, Direction)

spaceFilled :: Board -> Coord -> Bool
spaceFilled b c = S.member (c, East) b || S.member (c, South) b

getDirection :: Board -> Coord -> Direction
getDirection b c
  | (c, East) `S.member` b = East
  | (c, South) `S.member` b = South
  | otherwise = undefined

wrapAround :: Board -> Coord -> Coord
wrapAround b (x, y) = (x', y')
  where
    x' = if x > mx then 0 else x
    y' = if y > my then 0 else y
    mx = 138
    my = 136

stepOne :: Board -> (Coord, Direction) -> (Coord, Direction)
stepOne b (c@(x, y), d) = (if canMove then newCoord else c, d)
  where
    newCoord = case d of
      East -> wrapAround b (x + 1, y)
      South -> wrapAround b (x, y + 1)
    canMove = not $ spaceFilled b newCoord

stepAll :: Board -> Board
stepAll ib = S.map (stepSouth eastStepped) eastStepped
  where
    eastStepped = S.map (stepEast ib) ib
    stepEast b x@(_, East) = stepOne b x
    stepEast b x = x
    stepSouth b x@(_, South) = stepOne b x
    stepSouth b x = x

stepTillStationary :: Board -> (Board, Int)
stepTillStationary ib = keepStepping ib 0
  where
    keepStepping b n
      | b == b' = (b, n + 1)
      | otherwise = keepStepping b' (n + 1)
      where
        b' = stepAll b

parseFile :: String -> Board
parseFile str = S.fromList $ getFilled str 0 0
  where
    getFilled [] x y = []
    getFilled ('.' : cs) x y = getFilled cs (x + 1) y
    getFilled ('\n' : cs) x y = getFilled cs 0 (y + 1)
    getFilled ('>' : cs) x y = ((x, y), East) : getFilled cs (x + 1) y
    getFilled ('v' : cs) x y = ((x, y), South) : getFilled cs (x + 1) y
    getFilled (_ : cs) x y = undefined

printBoard :: Board -> String
printBoard b = intercalate "\n" $ map printLine [0 .. 10]
  where
    printLine y = map (printCell y) [0 .. 10]
    printCell y x
      | spaceFilled b (x, y) = case getDirection b (x, y) of
        East -> '>'
        South -> 'v'
      | otherwise = '.'

main :: IO ()
main = do
  input <- readFile "./input"
  let parsed = parseFile input
  print $ fromMaybe 0 $ S.lookupMax $ S.map (fst . fst) parsed
  print $ fromMaybe 0 $ S.lookupMax $ S.map (snd . fst) parsed
  -- putStrLn $ printBoard parsed
  -- putStrLn "---"
  -- putStrLn $ printBoard $ foldr (\x b -> stepAll b) parsed [0 .. 0]

  print $ snd $ stepTillStationary parsed