module BiDi.Runner
  ( ChannelActions (..),
    MessageLoops (..),
    mkChannelActions,
    mkFaiChannelActions,
    withBiDi,
    withBidiFailTest,
    unsubscribe,
    subscribe,
  )
where

import BiDi.Socket as S
  ( Channels (..),
    RegisteredSubscription (..),
    SocketActions (..),
    SocketSubscription (..),
    SocketSubscriptionId (..),
    SocketSubscriptionType (..),
    SocketUnregister (..),
    counterVar,
    initChannels,
    mkAtomicCounter,
    mkSocketActions,
  )
import Control.Exception (Exception (displayException), throw)
import Control.Monad (when, (>=>))
import Data.Aeson (FromJSON, Object, Value (..), encode, toJSON, withObject, (.:))
import Data.Aeson.Types (FromJSON (..), Parser)
import Data.ByteString.Lazy qualified as BL
import Data.Coerce (coerce)
import Data.Foldable (traverse_)
import Data.Function ((&))
import Data.Set qualified as Set
import Data.Text as T (Text, pack, take, unpack)
import Data.Word (Word64)
import GHC.Generics (Generic)
import IOUtils (DemoActions (..), Logger (..), catchLog, loopForever)
import Network.WebSockets (Connection, receiveData, runClient, sendTextData)
import UnliftIO (catchAny, throwIO, waitAnyCatch)
import UnliftIO.Async (Async, async, cancel)
import UnliftIO.STM
import BiDi.BiDiUrl (BiDiUrl (..))
import WebDriverPreCore.BiDi.Protocol as P
  ( JSUInt (..),
    SessionSubscribeResult (..),
    SessionSubscibe (..),
    SessionUnsubscribe (..),
    Subscription (..),
    SubscriptionId (..),
    SubscriptionType (..),
    subscriptionTypeToText
  )
import AesonUtils (jsonToText, objToText, parseThrow)
import Utils (txt)
import Prelude hiding (getLine, log, null, print, putStrLn)
import BiDi.Response
  ( ResponseObject (..),
    decodeResponse,
  )

-- | Run WebDriver BiDi client and return a client interface
withBiDi :: DemoActions -> BiDiUrl -> (DemoActions -> SocketActions -> IO ()) -> IO ()
withBiDi = withBidi' mkChannelActions

-- | Run WebDriver BiDi client and return a client interface
withBidi' :: (Logger -> IO ChannelActions) -> DemoActions -> BiDiUrl -> (DemoActions -> SocketActions -> IO ()) -> IO ()
withBidi'
  mkChActions
  demoActions@MkDemoActions {logTxt}
  bidiUrl
  action = do
    ca <- mkChActions logger
    withSocket bidiUrl logger ca.messageLoops $
      action demoActions ca.socketActions
    where
      logger = MkLogger logTxt

mkChannelActions :: Logger -> IO ChannelActions
mkChannelActions logger = do
  c <- initChannels
  pure $
    MkChannelActions
      { socketActions = mkSocketActions c,
        messageLoops = demoMessageLoops logger c
      }

data ChannelActions = MkChannelActions
  { messageLoops :: MessageLoops,
    socketActions :: SocketActions
  }

data MessageActions = MkMessageActions
  { send :: Connection -> IO (),
    get :: Connection -> IO (),
    eventHandler :: IO ()
  }

demoMessageActions :: Logger -> Channels -> MessageActions
demoMessageActions logger MkChannels {sendChan, receiveChan, eventChan, subscriptions} =
  MkMessageActions
    { send = \conn -> do
        msgToSend <- atomically $ readTChan sendChan
        log $ "Sending Message: " <> jsonToText msgToSend
        catchLog
          "Message Send Failed"
          log
          (sendTextData conn (BL.toStrict $ encode msgToSend)),
      --
      get = \conn -> do
        msg <- receiveData conn
        log $ "Received raw data: " <> T.take 100 (txt msg) <> "..."
        log $ "Received raw data: " <> txt msg
        let writeReceiveChan = atomically . writeTChan receiveChan
            writeEventChan = atomically . writeTChan eventChan
            r = decodeResponse msg
        case r of
          Left {} -> writeReceiveChan r
          Right r' -> case r' of
            NoID obj -> writeEventChan obj
            WithID {} -> writeReceiveChan r,
      --
      eventHandler = do
        obj <- atomically $ readTChan eventChan
        log $ "Event received: " <> jsonToText (toJSON obj)
        applySubscriptions log obj subscriptions
    }
  where
    log :: Text -> IO ()
    log = logger.log

--
data EventProps = MkEventProps
  { msgType :: Text,
    method :: Text,
    params :: Value,
    fullObj :: Value
  }
  deriving (Show, Generic)

instance FromJSON EventProps where
  parseJSON :: Value -> Parser EventProps
  parseJSON v =
    withObject
      "EventProps"
      ( \o ->
          MkEventProps
            <$> o .: "type"
            <*> o .: "method"
            <*> o .: "params"
            <*> pure v
      )
      v

unsubscribe :: SocketActions -> (SessionUnsubscribe -> IO ()) -> SessionUnsubscribe -> IO ()
unsubscribe sa callUnsubscribe unsub = do
  callUnsubscribe unsub
  atomically . sa.unregisterSubscription $ toSocketUnregister unsub
  where
    toSocketUnregister :: SessionUnsubscribe -> SocketUnregister
    toSocketUnregister = \case
      UnsubscribeById {subscriptions} ->
        UnregisterById . Set.fromList $ MkSocketSubscriptionId . coerce <$> subscriptions
      UnsubscribeByAttributes {unsubEvents} ->
        UnregisterByAttributes . Set.fromList $ MkSocketSubscriptionType . subscriptionTypeToText <$> unsubEvents

subscribe ::
  SocketActions ->
  (SessionSubscibe -> IO SessionSubscribeResult) ->
  Subscription IO ->
  IO SubscriptionId
subscribe sa callSubscribe subscription = do
  -- subscribe with a dummy id first so we don't miss any messages
  atomically $ subscribeWithId dummySubId
  catchAny
    ( do
        subId <- callSubscribe $ mkRequest subscription
        -- swap in the real id
        atomically $ do
          removeDummySub
          subscribeWithId $ coerce subId.subscription
        pure subId.subscription
    )
    ( \e -> do
        -- on error, remove the dummy subscription
        atomically removeDummySub
        throwIO e
    )
  where
    mkRequest :: Subscription IO -> SessionSubscibe
    mkRequest s = case s of
      P.SingleSubscription
        { subscriptionType
        } ->
          MkSessionSubscribe
            { events = [coerce subscriptionType],
              contexts,
              userContexts
            }
      P.MultiSubscription
        { subscriptionTypes
        } ->
          MkSessionSubscribe
            { events = coerce <$> subscriptionTypes,
              contexts,
              userContexts
            }
      P.OffSpecSubscription
        { subscriptionTypes
        } ->
          MkSessionSubscribe
            { events = coerce <$> subscriptionTypes,
              contexts,
              userContexts
            }
      where
        contexts = maybeList s.browsingContexts
        userContexts = maybeList s.userContexts
        maybeList :: [a] -> Maybe [a]
        maybeList = \case
          [] -> Nothing
          xs -> Just xs

    mkRegistration :: Subscription IO -> SocketSubscription
    mkRegistration = \case
      P.SingleSubscription
        { subscriptionType,
          action
        } ->
          S.SingleSubscription
            { subscriptionType = toSocketSubType subscriptionType,
              action
            }
      s' -> case s' of
        P.MultiSubscription
          { nAction
          } ->
            S.MultiSubscription
              { subscriptionTypes = socketSubtypes,
                nAction = parseThrow ("Could not parse Event when executing MultiSubscription action for: " <> socketSubtypesTxt) >=> nAction
              }
        P.OffSpecSubscription
          { nValueAction
          } ->
            S.MultiSubscription
              { subscriptionTypes = socketSubtypes,
                nAction = parseThrow ("Could not parse Value when executing OffSpecSubscription action for: " <> socketSubtypesTxt) >=> nValueAction
              }
        where
          socketSubtypes = Set.fromList $ toSocketSubType <$> s'.subscriptionTypes
          socketSubtypesTxt = txt (Set.toList socketSubtypes)

    dummySubId = MkSocketSubscriptionId "dummy"

    subscribeWithId :: SocketSubscriptionId -> STM ()
    subscribeWithId =
      sa.registerSubscription (mkRegistration subscription)

    removeDummySub :: STM ()
    removeDummySub = sa.unregisterSubscription . UnregisterById $ Set.singleton dummySubId

applySubscriptions :: (Text -> IO ()) -> Object -> TVar [RegisteredSubscription IO] -> IO ()
applySubscriptions _log obj subscriptions = do
  MkEventProps {msgType, method, fullObj, params} <-
    parseThrow "Could not parse event properties" (Object obj)
  when (msgType /= "event") $
    fail . unpack $
      "Event message expected. This is not an event message: "
        <> msgType
        <> "\n"
        <> objToText obj

  -- log $ "Parsed event: " <> txt eventProps
  subs <- readTVarIO subscriptions
  traverse_ (applySubscription (coerce method) params fullObj) ((.subscription) <$> subs)

applySubscription :: SocketSubscriptionType -> Value -> Value -> SocketSubscription -> IO ()
applySubscription subType params fullObj =
  \case
    S.SingleSubscription {subscriptionType, action} ->
      when (subType == subscriptionType) $
        parseThrow ("could not parse Event for: " <> txt subType) params
          >>= action
    S.MultiSubscription {subscriptionTypes, nAction} ->
      when (subType `Set.member` subscriptionTypes) $
        nAction fullObj

toSocketSubType :: SubscriptionType -> SocketSubscriptionType
toSocketSubType = MkSocketSubscriptionType . subscriptionTypeToText

loopActions :: Logger -> MessageActions -> MessageLoops
loopActions logger MkMessageActions {..} =
  MkMessageLoops
    { sendLoop = asyncLoop "Sender" . send,
      getLoop = asyncLoop "Receiver" . get,
      eventLoop = asyncLoop "EventHandler" eventHandler
    }
  where
    asyncLoop = loopForever logger.log

data MessageLoops = MkMessageLoops
  { sendLoop :: Connection -> IO (Async ()),
    getLoop :: Connection -> IO (Async ()),
    eventLoop :: IO (Async ())
  }

demoMessageLoops :: Logger -> Channels -> MessageLoops
demoMessageLoops logger channels =
  loopActions logger $ demoMessageActions logger channels

withSocket :: BiDiUrl -> Logger -> MessageLoops -> IO () -> IO ()
withSocket
  pth@MkBiDiUrl {host, port, path}
  logger
  messageLoops
  action =
    do
      log $ "Connecting to WebDriver at " <> txt pth
      runClient (unpack host) port (unpack path) $ \conn -> do
        eventLoop <- messageLoops.eventLoop
        getLoop <- messageLoops.getLoop conn
        sendLoop <- messageLoops.sendLoop conn

        log "WebSocket connection established"

        result <- async action

        (_asy, ethresult) <- waitAnyCatch [getLoop, sendLoop, result, eventLoop]

        -- cancelMany not reexported by UnliftIO
        traverse_ cancel [getLoop, sendLoop, result, eventLoop]

        ethresult
          & either
            ( \e -> do
                log $ "One of the BiDi client threads failed: \n" <> pack (displayException e)
                throw e
            )
            pure
    where
      log :: Text -> IO ()
      log = coerce logger

-- ####################################################################################
-- ########################  Fail Actions - for testing errors ########################
-- ####################################################################################

-- | Run WebDriver BiDi client and return a client interface
withBidiFailTest ::
  Word64 ->
  Word64 ->
  Word64 ->
  DemoActions ->
  BiDiUrl ->
  (DemoActions -> SocketActions -> IO ()) ->
  IO ()
withBidiFailTest failSendCount failGetCount failEventCount =
  withBidi' (mkFaiChannelActions failSendCount failGetCount failEventCount)

mkFaiChannelActions ::
  Word64 ->
  Word64 ->
  Word64 ->
  Logger ->
  IO ChannelActions
mkFaiChannelActions failSendCount failGetCount failEventCount logger = do
  c <- initChannels
  messageLoops <- failMessageLoops logger c failSendCount failGetCount failEventCount
  pure $
    MkChannelActions
      { socketActions = mkSocketActions c,
        messageLoops
      }

failMessageLoops ::
  Logger ->
  Channels ->
  Word64 ->
  Word64 ->
  Word64 ->
  IO MessageLoops
failMessageLoops logger channels failSendCount failGetCount failEventCount =
  loopActions logger
    <$> failMessageActions
      (demoMessageActions logger channels)
      failSendCount
      failGetCount
      failEventCount

failMessageActions :: MessageActions -> Word64 -> Word64 -> Word64 -> IO MessageActions
failMessageActions a failSendCount failGetCount failEventCount =
  do
    send <- failAction "send" failSendCount a.send
    get <- failAction "get" failGetCount a.get
    eventHandler' <- failAction "eventhandler" failEventCount $ const a.eventHandler
    pure $
      MkMessageActions
        { send,
          get,
          eventHandler = eventHandler' 1
        }

failAction :: Text -> Word64 -> (a -> IO ()) -> IO ((a -> IO ()))
failAction lbl failCallCount action = do
  counterVar' <- counterVar
  let counter = mkAtomicCounter counterVar'
  pure $ \a -> do
    n <- atomically $ counter
    if (coerce n :: Word64) == failCallCount
      then do
        fail $ "Forced failure for testing: " <> unpack lbl <> " (call #" <> show n <> ")"
      else do
        action a