-- Part 2: 17281 is too high

{-# LANGUAGE TupleSections #-}
module Main where

import Data.Foldable (toList)
import qualified Data.Map.Strict as Map
import Data.Maybe
import qualified Data.Sequence as S

import Debug.Trace

import Input
import IntCode


safeLast :: [a] -> Maybe a
safeLast [] = Nothing
safeLast [x] = Just x
safeLast (_:xs) = safeLast xs

data Computer
    = Computer { cCont :: Continuation
               , cQueue :: S.Seq Integer }

queueAdd :: Message -> Computer -> Computer
queueAdd (Message _ x y) comp = comp { cQueue = cQueue comp S.|> x S.|> y }

queueAdds :: [Message] -> Computer -> Computer
queueAdds msgs comp = foldr queueAdd comp msgs

data Message = Message Int Integer Integer
  deriving (Show)

parseMessages :: [Integer] -> [Message]
parseMessages [] = []
parseMessages (to:x:y:rest) = Message (fromIntegral to) x y : parseMessages rest
parseMessages _ = error "Non-3-blocked output"

runComputer :: Computer -> ([Message], Computer)
runComputer (Computer cont queue) =
    let queue' = if null queue then [-1] else toList queue
        (cont', output) = either id ((undefined,) . snd) (runContinue cont queue')
    in (parseMessages output, Computer cont' mempty)

step :: Map.Map Int Computer -> Int -> ([Message], Map.Map Int Computer)
step network cid =
    let (msgs, computer') = runComputer (network Map.! cid)
        network' = foldr (\msg@(Message to _ _) -> Map.adjust (queueAdd msg) to)
                         (Map.insert cid computer' network)
                         msgs
    in (msgs, network')

runUntilNAT :: Map.Map Int Computer -> Int -> Message
runUntilNAT network cid =
    let (msgs, network') = step network cid
    in case [msg | msg@(Message 255 _ _) <- msgs] of
           msg:_ -> msg
           [] -> runUntilNAT network' ((cid + 1) `mod` 50)

-- NAT (last recv'd msg (with to=0)) (last Y value sent to 0) (number of idle computers)
data NAT = NAT (Maybe Message) (Maybe Integer) Int

initNAT :: NAT
initNAT = NAT Nothing Nothing 0

natNotify :: NAT -> Int -> [Message] -> S.Seq Integer -> NAT
natNotify (NAT origBin zeroY idle) from msgs queue =
    let recv = [msg | msg@(Message 255 _ _) <- msgs]
        newBin = safeLast (maybeToList origBin ++ [Message 0 x y | Message _ x y <- recv])
        idle' = if null msgs && S.null queue then idle + 1 else 0
    in -- (if not (null recv) then trace ("NAT received " ++ show recv ++ " from " ++ show from) else id) $
       (if idle' >= 50 then trace "idle!" else id) $
       NAT newBin zeroY idle'

natPoll :: NAT -> Either Integer ([Message], NAT)
-- natPoll (NAT _ _ idle) | traceShow ("poll", idle) False = undefined
natPoll (NAT (Just bin) zeroY idle) | idle >= 50 =
    case (bin, zeroY) of
        (Message 0 _ y, Just y') | y == y' -> Left y
        (Message 0 _ y, _) -> Right ([bin], NAT Nothing (Just y) 0)
        (_, _) -> Right ([bin], NAT Nothing zeroY 0)
natPoll nat@(NAT _ _ _) = Right ([], nat)

traceMessages :: [(Int, Message)] -> a -> a
traceMessages = flip (foldr (\(from, Message to x y) -> trace ("Message " ++ show from ++ " -> " ++ show to ++ ": [" ++ show x ++ "," ++ show y ++ "]")))

runWithNAT :: NAT -> Map.Map Int Computer -> Int -> Integer
runWithNAT nat network cid =
    -- trace ("runWithNAT cid=" ++ show cid) $
    let (msgs, network') = step network cid
        nat' = natNotify nat cid msgs (cQueue (network Map.! cid))
    in trace ("step " ++ show cid ++ " in: " ++ show (cQueue (network Map.! cid)) ++ ", out: " ++ show msgs) $
       traceMessages (map (cid,) msgs) $
       case natPoll nat' of
           Left res -> res
           Right (natout, nat'') ->
               -- (if not (null natout) then trace ("NAT sent " ++ show natout) else id) $
               traceMessages (map (-1,) natout) $
               runWithNAT nat'' (Map.adjust (queueAdds natout) 0 network') ((cid + 1) `mod` 50)

main :: IO ()
main = do
    nicProgram <- parse . head <$> getInput 23

    let initComputer i = Computer (initialContinuation nicProgram) (S.singleton (fromIntegral i))
        initNetwork = Map.fromList [(i, initComputer i) | i <- [0..49]]
        Message _ _ y = runUntilNAT initNetwork 0
    print y

    print (runWithNAT initNAT initNetwork 0)