{-# LANGUAGE BangPatterns #-}
module Main where

import Data.List
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Data.Maybe

import Input
import IntCode


sortUniqBy :: Ord b => (a -> b) -> [a] -> [a]
sortUniqBy f = uniq . sortOn f
  where uniq (x:y:zs) | f x == f y = uniq (y:zs)
                      | otherwise = x : uniq (y:zs)
        uniq l = l

data Cell = Open | Wall | Unit deriving (Show, Eq)
type Pos = (Int, Int)

step :: Pos -> Int -> Pos
step (x, y) dir = case dir of
                      1 -> (x, y - 1)
                      2 -> (x, y + 1)
                      3 -> (x - 1, y)
                      4 -> (x + 1, y)
                      _ -> error "Invalid direction in 'step'"

showMap :: Map.Map Pos Cell -> Pos -> String
showMap mp droid =
    let keys = Map.keys mp
        (minx, maxx) = (minimum (map fst keys), maximum (map fst keys))
        (miny, maxy) = (minimum (map snd keys), maximum (map snd keys))
    in unlines [[if (x, y) `elem` [droid, (0, 0)]
                     then if (x, y) == (0, 0) then 'S' else 'D'
                     else case Map.lookup (x, y) mp of
                              Nothing -> ' '
                              Just Open -> '.'
                              Just Wall -> '#'
                              Just Unit -> '@'
                | x <- [minx..maxx]]
               | y <- [miny..maxy]]

performStep :: Continuation -> Pos -> Map.Map Pos Cell -> Int -> (Continuation, Pos, Map.Map Pos Cell)
performStep cont pos mp dir =
    let (cont', output) = case runContinue cont [fromIntegral dir] of
                              Left (c, out) -> (c, out)
                              Right (_, out) -> (undefined, out)
    in case output of
           [0] -> (cont', pos         , Map.insert (step pos dir) Wall mp)
           [1] -> (cont', step pos dir, Map.insert (step pos dir) Open mp)
           [2] -> (cont', step pos dir, Map.insert (step pos dir) Unit mp)
           _   -> error $ "Invalid output " ++ show output ++ " in performStep"

shortestPath :: Map.Map Pos Cell -> Pos -> Maybe Cell -> Either Int [Int]
shortestPath mp initPos target = reverse <$> search 0 (Set.singleton initPos) [(initPos, [])]
  where
    search :: Int -> Set.Set Pos -> [(Pos, [Int])] -> Either Int [Int]
    search !depth _ [] = Left (depth - 1)
    search !depth seen border =
        let boundary =
                sortUniqBy (fst . fst)
                    [((newPos, dir : path), mcell)
                    | (pos, path) <- border
                    , dir <- [1..4]
                    , let newPos = step pos dir
                    , newPos `Set.notMember` seen
                    , let mcell = Map.lookup newPos mp
                    , mcell /= Just Wall]
            newBorder = map fst boundary
        in case find ((== target) . snd) boundary of
               Just ((_, newPath), _) -> Right newPath
               Nothing -> search (depth + 1) (seen <> Set.fromList (map fst newBorder)) newBorder

determineMap :: Continuation -> Pos -> Map.Map Pos Cell -> Map.Map Pos Cell
determineMap cont pos mp =
    case shortestPath mp pos Nothing of
        Right path ->
            let (cont', pos', mp') =
                    foldl (\(c, p, m) dir -> performStep c p m dir) (cont, pos, mp) path
            in determineMap cont' pos' mp'
        Left _ ->
            mp

main :: IO ()
main = do
    program <- parse . head <$> getInput 15
    let initCont = initialContinuation program

    let mp = determineMap initCont (0, 0) (Map.singleton (0, 0) Open)
    let Right path = shortestPath mp (0, 0) (Just Unit)
    print (length path)

    let unitpos = fst (fromJust (find ((== Unit) . snd) (Map.assocs mp)))

    let Left maxdist = shortestPath mp unitpos Nothing
    print maxdist