{-# LANGUAGE MultiWayIf #-}
module Main where

import Control.Monad
import Data.Char
import Data.List
import Data.Maybe
import qualified Data.Map.Strict as Map

import Input
import IntCode


blockBy :: Int -> [a] -> [[a]]
blockBy n l =
    let (pre, post) = splitAt n l
    in if null post then [pre] else pre : blockBy n post

weightLimit :: Int
weightLimit = 20

numFunctions :: Int
numFunctions = 3

data Move = F Int | L | R deriving (Show, Eq)
type Path = [Move]

weight :: Path -> Int
weight l = length l - 1 + sum (map moveWeight l)

moveWeight :: Move -> Int
moveWeight (F n) = numberWeight n
moveWeight _ = 1

numberWeight :: Int -> Int
numberWeight n | n < 10 = 1
               | otherwise = 1 + numberWeight (n `quot` 10)

limitedPrefixes :: Path -> [(Path, Path)]
limitedPrefixes moves =
    let prefws = tail (scanl (+) (-1) (map (succ . moveWeight) moves))
        maxlen = length (takeWhile (<= weightLimit) prefws)
    in [splitAt i moves | i <- [1..maxlen]]

data State = State { pathFuncs :: [Path], pathRemain :: Path } deriving (Show)

compress' :: State -> [([Int], State)]
compress' state@(State _ []) = return ([], state)
compress' (State funcs remain) =
    let match idx func = do
            guard $ take (length func) remain == func
            return ([idx], State funcs (drop (length func) remain))
        matches = concatMap (uncurry match) (zip [0..] funcs)
        newFunc prefix rest = do
            let newidx = length funcs
            guard $ newidx < numFunctions
            return ([newidx], State (funcs ++ [prefix]) rest)
        newFuncs = concatMap (uncurry newFunc) (limitedPrefixes remain)
        options = matches ++ newFuncs
        developOption (yet, state) = do
            (yet', state') <- compress' state
            return (yet ++ yet', state')
    in concatMap developOption options

compress :: Path -> [([Int], [Path])]
compress path = do
    (calls, State funcs _) <- compress' (State [] path)
    let mainWeight = length calls - 1 + sum (map numberWeight calls)
    guard $ mainWeight <= weightLimit
    let restored = concatMap (funcs !!) calls
    when (restored /= path) (error "kaas!")
    return (calls, funcs)

type Pos = (Int, Int)
type Dir = (Int, Int)

add :: Pos -> Dir -> Pos
add (x, y) (dx, dy) = (x + dx, y + dy)

rotR, rotL :: Dir -> Dir
rotR (dx, dy) = (-dy, dx)
rotL (dx, dy) = (dy, -dx)

fourDirections :: [Dir]
fourDirections = [(-1,0), (1,0), (0,-1), (0,1)]

findPath :: Map.Map Pos Char -> Pos -> Dir -> Path
findPath bd pos dir =
    let cands = [([F 1], add pos dir, dir),
                 ([R, F 1], add pos (rotR dir), rotR dir),
                 ([L, F 1], add pos (rotL dir), rotL dir)]
        cands' = filter (\(_, p, _) -> Map.lookup p bd == Just '#') cands
    in case cands' of
           (moves, pos', dir') : _ -> moves ++ findPath bd pos' dir'
           _ -> []

simplifyPath :: Path -> Path
simplifyPath (F n : F m : rest) = simplifyPath (F (n + m) : rest)
simplifyPath (m : rest) = m : simplifyPath rest
simplifyPath [] = []

main :: IO ()
main = do
    program <- parse . head <$> getInput 17
    let bd = let output = map (chr . fromIntegral) (snd (run program []))
                 nlidx = fromJust (findIndex (== '\n') output)
             in blockBy nlidx (filter (/= '\n') output)
        w = length (head bd)
        h = length bd
        bdmap = Map.fromList [((x, y), c) | (y, row) <- zip [0..] bd, (x, c) <- zip [0..] row]

    -- mapM_ print bd

    let inters = [(x, y)
                 | y <- [1..h-2]
                 , x <- [1..w-2]
                 , all (== '#') [bdmap Map.! (x+dx, y+dy) | (dx, dy) <- (0,0) : fourDirections]]
    -- print inters

    -- putStr (unlines [concat [if (x, y) `elem` inters
    --                              then "\x1B[41;1m" ++ cell : "\x1B[0m"
    --                              else [cell]
    --                         | (x, cell) <- zip [0..] row]
    --                 | (y, row) <- zip [0..] bd])

    print (sum (map (uncurry (*)) inters))

    let startpos = fst (head (filter ((`elem` "^>v<") . snd) (Map.assocs bdmap)))
        startdir = fromJust (lookup (bdmap Map.! startpos) (zip "<>^v" fourDirections))
        path = simplifyPath (findPath bdmap startpos startdir)
        compressed = compress path
        (mainroutine, funcs) = head compressed
    -- mapM_ print compressed

    let program' = 2 : tail program
        moveInput = map (fromIntegral . ord) . unlines . map (intercalate ",") $
                        [map (pure . ("ABC" !!)) mainroutine]
                        ++ [[case m of {F n -> show n; R -> "R"; L -> "L"} | m <- func]
                           | func <- funcs]
                        ++ [["n"]]

    print (last (snd (run program' moveInput)))