module Main where

import qualified Data.Array as A
import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad
import qualified Data.IntSet as IntSet

-- import Debug.Trace

import Input
import IntCode


data Message = Message Int Integer Integer deriving (Show)
data CompIn = CompRecvMsg Message | IdlePoll | Quit deriving (Show)
data CompOut = CompSendMsg Int Message | Idle Int deriving (Show)
data NatIn = NatRecvMsg Message | NetworkIdle deriving (Show)
data NatOut = NatSendMsg Message | Answer 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"

chanFunnel :: (a -> c) -> TChan a -> (b -> c) -> TChan b -> IO (ThreadId, TChan c)
chanFunnel f1 in1Chan f2 in2Chan = do
    outChan <- atomically $ newTChan
    th1 <- forkIO $ forever $ atomically (readTChan in1Chan) >>= atomically . writeTChan outChan . f1
    th2 <- forkIO $ forever $ atomically (readTChan in2Chan) >>= atomically . writeTChan outChan . f2
    th <- forkFinally (forever yield) (const $ killThread th1 >> killThread th2)
    return (th, outChan)

threadComputer :: Int -> [Integer] -> TChan CompIn -> TChan CompOut -> IO ()
threadComputer myid nicProgram inChan outChan =
    case runInterruptible nicProgram [fromIntegral myid] of
        Left (cont, output) -> sendMsgs output >> loop cont
        Right _ -> undefined
  where
    loop :: Continuation -> IO ()
    loop cont = do
        incoming <- atomically $ readTChan inChan
        case incoming of
            CompRecvMsg (Message _ x y) -> do
                (cont', output) <-
                    case runContinue cont [x, y] of
                        Left pair -> return pair
                        Right (_, out) -> do
                            putStrLn $ "!! Computer " ++ show myid ++ ": intcode terminated! (in recvmsg)"
                            return (undefined, out)
                sendMsgs output
                loop cont'
            IdlePoll -> do
                (cont', output) <-
                    case runContinue cont [-1] of
                        Left pair -> return pair
                        Right (_, out) -> do
                            putStrLn $ "!! Computer " ++ show myid ++ ": intcode terminated! (in idlepoll)"
                            return (undefined, out)
                if null output
                    then atomically $ writeTChan outChan (Idle myid)
                    else sendMsgs output
                loop cont'
            Quit ->
                return ()

    sendMsgs :: [Integer] -> IO ()
    sendMsgs output = atomically $ sequence_ [writeTChan outChan (CompSendMsg myid msg)
                                             | msg <- parseMessages output]

threadNAT :: TChan NatIn -> TChan NatOut -> TVar (Maybe Integer) -> IO ()
threadNAT inChan outChan part1 = loop Nothing Nothing
  where
    loop :: Maybe Message -> Maybe Integer -> IO ()
    loop msave mlastY = do
        incoming <- atomically $ readTChan inChan
        case incoming of
            NatRecvMsg msg@(Message _ _ y) -> do
                -- Set 'part1' to contain the first y value ever sent to the NAT
                atomically $ modifyTVar' part1 (\m -> case m of
                                                          Just _ -> m
                                                          Nothing -> Just y)
                loop (Just msg) mlastY
            NetworkIdle -> do
                -- traceM "NAT: network idle received"
                case (msave, mlastY) of
                    (Just (Message _ _ y), Just y') | y == y' ->
                        atomically $ writeTChan outChan (Answer y)
                    (Just (Message _ x y), _) -> do
                        atomically $ writeTChan outChan (NatSendMsg (Message 0 x y))
                        loop Nothing (Just y)
                    (Nothing, _) -> do
                        putStrLn "!! NAT deadlock: network idle but no message queued in NAT!"
                        atomically $ writeTChan outChan (Answer (-1))

arbitrator :: TChan (Either NatOut CompOut) -> A.Array Int (TChan CompIn) -> TChan NatIn -> IO Integer
arbitrator inChan compChans natChan = loop Nothing
  where
    loop :: Maybe IntSet.IntSet -> IO Integer
    loop midle = do
        yield
        incoming <- atomically $ tryReadTChan inChan
        case incoming of
            Just (Left (NatSendMsg msg)) ->
                sendMsg (-1) msg >> loop Nothing
            Just (Left (Answer y)) -> do
                sequence_ [atomically $ writeTChan chan Quit | chan <- A.elems compChans]
                return y
            Just (Right (CompSendMsg from msg)) ->
                sendMsg from msg >> loop Nothing
            Just (Right (Idle from)) ->
                case midle of
                    Just idle ->
                        let idle' = IntSet.insert from idle
                        in if IntSet.size idle' == 50
                               then atomically (writeTChan natChan NetworkIdle) >> loop Nothing
                               else loop (Just idle')
                    Nothing ->
                        loop midle
            Nothing -> do
                -- traceM "ar: Sending deadlock poll"
                sequence_ [atomically $ writeTChan chan IdlePoll | chan <- A.elems compChans]
                loop (Just IntSet.empty)

    sendMsg :: Int -> Message -> IO ()
    -- sendMsg from msg | trace ("ar: send " ++ show from ++ " -> " ++ show msg) False = undefined
    sendMsg _ msg@(Message 255 _ _) = atomically $ writeTChan natChan (NatRecvMsg msg)
    sendMsg _ msg@(Message to _ _) = atomically $ writeTChan (compChans A.! to) (CompRecvMsg msg)

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

    part1 <- atomically $ newTVar Nothing

    compChans <- atomically $ A.listArray (0, 49) <$> sequence (replicate 50 newTChan)
    compOutChan <- atomically $ newTChan
    natChan <- atomically $ newTChan
    natOutChan <- atomically $ newTChan
    (_, collectChan) <- chanFunnel Left natOutChan Right compOutChan

    _ <- sequence [forkIO (threadComputer i nicProgram (compChans A.! i) compOutChan)
                  | i <- [0..49]]

    _ <- forkIO (threadNAT natChan natOutChan part1)

    y <- arbitrator collectChan compChans natChan
    Just part1y <- atomically (readTVar part1)
    print part1y
    print y