module Control.Actor
  ( module Control.Actor.Types
  , module Control.Actor.Transport
  , module Control.Actor.Runtime
  , module Control.Actor.Core
  , module Control.Actor.Supervision
  , module Control.Actor.Network
  -- Initialization
  , initRuntime
  -- Demo
  , pingActor
  , forwardActorWithCell
  , repeatActor
  , system
  , networkDemo
  ) where

import Control.Actor.Core
import Control.Actor.Network
import Control.Actor.Runtime
import Control.Actor.Supervision
import Control.Actor.Transport
import Control.Actor.Types

import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.STM
  ( TMVar, TVar
  , atomically, newEmptyTMVarIO, newTVarIO, readTMVar, readTVarIO, writeTQueue, writeTVar
  )
import Control.Monad ((<=<), void)
import Control.Monad.Reader (MonadIO (..))
import Unsafe.Coerce (unsafeCoerce)
import Control.Exception (SomeException, try)
import Control.Actor.Registry (createRegistry, registerActor, lookupRemoteActor)

initRuntime :: NodeAddr -> (NodeAddr -> IO (Transport, NodeAddr)) -> IO Runtime
initRuntime myAddr createTransport = do
  (transport, actualAddr) <- createTransport myAddr
  rt0 <- newRuntime actualAddr transport
  let rt = rt0 { rtSendRemote = \addr nm -> getOrCreateConn addr >>= cast' nm }
  withRuntime rt $ void $ spawnActor sysHandlerFn stopOnDeath ()
  withRuntime rt $ void createRegistry
  tListen transport $ \ch -> do
    result <- try @SomeException $ withRuntime rt $ handleNewConn ch
    case result of
      Left _   -> chClose ch
      Right () -> return ()
  return rt


-- Demo

pingActor :: Handler String () String
pingActor msg = continue ("Hello, " <> msg <> "!")

forwardActorWithCell :: TMVar (ActorRef String String) -> Handler String () String
forwardActorWithCell cell msg = do
  pingRef <- liftIO $ atomically $ readTMVar cell
  reply   <- call msg pingRef
  case reply of
    Nothing -> liftIO $ putStrLn "forwardActorWithCell: received empty reply!"
    Just x  -> liftIO $ putStrLn ("forwardActorWithCell: received reply - " <> x)
  maybeContinue reply

repeatActor :: String -> TMVar (ActorRef String String) -> Handler () () String
repeatActor r cell () = do
  ref             <- liftIO $ atomically $ readTMVar cell
  (SomeActorRef self) <- getSelf
  cast r ref
  castIn 1000 () (unsafeCoerce self)
  pass

system :: IO ()
system = do
  rt <- initRuntime (NodeAddr "localhost" 9000) createTCPTransport
  withRuntime rt $ do
    pingCell    <- liftIO newEmptyTMVarIO
    forwardCell <- liftIO newEmptyTMVarIO
    repeatCell  <- liftIO newEmptyTMVarIO

    _ <- liftIO $ forkIO $ do
      repeatRef <- atomically $ readTMVar repeatCell
      withRuntime rt $
        cast' () repeatRef

    supervise'
      OneForOne
      [ childWithRef pingActor stopOnDeath () pingCell
      , childWithRef (forwardActorWithCell pingCell) stopOnDeath () forwardCell
      , childWithRef (repeatActor "repeaaat" forwardCell) stopOnDeath () repeatCell
      ]

-- Network demo

newNodeActor :: Handler NodeAddr () NodeId
newNodeActor = continue <=< liftRuntime . connect Nothing

-- | Actor on node 2: echo — returns whatever it received.
echoActor :: Handler String () String
echoActor = continue

-- | Actor on node 2: printer — side-effects only.
printerActor :: Handler String () ()
printerActor msg = liftIO (putStrLn $ "  [remote/print] " <> msg) >> pass

-- | Trivial actor used as a stub when we only care about death notifications.
noopActor :: Handler () () ()
noopActor _ = pass

-- | Captures the received string into a TVar (used as a node-1 receiver actor).
recvActorFn :: TVar (Maybe String) -> Handler String () ()
recvActorFn var msg =
  liftIO (atomically $ writeTVar var (Just msg)) >> pass

-- | End-to-end networking demo.
--
-- Covers:
--   1. Remote cast  (node 1 → node 2)
--   2. Remote call  (node 1 ↔ node 2, request/reply)
--   3. Reverse cast (node 2 → node 1)
--   4. Cross-node death notification
networkDemo :: IO ()
networkDemo = do
  putStrLn "=== networking demo ==="

  -- Port 0 lets the OS pick a free port each run, avoiding stale-socket
  -- conflicts when re-running in the same GHCi session.
  rt1 <- initRuntime (NodeAddr "localhost" 0) createTCPTransport
  rt2 <- initRuntime (NodeAddr "localhost" 0) createTCPTransport
  let addr1 = rtNodeId rt1
      addr2 = rtNodeId rt2

  -- connect returns the locally-assigned NodeId for the peer.
  -- NodeId 0 is always self; remotes get ids ≥ 1.
  node2 <- withRuntime rt1 $ connect Nothing addr2
  node1 <- withRuntime rt2 $ connect Nothing addr1

  -- Spawn and register workers on node 2.
  echoRef  <- withRuntime rt2 $ spawnActor echoActor   stopOnDeath ()
  printRef <- withRuntime rt2 $ spawnActor printerActor stopOnDeath ()
  withRuntime rt2 $ registerActor "echo"    echoRef
  withRuntime rt2 $ registerActor "printer" printRef

  -- Small delay to let NMRegistryUpdate propagate to node 1.
  threadDelay 100_000

  -- 1. Remote cast: fire-and-forget from node 1 to node 2.
  putStrLn "\n[1] remote cast  node 1 -> node 2"
  Just remotePrint <- withRuntime rt1 $ lookupRemoteActor "printer"
  withRuntime rt1 $ cast' "hello from node 1" remotePrint
  threadDelay 300_000

  -- 2. Remote call: request/reply across nodes.
  putStrLn "\n[2] remote call  node 1 <-> node 2"
  Just remoteEcho' <- withRuntime rt1 $ lookupRemoteActor "echo"
  let
    remoteEcho :: ActorRef String String
    remoteEcho = (unsafeCoerce remoteEcho')
  reply <- withRuntime rt1 $ call' "ping" remoteEcho
  putStrLn $ "    reply: " <> show reply

  -- 3. Reverse cast: actor on node 2 sends to an actor on node 1.
  putStrLn "\n[3] reverse cast  node 2 -> node 1"
  recvVar <- newTVarIO (Nothing :: Maybe String)
  recvRef <- withRuntime rt1 $ spawnActor (recvActorFn recvVar) stopOnDeath ()
  withRuntime rt1 $ registerActor "receiver" recvRef
  threadDelay 100_000
  Just remoteRecv <- withRuntime rt2 $ lookupRemoteActor "receiver"
  withRuntime rt2 $ cast' "hello from node 2" remoteRecv
  threadDelay 300_000
  readTVarIO recvVar >>= \v -> putStrLn $ "    received: " <> show v

  -- 4. Cross-node death notification.
  --
  -- The mortal actor lives on node 2.  When it dies, it sends NMDeath to
  -- node 1.  The watcher actor on node 1 has registered interest in the
  -- mortal's ActorId; routeRemoteDeath finds it and delivers the message.
  putStrLn "\n[4] cross-node death notification"
  diedVar  <- newTVarIO False
  watchRef <- withRuntime rt1 $ spawnActor noopActor
    (\dm -> do
        liftIO $ putStrLn $ "  [watcher] death of " <> show (dmActorId dm)
        liftIO $ atomically $ writeTVar diedVar True
        return Stop)
    ()
  mortalRef <- withRuntime rt2 $ spawnActor noopActor stopOnDeath ()

  let mortalId = someActorId (SomeActorRef mortalRef)
      watchId  = someActorId (SomeActorRef watchRef)

  -- mortalRef's links: on death, send NMDeath to node 1 (by NodeId).
  withRuntime rt2 $ linkActorTo (RemoteTarget watchId  node1) mortalRef
  -- watchRef's links: declare interest in mortalId so routeRemoteDeath
  -- on node 1 can find this actor when NMDeath arrives.
  withRuntime rt1 $ linkActorTo (RemoteTarget mortalId node2) watchRef

  -- Kill the mortal actor by posting directly to its death queue.
  case mortalRef of
    LocalRef {arDeathQ, arState} ->
      atomically $ writeTQueue arDeathQ (DeathMessage (asId arState) Killed)
    RemoteRef _ -> return ()

  threadDelay 500_000
  readTVarIO diedVar >>= \d -> putStrLn $ "    death delivered: " <> show d

  putStrLn "\n=== done ==="