content: {-# OPTIONS_GHC -fplugin=Polysemy.Plugin #-}
  
  -- | The shard logic
  module Calamity.Gateway.Shard (
    Shard (..),
    newShard,
  ) where
  
  import Calamity.Gateway.DispatchEvents (
    CalamityEvent (Dispatch),
    DispatchData (Ready),
   )
  import Calamity.Gateway.Intents (Intents)
  import Calamity.Gateway.Types (
    ControlMessage (..),
    IdentifyData (
      IdentifyData,
      compress,
      intents,
      largeThreshold,
      presence,
      properties,
      shard,
      token
    ),
    IdentifyProps (IdentifyProps, browser, device),
    ReceivedDiscordMessage (
      EvtDispatch,
      HeartBeatAck,
      HeartBeatReq,
      Hello,
      InvalidSession,
      Reconnect
    ),
    ResumeData (ResumeData, seq, sessionID, token),
    SentDiscordMessage (HeartBeat, Identify, Resume, StatusUpdate),
    Shard (..),
    ShardC,
    ShardFlowControl (..),
    ShardMsg (..),
    ShardState (ShardState, wsConn),
    StatusUpdateData,
   )
  import Calamity.Internal.RunIntoIO (bindSemToIO)
  import Calamity.Internal.Utils (
    debug,
    error,
    info,
    leftToMaybe,
    swap,
    unlessM,
    untilJustFinalIO,
    whenJust,
    whileMFinalIO,
   )
  import Calamity.Metrics.Eff (
    MetricEff,
    modifyGauge,
    registerGauge,
   )
  import Calamity.Types.LogEff (LogEff)
  import Calamity.Types.Token (Token, rawToken)
  import Control.Concurrent (threadDelay)
  import Control.Concurrent.Async (Async, cancel)
  import Control.Concurrent.Chan.Unagi qualified as UC
  import Control.Concurrent.STM (STM, atomically, retry)
  import Control.Concurrent.STM.TBMQueue (
    TBMQueue,
    closeTBMQueue,
    newTBMQueueIO,
    readTBMQueue,
    tryWriteTBMQueue,
    writeTBMQueue,
   )
  import Control.Exception (
    Exception (fromException),
    SomeException,
   )
  import Control.Exception.Safe qualified as Ex
  import Control.Monad (void, when)
  import Control.Monad.State.Lazy (runState)
  import Data.Aeson qualified as A
  import Data.ByteString.Lazy qualified as LBS
  import Data.Default.Class (def)
  import Data.IORef (newIORef)
  import Data.Maybe (fromMaybe)
  import Data.Text qualified as T
  import DiPolysemy (attr, push)
  import Network.Connection qualified as NC
  import Network.TLS qualified as NT
  import Network.TLS.Extra qualified as NT
  import Network.WebSockets (
    Connection,
    ConnectionException (..),
    receiveData,
    sendCloseCode,
    sendTextData,
   )
  import Network.WebSockets qualified as NW
  import Network.WebSockets.Stream qualified as NW
  import Optics
  import Optics.State.Operators
  import Polysemy (Sem)
  import Polysemy qualified as P
  import Polysemy.Async qualified as P
  import Polysemy.AtomicState qualified as P
  import Polysemy.Error qualified as P
  import Polysemy.Resource qualified as P
  import System.X509 qualified as X509
  import TextShow (showt)
  import Prelude hiding (error)
  
  runWebsocket ::
    (P.Members '[LogEff, P.Final IO, P.Embed IO] r) =>
    T.Text ->
    T.Text ->
    (Connection -> P.Sem r a) ->
    P.Sem r (Maybe a)
  runWebsocket host path ma = do
    inner <- bindSemToIO ma
  
    -- We have to do this all ourself I think?
    -- TODO: see if this isn't needed
    let logExc e = debug $ "runWebsocket raised with " <> (T.pack . show $ e)
    logExc' <- bindSemToIO logExc
    let handler e = do
          void $ logExc' e
          pure Nothing
  
    P.embed . Ex.handleAny handler $ do
      ctx <- NC.initConnectionContext
      certStore <- X509.getSystemCertificateStore
      let clientParams =
            (NT.defaultParamsClient (T.unpack host) "443")
              { NT.clientSupported = def {NT.supportedCiphers = NT.ciphersuite_default}
              , NT.clientShared =
                  def
                    { NT.sharedCAStore = certStore
                    }
              }
      let tlsSettings = NC.TLSSettings clientParams
          connParams = NC.ConnectionParams (T.unpack host) 443 (Just tlsSettings) Nothing
  
      Ex.bracket
        (NC.connectTo ctx connParams)
        NC.connectionClose
        ( \conn -> do
            stream <-
              NW.makeStream
                (Just <$> NC.connectionGetChunk conn)
                (maybe (pure ()) (NC.connectionPut conn . LBS.toStrict))
            NW.runClientWithStream stream (T.unpack host) (T.unpack path) NW.defaultConnectionOptions [] inner
        )
  
  newShardState :: Shard -> ShardState
  newShardState shard = ShardState shard Nothing Nothing False Nothing Nothing Nothing
  
  -- | Creates and launches a shard
  newShard ::
    (P.Members '[LogEff, MetricEff, P.Embed IO, P.Final IO, P.Async] r) =>
    T.Text ->
    Int ->
    Int ->
    Token ->
    Maybe StatusUpdateData ->
    Intents ->
    UC.InChan CalamityEvent ->
    Sem r (UC.InChan ControlMessage, Async (Maybe ()))
  newShard gateway id count token presence intents evtIn = do
    (cmdIn, cmdOut) <- P.embed UC.newChan
    let shard = Shard id count gateway evtIn cmdOut (rawToken token) presence intents
    stateVar <- P.embed . newIORef $ newShardState shard
  
    let runShard = P.runAtomicStateIORef stateVar shardLoop
    let action = push "calamity-shard" . attr "shard-id" id $ runShard
  
    thread' <- P.async action
  
    pure (cmdIn, thread')
  
  sendToWs :: (ShardC r) => SentDiscordMessage -> Sem r ()
  sendToWs data' = do
    wsConn' <- P.atomicGets wsConn
    case wsConn' of
      Just wsConn -> do
        let encodedData = A.encode data'
        debug . T.pack $ "sending " <> show data' <> " encoded to " <> show encodedData <> " to gateway"
        P.embed . sendTextData wsConn $ encodedData
      Nothing -> debug "tried to send to closed WS"
  
  tryWriteTBMQueue' :: TBMQueue a -> a -> STM Bool
  tryWriteTBMQueue' q v = do
    v' <- tryWriteTBMQueue q v
    case v' of
      Just False -> retry
      Just True -> pure True
      Nothing -> pure False
  
  restartUnless :: (P.Members '[LogEff, P.Error ShardFlowControl] r) => T.Text -> Maybe a -> P.Sem r a
  restartUnless _ (Just a) = pure a
  restartUnless msg Nothing = do
    error msg
    P.throw ShardFlowRestart
  
  -- | The loop a shard will run on
  shardLoop :: (ShardC r) => Sem r ()
  shardLoop = do
    activeShards <- registerGauge "active_shards" mempty
    void $ modifyGauge (+ 1) activeShards
    void outerloop
    void $ modifyGauge (subtract 1) activeShards
    debug "Shard shut down"
    where
      controlStream :: Shard -> TBMQueue ShardMsg -> IO ()
      controlStream shard outqueue = inner
        where
          q = shard ^. #cmdOut
          inner = do
            v <- UC.readChan q
            r <- atomically $ tryWriteTBMQueue' outqueue (Control v)
            when r inner
  
      handleWSException :: SomeException -> IO (Either (ControlMessage, Maybe T.Text) a)
      handleWSException e = pure $ case fromException e of
        Just (CloseRequest code _)
          | code `elem` [4004, 4010, 4011, 4012, 4013, 4014] ->
              Left (ShutDownShard, Just . showt $ code)
        e -> Left (RestartShard, Just . T.pack . show $ e)
  
      discordStream :: (P.Members '[LogEff, MetricEff, P.Embed IO, P.Final IO] r) => Connection -> TBMQueue ShardMsg -> Sem r ()
      discordStream ws outqueue = inner
        where
          inner = do
            msg <- P.embed $ Ex.catchAny (Right <$> receiveData ws) handleWSException
  
            case msg of
              Left (c, reason) -> do
                whenJust reason (\r -> error . T.pack $ "Shard closed with reason: " <> show r)
                P.embed . atomically $ writeTBMQueue outqueue (Control c)
              Right msg' -> do
                -- debug [fmt|Got msg: {msg'}|]
                let decoded = A.eitherDecode msg'
                r <- case decoded of
                  Right a ->
                    P.embed . atomically $ tryWriteTBMQueue' outqueue (Discord a)
                  Left e -> do
                    error . T.pack $ "Failed to decode " <> e <> ": " <> show msg'
                    pure True
                when r inner
      outerloop :: (ShardC r) => Sem r ()
      outerloop = whileMFinalIO $ do
        shard :: Shard <- P.atomicGets (^. #shardS)
        let host = shard ^. #gateway
        let host' = fromMaybe host $ T.stripPrefix "wss://" host
        info . T.pack $ "starting up shard " <> show (shardID shard) <> " of " <> show (shardCount shard)
  
        innerLoopVal <- runWebsocket host' "/?v=9&encoding=json" innerloop
  
        case innerLoopVal of
          Just ShardFlowShutDown -> do
            info "Shutting down shard"
            pure False
          Just ShardFlowRestart -> do
            info "Restaring shard"
            pure True
          -- we restart normally when we loop
  
          Nothing -> do
            -- won't happen unless innerloop starts using a non-deterministic effect or connecting to the ws dies
            info "Restarting shard (abnormal reasons?)"
            pure True
  
      innerloop :: (ShardC r) => Connection -> Sem r ShardFlowControl
      innerloop ws = do
        debug "Entering inner loop of shard"
  
        shard <- P.atomicGets (^. #shardS)
        P.atomicModify' (#wsConn ?~ ws)
  
        seqNum' <- P.atomicGets (^. #seqNum)
        sessionID' <- P.atomicGets (^. #sessionID)
  
        case (seqNum', sessionID') of
          (Just n, Just s) -> do
            debug $ "Resuming shard (sessionID: " <> s <> ", seq: " <> T.pack (show n)
            sendToWs
              ( Resume
                  ResumeData
                    { token = shard ^. #token
                    , sessionID = s
                    , seq = n
                    }
              )
          _noActiveSession -> do
            debug "Identifying shard"
            sendToWs
              ( Identify
                  IdentifyData
                    { token = shard ^. #token
                    , properties =
                        IdentifyProps
                          { browser = "Calamity: https://github.com/simmsb/calamity"
                          , device = "Calamity: https://github.com/simmsb/calamity"
                          }
                    , compress = False
                    , largeThreshold = Nothing
                    , shard =
                        Just (shard ^. #shardID, shard ^. #shardCount)
                    , presence = shard ^. #initialStatus
                    , intents = shard ^. #intents
                    }
              )
  
        result <-
          P.resourceToIOFinal $
            P.bracket
              (P.embed $ newTBMQueueIO 1)
              (P.embed . atomically . closeTBMQueue)
              ( \q -> do
                  debug "handling events now"
                  _controlThread <- P.async . P.embed $ controlStream shard q
                  _discordThread <- P.async $ discordStream ws q
                  P.raise . untilJustFinalIO . (leftToMaybe <$>) . P.runError $ do
                    -- only we close the queue
                    msg <- P.embed . atomically $ readTBMQueue q
                    handleMsg =<< restartUnless "shard message stream closed by someone other than the sink" msg
              )
  
        debug "Exiting inner loop of shard"
  
        P.atomicModify' (#wsConn .~ Nothing)
        haltHeartBeat
        pure result
      handleMsg :: (ShardC r, P.Member (P.Error ShardFlowControl) r) => ShardMsg -> Sem r ()
      handleMsg (Discord msg) = case msg of
        EvtDispatch sn data' -> do
          -- trace $ "Handling event: ("+||data'||+")"
          P.atomicModify' (#seqNum ?~ sn)
  
          case data' of
            Ready rdata' ->
              P.atomicModify' (#sessionID ?~ (rdata' ^. #sessionID))
            _NotReady -> pure ()
  
          shard <- P.atomicGets (^. #shardS)
          P.embed $ UC.writeChan (shard ^. #evtIn) (Dispatch (shard ^. #shardID) data')
        HeartBeatReq -> do
          debug "Received heartbeat request"
          sendHeartBeat
        Reconnect -> do
          debug "Being asked to restart by Discord"
          P.throw ShardFlowRestart
        InvalidSession resumable -> do
          if resumable
            then info "Received resumable invalid session"
            else do
              info "Received non-resumable invalid session, sleeping for 15 seconds then retrying"
              P.atomicModify' (#sessionID .~ Nothing)
              P.atomicModify' (#seqNum .~ Nothing)
              P.embed $ threadDelay (15 * 1000 * 1000)
          P.throw ShardFlowRestart
        Hello interval -> do
          info . T.pack $ "Received hello, beginning to heartbeat at an interval of " <> show interval <> "ms"
          startHeartBeatLoop interval
        HeartBeatAck -> do
          debug "Received heartbeat ack"
          P.atomicModify' (#hbResponse .~ True)
      handleMsg (Control msg) = case msg of
        SendPresence data' -> do
          debug . T.pack $ "Sending presence: (" <> show data' <> ")"
          sendToWs $ StatusUpdate data'
        RestartShard -> P.throw ShardFlowRestart
        ShutDownShard -> P.throw ShardFlowShutDown
  
  startHeartBeatLoop :: (ShardC r) => Int -> Sem r ()
  startHeartBeatLoop interval = do
    haltHeartBeat -- cancel any currently running hb thread
    thread <- P.async $ heartBeatLoop interval
    P.atomicModify' (#hbThread ?~ thread)
  
  haltHeartBeat :: (ShardC r) => Sem r ()
  haltHeartBeat = do
    thread <- P.atomicState @ShardState . (swap .) . runState $ do
      thread <- use #hbThread
      #hbThread .= Nothing
      pure thread
    case thread of
      Just t -> do
        debug "Stopping heartbeat thread"
        P.embed (void $ cancel t)
      Nothing -> pure ()
  
  sendHeartBeat :: (ShardC r) => Sem r ()
  sendHeartBeat = do
    sn <- P.atomicGets (^. #seqNum)
    debug . T.pack $ "Sending heartbeat (seq: " <> show sn <> ")"
    sendToWs $ HeartBeat sn
    P.atomicModify' (#hbResponse .~ False)
  
  heartBeatLoop :: (ShardC r) => Int -> Sem r ()
  heartBeatLoop interval = untilJustFinalIO . (leftToMaybe <$>) . P.runError $ do
    sendHeartBeat
    P.embed . threadDelay $ interval * 1000
    unlessM (P.atomicGets (^. #hbResponse)) $ do
      debug "No heartbeat response, restarting shard"
      wsConn <- P.note () =<< P.atomicGets (^. #wsConn)
      P.embed $ sendCloseCode wsConn 4000 ("No heartbeat in time" :: T.Text)
      P.throw ()
  