{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Testlib.Cannon
  ( WebSocket (..),
    WSConnect (..),
    ToWSConnect (..),
    AwaitResult (..),
    withWebSocket,
    withWebSockets,
    awaitNMatchesResult,
    awaitNMatchesResultFor,
    awaitNMatches,
    awaitMatch,
    awaitMatchFor,
    awaitAnyEvent,
    awaitAtLeastNMatchesResult,
    awaitAtLeastNMatches,
    awaitNToMMatchesResult,
    awaitNToMMatches,
    assertAwaitResult,
    nPayload,
    printAwaitResult,
    printAwaitAtLeastResult,
    waitForResponse,
    assertNoEvent,
  )
where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Concurrent.STM.TChan
import Control.Exception (throwIO)
import qualified Control.Exception as E
import Control.Monad
import Control.Monad.Catch hiding (bracket)
import qualified Control.Monad.Catch as Catch
import Control.Monad.IO.Class
import Control.Monad.Reader (asks)
import Control.Monad.STM
import Data.Aeson hiding ((.=))
import Data.ByteString (ByteString)
import Data.ByteString.Conversion (fromByteString)
import Data.ByteString.Conversion.To
import Data.Function
import Data.Maybe
import Data.Traversable
import Data.Word
import GHC.Records
import GHC.Stack
import qualified Network.HTTP.Client as HTTP
import qualified Network.HTTP.Client as Http
import qualified Network.WebSockets as WS
import System.Random (randomIO)
import System.Timeout (timeout)
import Testlib.App
import Testlib.Assertions
import Testlib.Env
import Testlib.HTTP
import Testlib.JSON
import Testlib.Printing
import Testlib.Types
import UnliftIO (withRunInIO)
import Prelude

data WebSocket = WebSocket
  { wsConnect :: WSConnect,
    wsChan :: TChan Value,
    wsCloseLatch :: MVar (),
    wsAppThread :: Async ()
  }

instance HasField "client" WebSocket (Maybe ClientIdentity) where
  getField ws = do
    c <- ws.wsConnect.client
    pure
      ClientIdentity
        { domain = ws.wsConnect.domain,
          user = ws.wsConnect.user,
          client = c
        }

instance HasField "user" WebSocket Value where
  getField ws =
    object
      [ "domain" .= ws.wsConnect.domain,
        "id" .= ws.wsConnect.user
      ]

-- Specifies how a Websocket at cannon should be opened
data WSConnect = WSConnect
  { user :: String,
    domain :: String,
    client :: Maybe String,
    -- | If this is Nothing then a random Z-Connection will be used
    conn :: Maybe String
  }

class ToWSConnect a where
  toWSConnect :: (HasCallStack) => a -> App WSConnect

instance {-# OVERLAPPING #-} ToWSConnect WSConnect where
  toWSConnect = pure

instance {-# OVERLAPPABLE #-} (MakesValue user) => ToWSConnect user where
  toWSConnect u = do
    (domain, uid) <- objQid u
    mc <- lookupField u "client_id"
    c <- traverse asString mc
    pure (WSConnect uid domain c Nothing)

instance (MakesValue user, MakesValue conn) => ToWSConnect (user, conn) where
  toWSConnect (u, c) = do
    (domain, uid) <- objQid u
    conn <- make c & asString
    pure (WSConnect uid domain Nothing (Just conn))

instance (MakesValue user, MakesValue conn, MakesValue client) => ToWSConnect (user, conn, client) where
  toWSConnect (u, c, cl) = do
    (domain, uid) <- objQid u
    client <- make cl & asString
    conn <- make c & asString
    pure (WSConnect uid domain (Just client) (Just conn))

connect :: (HasCallStack) => WSConnect -> App WebSocket
connect wsConnect = do
  nchan <- liftIO newTChanIO
  latch <- liftIO newEmptyMVar
  wsapp <- run wsConnect (clientApp nchan latch)
  pure $ WebSocket wsConnect nchan latch wsapp

clientApp :: (HasCallStack) => TChan Value -> MVar () -> WS.ClientApp ()
clientApp wsChan latch conn = do
  r <- async wsRead
  w <- async wsWrite
  void $ waitEitherCancel r w
  where
    wsRead = forever $ do
      bs <- WS.receiveData conn
      case decodeStrict' bs of
        Just n -> atomically $ writeTChan wsChan n
        Nothing -> putStrLn $ "Failed to decode notification: " ++ show bs
    wsWrite = forever $ do
      takeMVar latch
      WS.sendClose conn ("close" :: ByteString)

-- | Start a client thread in 'Async' that opens a web socket to a Cannon, wait
--   for the connection to register with Gundeck, and return the 'Async' thread.
run ::
  (HasCallStack) =>
  WSConnect ->
  WS.ClientApp () ->
  App (Async ())
run wsConnect app = do
  domain <- asString wsConnect.domain
  serviceMap <- getServiceMap domain

  let HostPort caHost caPort = serviceHostPort serviceMap Cannon
  latch <- liftIO newEmptyMVar

  connId <- case wsConnect.conn of
    Just c -> pure c
    Nothing -> show <$> liftIO (randomIO :: IO Word32)

  let path =
        "/await"
          <> ( case wsConnect.client of
                 Nothing -> ""
                 Just client -> fromJust . fromByteString $ Http.queryString (Http.setQueryString [("client", Just (toByteString' client))] Http.defaultRequest)
             )
      caHdrs =
        [ ("Z-User", toByteString' (wsConnect.user)),
          ("Z-Connection", toByteString' connId)
        ]
  request <- do
    r <- rawBaseRequest domain Cannon Versioned path
    pure r {HTTP.requestHeaders = caHdrs}

  wsapp <-
    liftIO
      $ async
      $ catch
        ( WS.runClientWith
            caHost
            (fromIntegral caPort)
            path
            WS.defaultConnectionOptions
            caHdrs
            app
        )
      $ \(e :: SomeException) -> putMVar latch e

  presenceRequest <-
    baseRequest domain Cannon Unversioned $
      "/i/presences/" <> wsConnect.user <> "/" <> connId

  waitForPresence <- appToIO $ retryT $ do
    response <- submit "HEAD" presenceRequest
    status response `shouldMatchInt` 200
  let waitForException = do
        ex <- takeMVar latch
        -- Construct a "fake" response. We do not really have access to the
        -- websocket connection requests and response, unfortunately, but it is
        -- useful to display some information about the request in case an
        -- exception occurs.
        let r =
              Response
                { json = Nothing,
                  body = "This is a fake response. The actual response from cannon is not available.",
                  status = 101,
                  headers = mempty,
                  request = request
                }
        throwIO (AssertionFailure callStack (Just r) Nothing (displayException ex))

  liftIO $ race_ waitForPresence waitForException
  pure wsapp

close :: (MonadIO m) => WebSocket -> m ()
close ws = liftIO $ do
  putMVar (wsCloseLatch ws) ()
  void $ waitCatch (wsAppThread ws)

withWebSocket :: (HasCallStack, ToWSConnect w) => w -> (WebSocket -> App a) -> App a
withWebSocket w k = do
  wsConnect <- toWSConnect w
  Catch.bracket (connect wsConnect) close k

withWebSockets :: forall a w. (HasCallStack, ToWSConnect w) => [w] -> ([WebSocket] -> App a) -> App a
withWebSockets twcs k = do
  wcs <- for twcs toWSConnect
  go wcs []
  where
    go :: (HasCallStack) => [WSConnect] -> [WebSocket] -> App a
    go [] wss = k (reverse wss)
    go (wc : wcs) wss = withWebSocket wc (\ws -> go wcs (ws : wss))

data AwaitResult = AwaitResult
  { success :: Bool,
    nMatchesExpected :: Int,
    matches :: [Value],
    nonMatches :: [Value]
  }

data AwaitAtLeastResult = AwaitAtLeastResult
  { success :: Bool,
    nMatchesExpectedMin :: Int,
    nMatchesExpectedMax :: Maybe Int,
    matches :: [Value],
    nonMatches :: [Value]
  }

prettyAwaitResult :: AwaitResult -> App String
prettyAwaitResult r = do
  matchesS <- for r.matches prettyJSON
  nonMatchesS <- for r.nonMatches prettyJSON
  pure $
    "AwaitResult\n"
      <> indent
        4
        ( unlines
            [ "success: " <> show (r.success),
              "matches:\n" <> unlines matchesS,
              "non-matches:\n" <> unlines nonMatchesS
            ]
        )

prettyAwaitAtLeastResult :: AwaitAtLeastResult -> App String
prettyAwaitAtLeastResult r = do
  matchesS <- for r.matches prettyJSON
  nonMatchesS <- for r.nonMatches prettyJSON
  pure $
    "AwaitAtLeastResult\n"
      <> indent
        4
        ( unlines
            [ "min expected:" <> show r.nMatchesExpectedMin,
              "max expected:" <> show r.nMatchesExpectedMax,
              "success: " <> show (r.success),
              "matches:\n" <> unlines matchesS,
              "non-matches:\n" <> unlines nonMatchesS
            ]
        )

printAwaitResult :: AwaitResult -> App ()
printAwaitResult = prettyAwaitResult >=> liftIO . putStrLn

printAwaitAtLeastResult :: AwaitAtLeastResult -> App ()
printAwaitAtLeastResult = prettyAwaitAtLeastResult >=> liftIO . putStrLn

awaitAnyEvent :: Int -> WebSocket -> App (Maybe Value)
awaitAnyEvent tSecs = liftIO . timeout (tSecs * 1000 * 1000) . atomically . readTChan . wsChan

-- | 'await' an expected number of notification events on the websocket that
-- satisfy the provided predicate. If there isn't any new event (matching or
-- non-matching) for a 'tSecs' seconds then AwaitResult is a failure. This
-- function will never terminate if there is a constant stream of events
-- received. When this functions returns it will push any non-matching
-- events back to the websocket.
awaitNMatchesResult ::
  (HasCallStack) =>
  -- | Number of matches
  Int ->
  -- | Selection function. Exceptions are *not* caught.
  (Value -> App Bool) ->
  WebSocket ->
  App AwaitResult
awaitNMatchesResult nExpected checkMatch ws = do
  tSecs <- asks timeOutSeconds
  awaitNMatchesResultFor tSecs nExpected checkMatch ws

awaitNMatchesResultFor ::
  (HasCallStack) =>
  Int ->
  Int ->
  (Value -> App Bool) ->
  WebSocket ->
  App AwaitResult
awaitNMatchesResultFor tSecs nExpected checkMatch ws = go nExpected [] []
  where
    go 0 nonMatches matches = do
      refill nonMatches
      pure $
        AwaitResult
          { success = True,
            nMatchesExpected = nExpected,
            matches = reverse matches,
            nonMatches = reverse nonMatches
          }
    go nLeft nonMatches matches = do
      mEvent <- awaitAnyEvent tSecs ws
      case mEvent of
        Just event ->
          do
            isMatch <- checkMatch event
            if isMatch
              then go (nLeft - 1) nonMatches (event : matches)
              else go nLeft (event : nonMatches) matches
        Nothing -> do
          refill nonMatches
          pure $
            AwaitResult
              { success = False,
                nMatchesExpected = nExpected,
                matches = reverse matches,
                nonMatches = reverse nonMatches
              }
    refill = mapM_ (liftIO . atomically . writeTChan (wsChan ws))

awaitAtLeastNMatchesResult ::
  (HasCallStack) =>
  -- | Minimum number of matches
  Int ->
  -- | Selection function. Exceptions are *not* caught.
  (Value -> App Bool) ->
  WebSocket ->
  App AwaitAtLeastResult
awaitAtLeastNMatchesResult nExpected checkMatch ws = go 0 [] []
  where
    go nSeen nonMatches matches = do
      tSecs <- asks timeOutSeconds
      mEvent <- awaitAnyEvent tSecs ws
      case mEvent of
        Just event ->
          do
            isMatch <- checkMatch event
            if isMatch
              then go (nSeen + 1) nonMatches (event : matches)
              else go nSeen (event : nonMatches) matches
        Nothing -> do
          refill nonMatches
          pure $
            AwaitAtLeastResult
              { success = nSeen >= nExpected,
                nMatchesExpectedMin = nExpected,
                nMatchesExpectedMax = Nothing,
                matches = reverse matches,
                nonMatches = reverse nonMatches
              }
    refill = mapM_ (liftIO . atomically . writeTChan (wsChan ws))

awaitNToMMatchesResult ::
  (HasCallStack) =>
  -- | Minimum number of matches
  Int ->
  -- | Maximum number of matches
  Int ->
  -- | Selection function. Exceptions are *not* caught.
  (Value -> App Bool) ->
  WebSocket ->
  App AwaitAtLeastResult
awaitNToMMatchesResult nMin nMax checkMatch ws = go 0 [] []
  where
    go nSeen nonMatches matches = do
      tSecs <- asks timeOutSeconds
      mEvent <- awaitAnyEvent tSecs ws
      case mEvent of
        Just event ->
          do
            isMatch <- checkMatch event
            if isMatch
              then go (nSeen + 1) nonMatches (event : matches)
              else go nSeen (event : nonMatches) matches
        Nothing -> do
          refill nonMatches
          pure $
            AwaitAtLeastResult
              { success = nMin <= nSeen && nSeen <= nMax,
                nMatchesExpectedMin = nMin,
                nMatchesExpectedMax = pure nMax,
                matches = reverse matches,
                nonMatches = reverse nonMatches
              }
    refill = mapM_ (liftIO . atomically . writeTChan (wsChan ws))

awaitNMatches ::
  (HasCallStack) =>
  -- | Number of matches
  Int ->
  -- | Selection function. Should not throw any exceptions
  (Value -> App Bool) ->
  WebSocket ->
  App [Value]
awaitNMatches nExpected checkMatch ws = do
  res <- awaitNMatchesResult nExpected checkMatch ws
  withWebSocketFailureContext ws $
    assertAwaitResult res

withWebSocketFailureContext :: WebSocket -> App a -> App a
withWebSocketFailureContext ws = addFailureContext ("on websocket for user: " <> ws.wsConnect.user <> "@" <> ws.wsConnect.domain)

assertAwaitResult :: (HasCallStack) => AwaitResult -> App [Value]
assertAwaitResult res = do
  if res.success
    then pure res.matches
    else do
      let msgHeader = "Expected " <> show res.nMatchesExpected <> " matching events, but got " <> show (length res.matches) <> "."
      details <- ("Details:\n" <>) <$> prettyAwaitResult res
      assertFailure $ unlines [msgHeader, details]

awaitAtLeastNMatches ::
  (HasCallStack) =>
  -- | Minumum number of matches
  Int ->
  -- | Selection function. Should not throw any exceptions
  (Value -> App Bool) ->
  WebSocket ->
  App [Value]
awaitAtLeastNMatches nExpected checkMatch ws = do
  res <- awaitAtLeastNMatchesResult nExpected checkMatch ws
  if res.success
    then pure res.matches
    else do
      let msgHeader = "Expected " <> show nExpected <> " matching events, but got " <> show (length res.matches) <> "."
      details <- ("Details:\n" <>) <$> prettyAwaitAtLeastResult res
      assertFailure $ unlines [msgHeader, details]

awaitNToMMatches ::
  (HasCallStack) =>
  -- | Minimum Number of matches
  Int ->
  -- | Maximum Number of matches
  Int ->
  -- | Selection function. Should not throw any exceptions
  (Value -> App Bool) ->
  WebSocket ->
  App [Value]
awaitNToMMatches nMin nMax checkMatch ws = do
  res <- awaitNToMMatchesResult nMin nMax checkMatch ws
  if res.success
    then pure res.matches
    else do
      let msgHeader = "Expected between" <> show nMin <> " to " <> show nMax <> " matching events, but got " <> show (length res.matches) <> "."
      details <- ("Details:\n" <>) <$> prettyAwaitAtLeastResult res
      assertFailure $ unlines [msgHeader, details]

awaitMatch ::
  (HasCallStack) =>
  -- | Selection function. Should not throw any exceptions
  (Value -> App Bool) ->
  WebSocket ->
  App Value
awaitMatch checkMatch ws = head <$> awaitNMatches 1 checkMatch ws

awaitMatchFor ::
  (HasCallStack) =>
  Int ->
  (Value -> App Bool) ->
  WebSocket ->
  App Value
awaitMatchFor tSecs checkMatch ws = do
  res <- awaitNMatchesResultFor tSecs 1 checkMatch ws
  withWebSocketFailureContext ws $ head <$> assertAwaitResult res

assertNoEvent ::
  (HasCallStack) =>
  Int ->
  WebSocket ->
  App ()
assertNoEvent to ws = withWebSocketFailureContext ws $ do
  mEvent <- awaitAnyEvent to ws
  case mEvent of
    Just event -> assertFailure $ "Expected no event, but got: " <> show event
    Nothing -> pure ()

nPayload :: (MakesValue a) => a -> App Value
nPayload event = do
  payloads <- event %. "payload" & asList
  assertOne payloads

-- | waits for an http response to satisfy a predicate
waitForResponse :: (HasCallStack) => App Response -> (Response -> App r) -> App r
waitForResponse act p = do
  tSecs <- asks timeOutSeconds
  r <- withRunInIO \inIO ->
    timeout (1000 * 1000 * tSecs) do
      let go = do
            inIO (bindResponse act p) `E.catch` \(_ :: AssertionFailure) -> do
              threadDelay 1000
              go
      go
  let err = unwords ["Expected event didn't come true after", show tSecs, "seconds"]
  maybe (assertFailure err) pure r
