{-# LANGUAGE OverloadedStrings #-}

module Test.Driver.Dialogue.Generation (
    Dialogue(..)
  , dialogueGlobalSteps
  , DialogueWithoutExceptions(..)
  , DialogueWithExceptions(..)
  , ensureCorrectUsage
  ) where

import Control.Monad
import Data.ByteString qualified as BS.Strict
import Data.ByteString qualified as Strict (ByteString)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe)
import GHC.Stack
import Test.QuickCheck

import Network.GRPC.Common

import Test.Driver.Dialogue.Definition
import Test.Driver.Dialogue.TestClock qualified as TestClock

{-------------------------------------------------------------------------------
  Metadata

  We don't (usually) want to have to look at tests failures with unnecessarily
  complicated names or values, this simply distracts. The only case where we
  /do/ want that is if we are specifically testing parsers/builders; we do this
  in separate serialization tests ("Test.Prop.Serialization").
-------------------------------------------------------------------------------}

simpleAsciiValue :: [Strict.ByteString]
simpleAsciiValue = ["a", "b", "c"]

genMetadata :: Gen TestMetadata
genMetadata =
    TestMetadata
      <$> liftArbitrary genAsciiValue
      <*> liftArbitrary genAsciiValue
      <*> liftArbitrary genBinaryValue
      <*> liftArbitrary genBinaryValue
  where
    genAsciiValue :: Gen Strict.ByteString
    genAsciiValue = elements simpleAsciiValue

    genBinaryValue :: Gen Strict.ByteString
    genBinaryValue = sized $ \sz -> do
        n <- choose (0, sz)
        BS.Strict.pack <$> replicateM n arbitrary

{-------------------------------------------------------------------------------
  Local steps
-------------------------------------------------------------------------------}

-- | Generate 'LocalStep's
--
-- We do not insert timing here; we do this globally.
genLocalSteps ::
     Bool  -- ^ Should we generate exceptions?
  -> Gen [LocalStep]
genLocalSteps genExceptions = sized $ \sz -> do
    n <- choose (0, sz)
    (:) <$> genInitStep <*> replicateM n genStep
  where
    genInitStep :: Gen LocalStep
    genInitStep = do
        metadata <- genMetadata
        rpc      <- elements [RPC1, RPC2, RPC3]
        return $ ClientAction $ Initiate (metadata, rpc)

    genStep :: Gen LocalStep
    genStep = frequency [
          (if genExceptions then 1 else 0, genException)
        , (3, ServerAction . Initiate <$> genMetadata)
        , (3, ClientAction . Send     <$> genElem (pure NoMetadata))
        , (3, ServerAction . Send     <$> genElem genMetadata)
        ]

    genException :: Gen LocalStep
    genException = oneof [
          ClientAction . Terminate . Just . SomeClientException <$> choose (0, 5)
        , ServerAction . Terminate . Just . SomeServerException <$> choose (0, 5)
        , pure $ ClientAction . Terminate $ Nothing
        , pure $ ServerAction . Terminate $ Nothing
        ]

    genElem :: Gen b -> Gen (StreamElem b Int)
    genElem genTrailers = oneof [
           StreamElem  <$> genMsg
         , FinalElem   <$> genMsg <*> genTrailers
         , NoMoreElems <$> genTrailers
        ]

    genMsg :: Gen Int
    genMsg = choose (0, 99)

{-------------------------------------------------------------------------------
  Ensure correct library usage

  We have two essentially different kinds of tests: correctness of the library
  given correct library usage, and reasonable error reporting given incorrect
  library usage. We focus on the former here. For example:

  * We /are/ interested in testing what happens when a client sends a message
    to the server, but the server handler threw an exception.
  * We are /not/ interested in testing what happens when a client tries to send
    another message after having told the server that they sent their last
    message, or after the server told the client that the call is over.

  We also want to make sure the tests are not non-sensical (for example, it does
  not make sense for the client to send a message after it has terminated).
-------------------------------------------------------------------------------}

-- | State of a single RPC call
--
-- We use this to determine when certain actions can happen.
--
-- Invariants:
--
-- * The server response cannot be initiated until the client request has been
--   (until that time the server handler is not even running).
-- * The client request must be closed when the server response closes.
--
-- That second point is a bit subtle. Normally it is the responsibility of both
-- the client and the server to indicate when they won't send any more messages;
-- if they do not, an exception is raised. There is however one exception to
-- this rule: the server can unilaterally decide to close the entire RPC. When
-- this happens, the client (of course) does not have to send any more messages.
-- (Under normal circumstance this does not happen: the server would not close
-- the RPC until the client has closed their end.) We treat the case where the
-- server throws an exception the same: in both cases the RPC is closed and the
-- client does not need to send its final message.
--
-- Of course, this /does/ mean that the client needs to /notice/ that the server
-- has closed the call, even when it's an exception. We therefore implement a
-- 'Terminate' on one side as a receive on the other (see 'clientLocal' and
-- 'serverLocal').
--
-- One consequence of all this worth spelling out once more is that we rule
-- out tests such as
--
-- > NormalizedDialogue [
-- >    (0, ClientAction $ Initiate (def, RPC1))
-- >  , (0, ServerAction $ Terminate (Just $ SomeServerException 0))
-- >  , (0, ServerAction $ Send (StreamElem 1))
-- >  ]
--
-- where the server sends a message after throwing an exception; at one level,
-- it may seem obvious that we must rule such tests out (it doesn't make sense
-- for the server to send messages after terminating, after all). However, a
-- server 'Send' action is also a client @recv@ action, so one could think that
-- we need such tests to make sure that clients get an appropriate exception
-- when they try to receive a message but the server aborted. However, this case
-- is dealt with by the fact that that we treat a 'Terminate' on one side as a
-- receive on the other, as discussed above. This means that in the tests, a
-- 'Send' on one side /must/ be a /successful/ receive on the other.
data LocalGenState = LocalGenState {
      localGenClient :: LocalUniState
    , localGenServer :: LocalUniState
    }
  deriving (Show)

-- | Unidirectional state
--
-- The tests describe a dialogue in a one-sided manner: we talk about /sending/,
-- but not about receiving. In a way \"sending\" is something you /do/, whereas
-- \"receiving\" is something that's done /to/ you. The state we record here is
-- therefore the /outbound/ state only.
data LocalUniState =
    -- | The outbound stream has yet been initialized
    --
    -- Nothing can happen until the client initiates the request. The server can
    -- choose to initiate the response at any point during the conversation.
    UniUninit

    -- | The outbound stream has been established
  | UniOpen

    -- | The outbound stream has been closed
    --
    -- In the case of the client, closing the stream (cleanly) is referred to
    -- in gRPC documentation as \"putting the call in half-closed state\".
    -- In the case of the server, closing the stream involves sending the
    -- trailers and signals the end of the RPC.
    --
    -- The stream is closed \"uncleanly\" if the client or server handler
    -- simply disappears, or when it throws an exception.
  | UniClosed
  deriving (Show)

initLocalGenState :: LocalGenState
initLocalGenState = LocalGenState {
      localGenClient = UniUninit
    , localGenServer = UniUninit
    }

establishInvariant :: HasCallStack => LocalGenState -> LocalGenState
establishInvariant st =
    case (localGenClient st, localGenServer st) of
      (UniUninit , UniUninit) -> st
      (UniUninit , _)         -> error "Response initiated before request"
      (_         , UniClosed) -> st { localGenClient = UniClosed }
      _otherwise              -> st

ensureCorrectUsage :: [(Int, LocalStep)] -> [(Int, LocalStep)]
ensureCorrectUsage = go Map.empty []
  where
    go ::
         Map Int LocalGenState  -- State of each concurrent conversation
      -> [(Int, LocalStep)]     -- Accumulator (reverse order)
      -> [(Int, LocalStep)]     -- Still to consider
      -> [(Int, LocalStep)]     -- Result
    go sts acc [] = concat [
          reverse acc
        , concatMap (\(i, st) -> (i,) <$> ensureCleanClose st) $ Map.toList sts
        ]
      where
        -- Make sure all channels are closed cleanly
        --
        -- We could do this in two ways: we can either have the server close
        -- the call unilaterally, or we could first have the client close their
        -- end and then the server its own. We opt for the latter, as it is the
        -- more "clean" one, but if a particular test case is generated in which
        -- the client does not close its own end before the server does, that is
        -- also ok (see further discussion in 'LocalGenState').
        ensureCleanClose :: LocalGenState -> [LocalStep]
        ensureCleanClose st = concat [
              [ ClientAction $ Send $ NoMoreElems NoMetadata
              | case localGenClient st of
                  UniUninit -> False
                  UniOpen   -> True
                  UniClosed -> False
              ]
            , [ ServerAction $ Send $ NoMoreElems def
              | case (localGenClient st, localGenServer st) of
                  (UniUninit, _) -> False
                  (_, UniUninit) -> True
                  (_, UniOpen)   -> True
                  (_, UniClosed) -> False
              ]
            ]
    go sts acc ((i, s):ss) =
        case s of
          ClientAction Initiate{} ->
            case localGenClient st of
              UniUninit  -> contWith $ updClient UniOpen
              _otherwise -> skip

          ClientAction Terminate{} ->
            case localGenClient st of
              UniUninit  -> skip
              UniOpen    -> contWith $ updClient UniClosed
              UniClosed  -> skip

          ClientAction (Send (StreamElem _)) ->
            case localGenClient st of
              UniUninit  -> insert $ ClientAction (Initiate (def, RPC1))
              UniOpen    -> contWith $ id
              UniClosed  -> skip

          ClientAction (Send _) -> -- FinalElem or NoMoreElems
            case localGenClient st of
              UniUninit  -> insert (ClientAction (Initiate (def, RPC1)))
              UniOpen    -> contWith $ updClient UniClosed
              UniClosed  -> skip

          -- The server cannot do anything until the request is initiated
          -- (until that point the server handler is not even running)
          ServerAction _ | UniUninit <- localGenClient st ->
            skip

          ServerAction Initiate{} ->
            case localGenServer st of
             UniUninit   -> contWith $ updServer UniOpen
             _otherwise  -> skip

          ServerAction (Send (StreamElem _)) ->
            case localGenServer st of
              UniUninit  -> contWith $ updServer UniOpen -- implicitly opened
              UniOpen    -> contWith $ id
              UniClosed  -> skip

          ServerAction (Send _) -> -- FinalElem or NoMoreElems
            case localGenServer st of
              UniClosed  -> skip
              _otherwise -> contWith $ updServer UniClosed

          ServerAction Terminate{} ->
            case localGenServer st of
              UniClosed  -> skip
              _otherwise -> contWith $ updServer UniClosed
      where
        --
        -- Three different ways to continue
        --

        -- 1. Normal case: update the state and move on the next action
        contWith ::
             (Map Int LocalGenState -> Map Int LocalGenState)
          -> [(Int, LocalStep)]
        contWith f = go (Map.map establishInvariant $ f sts) ((i, s) : acc) ss

        -- 2. Skip this action
        skip :: [(Int, LocalStep)]
        skip = go sts acc ss

        -- 3. First execute a different action
        insert :: LocalStep -> [(Int, LocalStep)]
        insert newStep = go sts acc ((i, newStep) : (i, s) : ss)

        --
        -- Updating the state
        --

        st :: LocalGenState
        st = Map.findWithDefault initLocalGenState i sts

        updClient ::
             LocalUniState
          -> Map Int LocalGenState -> Map Int LocalGenState
        updClient client' =
            Map.alter (Just . aux . fromMaybe initLocalGenState) i
          where
            aux :: LocalGenState -> LocalGenState
            aux st' = st' { localGenClient = client' }

        updServer ::
             LocalUniState
          -> Map Int LocalGenState -> Map Int LocalGenState
        updServer server' =
            Map.alter (Just . aux . fromMaybe initLocalGenState) i
          where
            aux :: LocalGenState -> LocalGenState
            aux st' = st' { localGenServer = server' }

{-------------------------------------------------------------------------------
  Dialogue

  We generate the steps for each channel separately first, and then choose a
  particular interleaving ('interleave'). Before execution, we then assign
  timings ('assignTimings'), separating the channels again. The advantage of
  this representation is that it is easier to shrink: we can shrink the /global/
  step of steps, whilst keeping the choice of interleaving; that is much harder
  to do when the channels are kept separate.

  In addition, we apply 'ensureCorrectUsage' just before execution. We do /not/
  do this as part of generation/shrinking, as doing so could shrinking
  non-wellfounded (shrinking might remove some action which gets reinserted by
  'ensureCorrectUsage', and shrinking will loop).
-------------------------------------------------------------------------------}

data Dialogue =
    -- | Dialogue that is already in "correct" form
    --
    -- Calling 'ensureCorrectUsage' on such a dialogue should be a no-op.
    NormalizedDialogue [(Int, LocalStep)]

    -- | Dialogue that may still need to be corrected
  | UnnormalizedDialogue [(Int, LocalStep)]
  deriving stock (Show, Eq)

getNormalizedDialogue :: HasCallStack => Dialogue -> [(Int, LocalStep)]
getNormalizedDialogue (UnnormalizedDialogue steps) = ensureCorrectUsage steps
getNormalizedDialogue (NormalizedDialogue   steps)
  | steps == normalized
  = steps

  | otherwise
  = error $ "getNormalizedDialogue: normal form is " ++ show normalized
  where
    normalized :: [(Int, LocalStep)]
    normalized = ensureCorrectUsage steps

getRawDialogue :: Dialogue -> [(Int, LocalStep)]
getRawDialogue (NormalizedDialogue   steps) = steps
getRawDialogue (UnnormalizedDialogue steps) = steps

dialogueGlobalSteps :: Dialogue -> GlobalSteps
dialogueGlobalSteps =
      GlobalSteps
    . map LocalSteps
    . TestClock.assignTimings
    . getNormalizedDialogue

-- | Shrink dialogue
shrinkDialogue :: Dialogue -> [Dialogue]
shrinkDialogue =
      map UnnormalizedDialogue
    . shrinkList (shrinkInterleaved shrinkLocalStep)
    . getRawDialogue

{-------------------------------------------------------------------------------
  Shrinking
-------------------------------------------------------------------------------}

shrinkLocalStep :: LocalStep -> [LocalStep]
shrinkLocalStep = \case
    ClientAction (Initiate (metadata, rpc)) -> concat [
        [ ClientAction $ Initiate (metadata', rpc)
        | metadata' <- shrinkMetadata metadata
        ]
      , [ ClientAction $ Initiate (metadata, rpc')
        | rpc' <- shrinkRPC rpc
        ]
      ]
    ServerAction (Initiate metadata) ->
      map (ServerAction . Initiate) $ shrinkMetadata metadata
    ClientAction (Send x) ->
      map (ClientAction . Send) $ shrinkElem (const []) x
    ServerAction (Send x) ->
      map (ServerAction . Send) $ shrinkElem shrinkMetadata x
    ClientAction (Terminate (Just (SomeClientException n))) ->
      map (ClientAction . Terminate . Just . SomeClientException) (shrink n)
    ServerAction (Terminate (Just (SomeServerException n))) ->
      map (ServerAction . Terminate . Just . SomeServerException) (shrink n)
    ClientAction (Terminate Nothing) ->
      []
    ServerAction (Terminate Nothing) ->
      []

shrinkRPC :: RPC -> [RPC]
shrinkRPC RPC1 = []
shrinkRPC RPC2 = [RPC1]
shrinkRPC RPC3 = [RPC1, RPC2]

-- | Shrink metadata
--
-- For now we just shrink individual values. In principle we could also try to
-- /replace/ a binary header with an ASCII one.
shrinkMetadata :: TestMetadata -> [TestMetadata]
shrinkMetadata md = concat [
      [ md{metadataAsc1 = x}
      | x <- liftShrink shrinkAsciiValue (metadataAsc1 md)
      ]
    , [ md{metadataAsc2 = x}
      | x <- liftShrink shrinkAsciiValue (metadataAsc2 md)
      ]
    , [ md{metadataBin3 = x}
      | x <- liftShrink shrinkBinaryValue (metadataBin3 md)
      ]
    , [ md{metadataBin4 = x}
      | x <- liftShrink shrinkBinaryValue (metadataBin4 md)
      ]
    ]
  where
    shrinkAsciiValue :: Strict.ByteString -> [Strict.ByteString]
    shrinkAsciiValue val = filter (< val) simpleAsciiValue

    shrinkBinaryValue :: Strict.ByteString -> [Strict.ByteString]
    shrinkBinaryValue val = concat [
          -- aggressively try to shrink to a single byte
          [ BS.Strict.pack [x]
          | x:_:_ <- [BS.Strict.unpack val]
          ]

          -- normal shrinking of binary values
        , [ BS.Strict.pack val'
          | val' <- shrink (BS.Strict.unpack val)
          ]
        ]

-- | Shrink element
--
-- For now we don't change the nature of the elem. Not sure what the right
-- definition of "simpler" is here.
shrinkElem :: (a -> [a]) -> StreamElem a Int -> [StreamElem a Int]
shrinkElem _ (StreamElem x) = concat [
      [ StreamElem x'
      | x' <- shrink x
      ]
    ]
shrinkElem f (FinalElem x y) = concat [
      [ FinalElem x' y
      | x' <- shrink x
      ]
    , [ FinalElem x y'
      | y' <- f y
      ]
    ]
shrinkElem f (NoMoreElems y) = concat [
      [ NoMoreElems y'
      | y' <- f y
      ]
    ]

shrinkInterleaved :: (a -> [a]) -> (Int, a) -> [(Int, a)]
shrinkInterleaved f (i, a) = (i, ) <$> f a

{-------------------------------------------------------------------------------
  Arbitrary instance
-------------------------------------------------------------------------------}

newtype DialogueWithoutExceptions = DialogueWithoutExceptions Dialogue
  deriving stock (Show, Eq)

newtype DialogueWithExceptions = DialogueWithExceptions Dialogue
  deriving stock (Show, Eq)

instance Arbitrary DialogueWithoutExceptions where
  arbitrary = do
      concurrency <- choose (1, 3)
      threads     <- replicateM concurrency $ genLocalSteps False
      DialogueWithoutExceptions . UnnormalizedDialogue <$>
        TestClock.interleave threads

  shrink (DialogueWithoutExceptions dialogue) =
      DialogueWithoutExceptions <$> shrinkDialogue dialogue

instance Arbitrary DialogueWithExceptions where
  arbitrary = do
      concurrency <- choose (1, 3)
      threads     <- replicateM concurrency $ genLocalSteps True
      DialogueWithExceptions . UnnormalizedDialogue <$>
        TestClock.interleave threads

  shrink (DialogueWithExceptions dialogue) =
      DialogueWithExceptions <$> shrinkDialogue dialogue