module Days.Day14 where

import AOCUtil
import Data.Bool
import Data.Char
import Data.Either
import Data.IntSet qualified as Data.Map
import Data.List
import Data.Maybe
import Data.Set qualified
import Debug.Trace
import GHC.Arr
import GHC.Utils.Misc

runA :: IO ()
runA = interactF "data/day14.txt" solve

runB :: IO ()
runB = interactF "data/day14.txt" solveB

type Pos = (Int, Int)

type Path = [Pos]

type CMap = Data.Set.Set Pos -- occupied

solve :: String -> String
solve = show . fst . head . dropWhile (isRight . snd) . zip [0 ..] . tail . iterate (insertSand . fromEither) . Right . readMap . readPaths

solveB :: String -> String
solveB s = show . fst . head . dropWhile (not . Data.Set.member source . snd) . zip [0 ..] . iterate (insertSand2 fl) $ m
  where
    m = readMap . readPaths $ s
    fl = floorY m

insertSand :: CMap -> Either CMap CMap -- Left ^= Could not insert, Right ^=could
insertSand set = maybe (Left set) (Right . flip Data.Set.insert set) $ newPos set source

insertSand2 :: Int -> CMap -> CMap -- Left ^= Could not insert, Right ^=could
insertSand2 fl set = flip Data.Set.insert set . newPos2 fl set $ source

newPos :: CMap -> Pos -> Maybe Pos
newPos set (x, y)
  | y > maxY set = Nothing
  | not $ Data.Set.member (x, y + 1) set = newPos set (x, y + 1)
  | not $ Data.Set.member (x - 1, y + 1) set = newPos set (x - 1, y + 1)
  | not $ Data.Set.member (x + 1, y + 1) set = newPos set (x + 1, y + 1)
  | otherwise = Just (x, y)

newPos2 :: Int -> CMap -> Pos -> Pos
newPos2 fl set (x, y)
  | y + 1 == fl = (x, y) -- floor reached
  | not $ Data.Set.member (x, y + 1) set = newPos2 fl set (x, y + 1)
  | not $ Data.Set.member (x - 1, y + 1) set = newPos2 fl set (x - 1, y + 1)
  | not $ Data.Set.member (x + 1, y + 1) set = newPos2 fl set (x + 1, y + 1)
  | otherwise = (x, y)

readPaths :: String -> [Path]
readPaths = map (map (first2 . map read . splitOn ',') . filter (/= "->") . words) . lines

expandPath :: Path -> [Pos]
expandPath [] = []
expandPath [p] = [p]
expandPath ((x1, y1) : (x2, y2) : ps) = [(x, y) | x <- [min x1 x2 .. max x1 x2], y <- [min y1 y2 .. max y1 y2]] ++ expandPath ((x2, y2) : ps)

readMap :: [Path] -> CMap
readMap = Data.Set.fromList . concatMap expandPath

source :: Pos
source = (500, 0)

maxY :: CMap -> Int
maxY = maximum . map snd . Data.Set.toList

floorY :: CMap -> Int
floorY = (+ 2) . maxY
