{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Days.D12 where

import Common
import Control.Monad (filterM, foldM)
import Control.Monad.Extra (concatMapM)
import Control.Monad.ST (ST, runST)
import Data.ByteString (ByteString)
import Data.ByteString.Char8 qualified as BS
import Data.ByteString.Unsafe qualified as BS
import Data.Char (ord)
import Data.Functor ((<&>))
import Data.List (findIndex)
import Data.Maybe (fromMaybe)
import Data.Vector.Storable (Vector)
import Data.Vector.Storable qualified as Vec
import Data.Vector.Storable.Mutable qualified as MutVec
import Foreign.C (castCCharToChar)
import Foreign.ForeignPtr (newForeignPtr_)
import Foreign.Storable (Storable)
import System.IO.Unsafe (unsafePerformIO)

data Point = P !Int !Int
  deriving (Show, Eq, Ord)

data Grid t = Grid Int Int (Vector t)

data MGridSet s = MGridSet Int Int (MutVec.MVector s Bool)

neigh4 :: Point -> [Point]
neigh4 (P x y) =
  [ P (x - 1) y,
    P (x + 1) y,
    P x (y - 1),
    P x (y + 1)
  ]

(@) :: (Storable t) => Grid t -> Point -> t
(@) (Grid _ w m) (P x y) = m Vec.! (w * x + y)

get :: MGridSet s -> Point -> ST s Bool
get (MGridSet _ w m) (P x y) = MutVec.unsafeRead m (w * x + y)

set :: MGridSet s -> Point -> ST s Bool
set (MGridSet _ w m) (P x y) = MutVec.unsafeExchange m (w * x + y) True

dims :: Grid t -> (Int, Int)
dims (Grid h w _) = (h, w)

mkGridSet :: (Int, Int) -> ST s (MGridSet s)
mkGridSet (h, w) = MGridSet h w <$> MutVec.replicate (h * w) False

contains :: (Storable t) => Grid t -> Point -> Bool
contains (Grid h w _) (P x y) =
  0 <= x
    && 0 <= y
    && x < h
    && y < w

bfs :: (Storable t) => Grid t -> (Point -> Point -> Bool) -> Point -> [[Point]]
bfs g f start =
  runST $
    mkGridSet (dims g)
      >>= (\mgs -> set mgs start >> return mgs)
      >>= flip go [start]
  where
    go :: MGridSet s -> [Point] -> ST s [[Point]]
    go seen l
      | null l = pure []
      | otherwise = (l :) <$> (advance seen l >>= go seen)

    -- advance :: MGridSet s -> [Point] -> ST s [Point]
    -- advance seen points =
    --   points
    --     & concatMap (\point -> [p | p <- neigh4 point, g `contains` p, f point p])
    --     & foldM
    --       (\l p -> set seen p <&> (\isNew -> if isNew then l else p : l))
    --       []

    -- Uglier, but unfortunately significantly faster
    advance :: MGridSet s -> [Point] -> ST s [Point]
    advance seen =
      concatMapM
        ( \point ->
            filterM (fmap not . get seen) [p | p <- neigh4 point, g `contains` p]
              >>= foldM
                (\l p -> set seen p <&> \wasSeen -> if wasSeen then l else p : l)
                []
                . filter (f point)
        )

find :: (Storable t, Eq t) => Grid t -> t -> Maybe Point
find (Grid _ w m) v = (\pos -> P (pos `div` w) (pos `mod` w)) <$> Vec.elemIndex v m

elevation :: Char -> Int
elevation 'S' = ord 'a'
elevation 'E' = ord 'z'
elevation c = ord c

needGear :: Grid Char -> Point -> Point -> Bool
needGear grid from to = elevation (grid @ to) <= elevation (grid @ from) + 1

part1 :: Grid Char -> Maybe Int
part1 grid = do
  start <- find grid 'S'
  end <- find grid 'E'
  let bfsResult = bfs grid (needGear grid) start
  findIndex (elem end) bfsResult

part2 :: Grid Char -> Maybe Int
part2 grid =
  find grid 'E'
    >>= findIndex (any (\p -> elevation (grid @ p) == elevation 'a'))
      . bfs grid (flip $ needGear grid)

gridify :: ByteString -> Grid Char
gridify bytes =
  let numBytes = BS.length bytes
      w = fromMaybe numBytes (BS.elemIndex '\n' bytes)
      h = (numBytes + 1) `div` (w + 1)
      joinedBytes = BS.filter (/= '\n') bytes
   in Grid
        h
        w
        ( unsafePerformIO $
            BS.unsafeUseAsCStringLen
              joinedBytes
              ( \(mem, len) ->
                  newForeignPtr_ mem
                    >>= Vec.unsafeFreeze . flip MutVec.unsafeFromForeignPtr0 len
                    <&> Vec.map castCCharToChar
              )
        )

day :: Day
day =
  ByteStringDay
    ( \text ->
        let grid = gridify text
         in Right (Answer <$> part1 grid, Answer <$> part2 grid)
    )