{-# LANGUAGE MultiWayIf #-}
module Main where

import Control.Concurrent
import Control.Monad
import Data.List
import qualified Data.Map.Strict as Map
import Data.Maybe
import System.IO

import Input
import IntCode


showGame :: Bool
showGame = False


blockBy :: Show a => Int -> [a] -> [[a]]
blockBy _ [] = []
blockBy n l = case splitAt n l of
                  (pre, post@(_:_)) -> pre : blockBy n post
                  (pre, []) -> [pre]

printMap :: (Integral i, Integral j) => Map.Map (i, i) j -> [String]
printMap mp =
    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 [[".#+=O" !! (fromIntegral (fromMaybe 0 (Map.lookup (x, y) mp)))
        | x <- [minx..maxx]]
       | y <- [miny..maxy]]

type Pos = (Int, Int)
data AI =
    AI { aiBall :: [Pos]
       , aiPad :: [Int]
       , aiScreen :: Map.Map Pos Int
       , aiScore :: Integer }

pos :: Integral i => (i, i) -> Pos
pos (a, b) = (fromIntegral a, fromIntegral b)

initAI :: AI
initAI = AI [] [] Map.empty (-1)

processOutput :: (Integer, Integer, Integer) -> AI -> AI
processOutput (-1, 0, score) ai = ai { aiScore = score }
processOutput (x, y, v) ai =
    let p = pos (x, y)
        removed = fromMaybe 0 (Map.lookup p (aiScreen ai))
        screen' = Map.insert (pos (x, y)) (fromIntegral v) (aiScreen ai)
        ball' = if removed == 4 then aiBall ai \\ [p] else aiBall ai
        ball'' = if v == 4 then p : ball' else ball'
        pad' = if removed == 3 then aiPad ai \\ [fromIntegral x] else aiPad ai
        pad'' = if v == 3 then fromIntegral x : pad' else pad'
    in ai { aiScreen = screen', aiBall = ball'', aiPad = pad'' }

provideInput :: AI -> Integer
provideInput ai =
    case (aiBall ai, aiPad ai) of
        ([(x, _)], [pad]) -> if | x < pad -> -1
                                | x > pad -> 1
                                | otherwise -> 0
        _ -> error "Ambiguous pad/ball!"

main :: IO ()
main = do
    program <- parse . head <$> getInput 13

    print $ let output = snd (run program [])
                bd = foldl (\mp [x,y,t] -> Map.insert (x, y) t mp)
                           Map.empty (blockBy 3 output)
            in length (filter (== 2) (Map.elems bd))

    when showGame $ putStr "\x1B[H\x1B[2J"

    let program' = 2 : tail program
        loop ai mcont output = do
            let ai1 = foldl' (\ai' [a,b,c] -> processOutput (a, b, c) ai') ai (blockBy 3 output)

            when showGame $ do
                -- putStr "\x1B[1;1H\x1B[2J"
                -- putStr (unlines (printMap (aiScreen ai1)))
                flip mapM_ (blockBy 3 output) $ \[x,y,v] ->
                    if x /= -1
                        then putStr ("\x1B[" ++ show (y + 1) ++ ";" ++ show (x + 1) ++ "H" ++ [".#+=O" !! fromIntegral v])
                        else return ()
                putStrLn ("\x1B[24;1H" ++ show (aiScore ai1))
                hFlush stdout
                threadDelay 2000

            case mcont of
                Just cont -> case runContinue cont [provideInput ai1] of
                                 Left (cont', out) -> loop ai1 (Just cont') out
                                 Right (_, out) -> loop ai1 Nothing out
                Nothing -> return (aiScore ai1)

    score <- case runInterruptible program' [] of
                 Left (cont, output) -> loop initAI (Just cont) output
                 Right (_, output) -> loop initAI Nothing output
    print score