import Control.Lens
import qualified Data.HashSet as H
import Data.List
import Data.List.Split
import Linear.V2

type Cave = H.HashSet (V2 Int)

main :: IO ()
main = do
  input <- map (map (splitOn ",") . splitOn " -> ") . lines <$> readFile "input"
  let cave = H.fromList $ concatMap toV2 input
  let getMaxY = maximum . map (^. _y) . H.toList
  let f x = length (H.toList x \\ H.toList cave)
  let part1 = f $ day14 (getMaxY cave) True cave
  let part2 = f $ day14 (getMaxY cave) False cave
  putStrLn $ "Part 1: " ++ show part1 ++ "\nPart 2: " ++ show part2

day14 :: Int -> Bool -> Cave -> Cave
day14 maxY mode cave =
  if cave == newCave
    then cave
    else day14 maxY mode newCave
  where
    newCave = fallDown maxY mode (V2 500 0) cave

toV2 :: [[String]] -> [V2 Int]
toV2 ([x1, y1]:[x2, y2]:xs) =
  [V2 i j | i <- f [x1, x2], j <- f [y1, y2]] ++ toV2 ([x2, y2] : xs)
  where
    f = (\[x, y] -> [x .. y]) . sort . map read
toV2 _ = []

fallDown :: Int -> Bool -> V2 Int -> Cave -> Cave
fallDown maxY mode start cave
  | not mode && start ^. _y == maxY + 1 = H.insert start cave
  | mode && start ^. _y >= maxY = cave
  | f (V2 0 1) cave = fallDown maxY mode (start + V2 0 1) cave
  | f (V2 (-1) 1) cave = fallDown maxY mode (start + V2 (-1) 1) cave
  | f (V2 1 1) cave = fallDown maxY mode (start + V2 1 1) cave
  | otherwise = H.insert start cave
  where
    f x = not . H.member (start + x)