Client.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

-- | Test runner and lifecycle management for Hegel.
--
-- This module implements the client-side logic for running property-based
-- tests against a Hegel server. It manages:
--
-- * Test lifecycle (run_test, test_case events, mark_complete)
-- * Helper functions (assume, note, target, generate_from_schema)
-- * Origin extraction for error reporting
module Hegel.Client
  ( -- * Exceptions
    AssumeRejected (..)
  , DataExhausted (..)
  , HegelTestFailure (..)

    -- * Health checks
  , HealthCheck (..)
  , healthCheckToString

    -- * Verbosity and database settings
  , Verbosity (..)
  , Database (..)
  , Settings
      ( settingsTestCases
      , settingsVerbosity
      , settingsSeed
      , settingsDerandomize
      , settingsDatabase
      , settingsSuppressHealthCheck
      )
  , defaultSettings

    -- * Test case state
  , TestCase (..)

    -- * Client
  , HegelClient (..)
  , createClient

    -- * Test operations
  , generateFromSchema
  , assume
  , assert
  , failure
  , note
  , target
  , startSpan
  , stopSpan
  , runTestCase
  , runTest
  ) where

import Codec.CBOR.Term (Term (..))
import Control.Concurrent.MVar
import Control.Exception hiding (assert)
import Control.Monad (replicateM_, when, unless, void)
import Data.Default (Default (def))
import Data.IORef
import Data.Text (Text)
import qualified Data.Text as T
import GHC.Stack
  ( CallStack
  , HasCallStack
  , SrcLoc (..)
  , callStack
  , getCallStack
  , prettyCallStack
  , withFrozenCallStack
  )
import System.IO (hPutStrLn, stderr)

import Hegel.Connection
import Hegel.Protocol (encodeTerm, extractInt, extractMap, extractText)


-- ---------------------------------------------------------------------------
-- Exceptions
-- ---------------------------------------------------------------------------

-- | Raised when 'assume' is called with a 'False' condition.
data AssumeRejected = AssumeRejected
  deriving (Show)

instance Exception AssumeRejected

-- | Raised when the server runs out of test data (StopTest).
data DataExhausted = DataExhausted
  deriving (Show)

instance Exception DataExhausted

-- | Raised when a property test fails.
newtype HegelTestFailure = HegelTestFailure String
  deriving (Show)

instance Exception HegelTestFailure

-- | Raised by Hegel's assertion helpers. Carries a call stack so we can report
-- a stable source location to the server while still surfacing the full error
-- message on the final replayed failure.
data AssertionFailure = AssertionFailure
  { assertionFailureMessage :: String
  , assertionFailureCallStack :: CallStack
  }

instance Show AssertionFailure where
  show AssertionFailure {..}
    | null (getCallStack assertionFailureCallStack) = assertionFailureMessage
    | otherwise =
        assertionFailureMessage
          ++ "\n"
          ++ prettyCallStack assertionFailureCallStack

instance Exception AssertionFailure where
  displayException = show


-- ---------------------------------------------------------------------------
-- Health checks
-- ---------------------------------------------------------------------------

-- | Health checks that can be suppressed during test execution.
data HealthCheck
  = FilterTooMuch
  | TooSlow
  | TestCasesTooLarge
  | LargeInitialTestCase
  deriving (Show, Eq)

-- | Returns the wire protocol name for a health check.
healthCheckToString :: HealthCheck -> Text
healthCheckToString FilterTooMuch = "filter_too_much"
healthCheckToString TooSlow = "too_slow"
healthCheckToString TestCasesTooLarge = "test_cases_too_large"
healthCheckToString LargeInitialTestCase = "large_initial_test_case"


-- ---------------------------------------------------------------------------
-- Verbosity and database settings
-- ---------------------------------------------------------------------------

-- | Controls how much output Hegel produces during test runs.
data Verbosity = Quiet | Normal | Verbose | Debug
  deriving (Show, Eq)

-- | The database setting: unset, disabled, or a path.
data Database
  = Unset
  | Disabled
  | DatabasePath FilePath
  deriving (Show, Eq)

-- | Configuration for a Hegel test run.
data Settings = Settings
  { settingsTestCases :: Int
  , settingsVerbosity :: Verbosity
  , settingsSeed :: Maybe Int
  , settingsDerandomize :: Bool
  , settingsDatabase :: Database
  , settingsSuppressHealthCheck :: [HealthCheck]
  } deriving (Show, Eq)

instance Default Settings where
  def = Settings
    { settingsTestCases = 100
    , settingsVerbosity = Normal
    , settingsSeed = Nothing
    , settingsDerandomize = False
    , settingsDatabase = Unset
    , settingsSuppressHealthCheck = []
    }


-- | Pure baseline settings before any runtime-specific overrides are applied.
defaultSettings :: Settings
defaultSettings = def


-- ---------------------------------------------------------------------------
-- Test case state
-- ---------------------------------------------------------------------------

-- | Per-test-case state passed to the test function.
data TestCase = TestCase
  { tcChannel :: Channel
    -- ^ The data channel for this test case.
  , tcIsFinal :: Bool
    -- ^ Whether this is the final (replay) run.
  , tcTestAborted :: IORef Bool
    -- ^ Set to 'True' when the server signals StopTest.
  }


-- ---------------------------------------------------------------------------
-- Client
-- ---------------------------------------------------------------------------

-- | Supported protocol version range.
supportedProtocolLo :: Double
supportedProtocolLo = 0.6

supportedProtocolHi :: Double
supportedProtocolHi = 0.7

-- | Hegel client for running property-based tests.
data HegelClient = HegelClient
  { clientConnection :: Connection
  , clientControl :: Channel
  , clientLock :: MVar ()
  }

-- | Creates a new client from a connection. The connection must not yet have
-- had its handshake performed. Validates the protocol version is in range
-- 0.1 through 0.7.
createClient :: Connection -> IO HegelClient
createClient conn = do
  versionStr <- sendHandshake conn
  let serverVersion = read versionStr :: Double
  when (serverVersion < supportedProtocolLo || serverVersion > supportedProtocolHi) $
    fail $ "hegel-haskell supports protocol versions "
      ++ show supportedProtocolLo ++ " through " ++ show supportedProtocolHi
      ++ ", but got server version " ++ show serverVersion
      ++ ". Upgrading hegel-haskell or downgrading your hegel cli might help."
  ctrl <- controlChannel conn
  lock <- newMVar ()
  pure HegelClient
    { clientConnection = conn
    , clientControl = ctrl
    , clientLock = lock
    }


-- ---------------------------------------------------------------------------
-- Test operations
-- ---------------------------------------------------------------------------

-- | Generates a value from a schema by sending a generate command to the
-- server. Raises 'DataExhausted' if the server signals StopTest.
generateFromSchema :: Term -> TestCase -> IO Term
generateFromSchema schema tc = do
  let ch = tcChannel tc
  let msg = TMap
        [ (TString "command", TString "generate")
        , (TString "schema", schema)
        ]
  result <- try (sendRequestAndWait ch msg) :: IO (Either RequestError Term)
  case result of
    Right val -> pure val
    Left (RequestError _ errType _)
      | errType == "StopTest" -> do
          writeIORef (tcTestAborted tc) True
          throwIO DataExhausted
    Left e -> throwIO e

-- | Rejects the current test case if the condition is 'False'.
assume :: TestCase -> Bool -> IO ()
assume _tc condition = unless condition $ throwIO AssumeRejected

-- | Fails the current property with a message and captured call stack.
failure :: HasCallStack => String -> IO a
failure message =
  withFrozenCallStack $
    throwIO
      AssertionFailure
        { assertionFailureMessage = message
        , assertionFailureCallStack = callStack
        }

-- | Asserts a property condition, reporting the call site on failure.
assert :: HasCallStack => Bool -> String -> IO ()
assert condition message = unless condition (withFrozenCallStack $ failure message)

-- | Records a message that will be printed on the final (failing) run.
note :: TestCase -> String -> IO ()
note tc message = when (tcIsFinal tc) $ hPutStrLn stderr message

-- | Sends a target command to guide the search engine toward higher values.
target :: TestCase -> Double -> String -> IO ()
target tc value label = do
  let ch = tcChannel tc
  let msg = TMap
        [ (TString "command", TString "target")
        , (TString "value", TDouble value)
        , (TString "label", TString (T.pack label))
        ]
  void $ sendRequestAndWait ch msg

-- | Starts a generation span for better shrinking.
startSpan :: Int -> TestCase -> IO ()
startSpan label tc = do
  aborted <- readIORef (tcTestAborted tc)
  unless aborted $ do
    let ch = tcChannel tc
    let msg = TMap
          [ (TString "command", TString "start_span")
          , (TString "label", TInt label)
          ]
    void $ sendRequestAndWait ch msg

-- | Ends the current generation span.
stopSpan :: Bool -> TestCase -> IO ()
stopSpan discard tc = do
  aborted <- readIORef (tcTestAborted tc)
  unless aborted $ do
    let ch = tcChannel tc
    let msg = TMap
          [ (TString "command", TString "stop_span")
          , (TString "discard", TBool discard)
          ]
    void $ sendRequestAndWait ch msg


-- ---------------------------------------------------------------------------
-- Test case outcomes
-- ---------------------------------------------------------------------------

data TestOutcome
  = Valid
  | Invalid
  | DataWasExhausted
  | Interesting Text (Maybe SomeException)
    -- ^ origin text, and optionally the exception (on final run)

-- | Extracts a human-readable origin string from an exception.
extractOrigin :: SomeException -> Text
extractOrigin exn =
  case fromException exn of
    Just AssertionFailure {assertionFailureCallStack = stack} ->
      maybe fallback (T.pack . formatSrcLoc) (firstSrcLoc stack)
    Nothing -> fallback
  where
    fallback = T.pack (displayException exn)

firstSrcLoc :: CallStack -> Maybe SrcLoc
firstSrcLoc stack = snd <$> case getCallStack stack of
  [] -> Nothing
  frame : _ -> Just frame

formatSrcLoc :: SrcLoc -> String
formatSrcLoc SrcLoc {..} =
  srcLocFile ++ ":" ++ show srcLocStartLine ++ ":" ++ show srcLocStartCol

-- | Send mark_complete, swallowing StopTest errors.
markComplete :: Channel -> Text -> Term -> IO ()
markComplete ch status origin = do
  let msg = TMap
        [ (TString "command", TString "mark_complete")
        , (TString "status", TString status)
        , (TString "origin", origin)
        ]
  result <- try (sendRequestAndWait ch msg) :: IO (Either RequestError Term)
  case result of
    Left (RequestError _ errType _)
      | errType == "StopTest" -> pure ()
    Left e -> throwIO e
    Right _ -> pure ()

-- | Runs a single test case. Creates a 'TestCase', runs the user function,
-- catches exceptions, and sends mark_complete with the appropriate status
-- (VALID, INVALID, or INTERESTING). Handles StopTest on mark_complete.
runTestCase :: HegelClient -> Channel -> (TestCase -> IO ()) -> Bool -> IO ()
runTestCase _client ch testFn isFinal = do
  abortedRef <- newIORef False
  let tc = TestCase
        { tcChannel = ch
        , tcIsFinal = isFinal
        , tcTestAborted = abortedRef
        }
  outcome <- (testFn tc >> pure Valid) `catches`
    [ Handler $ \AssumeRejected -> pure Invalid
    , Handler $ \DataExhausted -> pure DataWasExhausted
    , Handler $ \(e :: SomeException) -> do
        let origin = extractOrigin e
        let ex = if isFinal then Just e else Nothing
        pure (Interesting origin ex)
    ]
  case outcome of
    DataWasExhausted -> pure ()
    Valid -> markComplete ch "VALID" TNull
    Invalid -> markComplete ch "INVALID" TNull
    Interesting originText _ -> markComplete ch "INTERESTING" (TString originText)
  closeChannel ch
  -- On final run, re-raise the interesting exception
  case outcome of
    Interesting _ (Just e) -> throwIO e
    _ -> pure ()


-- ---------------------------------------------------------------------------
-- Full test lifecycle
-- ---------------------------------------------------------------------------

-- | Runs a full property test. Sends run_test on the control channel with
-- all settings fields, then processes test_case and test_done events.
-- Checks error, health_check_failure, flaky, and passed in results.
-- Replays interesting cases on the final run.
runTest :: HegelClient -> Settings -> (TestCase -> IO ()) -> IO ()
runTest client stgs testFn = do
  let conn = clientConnection client
  testChannel <- newChannel conn

  -- Send run_test command under the client lock
  withMVar (clientLock client) $ \_ -> do
    let seedValue = maybe TNull TInt (settingsSeed stgs)
    let databaseKeyValue = TNull
    let baseFields =
          [ (TString "command", TString "run_test")
          , (TString "test_cases", TInt (settingsTestCases stgs))
          , (TString "seed", seedValue)
          , (TString "channel_id", TInt (fromIntegral (channelId testChannel)))
          , (TString "database_key", databaseKeyValue)
          , (TString "derandomize", TBool (settingsDerandomize stgs))
          ]
    let databaseField = case settingsDatabase stgs of
          Unset          -> []
          Disabled       -> [(TString "database", TNull)]
          DatabasePath p -> [(TString "database", TString (T.pack p))]
    let suppressField = case settingsSuppressHealthCheck stgs of
          [] -> []
          checks ->
            [ ( TString "suppress_health_check"
              , TList (map (TString . healthCheckToString) checks)
              )
            ]
    let fields = baseFields ++ databaseField ++ suppressField
    void $ sendRequestAndWait (clientControl client) (TMap fields)

  -- Helper to receive and run a test case (used for replay)
  let receiveAndRunTestCase isFinal = do
        (msgId, msg) <- receiveRequest testChannel
        let pairs = extractMap msg
        let chId = extractInt (lookupCBOR "channel_id" pairs)
        sendResponseValue testChannel msgId TNull
        tcChan <- connectChannel conn (fromIntegral chId)
        runTestCase client tcChan testFn isFinal

  -- Event loop
  let receiveEvents = do
        (msgId, msg) <- receiveRequest testChannel
        let pairs = extractMap msg
        let event = extractText (lookupCBOR "event" pairs)
        case event of
          "test_case" -> do
            let chId = extractInt (lookupCBOR "channel_id" pairs)
            sendResponseValue testChannel msgId TNull
            tcChan <- connectChannel conn (fromIntegral chId)
            runTestCase client tcChan testFn False
            exited <- serverHasExited conn
            when exited $ fail serverCrashedMessage
            receiveEvents
          "test_done" -> do
            sendResponseValue testChannel msgId (TBool True)
            pure (extractMap (lookupCBOR "results" pairs))
          other -> do
            let errPayload = encodeTerm $ TMap
                  [ (TString "error", TString ("Unrecognised event " <> other))
                  , (TString "type", TString "InvalidMessage")
                  ]
            sendResponseRaw testChannel msgId errPayload
            receiveEvents

  results <- receiveEvents

  -- Check for server-side errors
  case lookupCBORMaybe "error" results of
    Just errorVal -> do
      let errorMsg = extractText errorVal
      fail ("Server error: " ++ T.unpack errorMsg)
    Nothing -> pure ()

  -- Check for health check failure
  case lookupCBORMaybe "health_check_failure" results of
    Just failureVal -> do
      let failureMsg = extractText failureVal
      fail ("Health check failure:\n" ++ T.unpack failureMsg)
    Nothing -> pure ()

  -- Check for flaky test detection
  case lookupCBORMaybe "flaky" results of
    Just flakyVal -> do
      let flakyMsg = extractText flakyVal
      fail ("Flaky test detected: " ++ T.unpack flakyMsg)
    Nothing -> pure ()

  -- Check passed flag
  let passed = case lookupCBORMaybe "passed" results of
        Just (TBool b) -> b
        _              -> True

  let nInteresting = extractInt (lookupCBOR "interesting_test_cases" results)

  if nInteresting == 0 && passed
    then pure ()
    else if nInteresting >= 1
      then do
        firstFailure <- try (receiveAndRunTestCase True)
        -- Drain any extra interesting cases quietly to keep the control
        -- channel in sync, but only surface the smallest final failure.
        replicateM_ (nInteresting - 1) (receiveAndRunTestCase False)
        case firstFailure of
          Left (e :: SomeException) -> throwIO e
          Right () -> unless passed $ throwIO (HegelTestFailure "Property test failed")
      else unless passed $ throwIO (HegelTestFailure "Property test failed")


-- ---------------------------------------------------------------------------
-- CBOR helpers (local to this module)
-- ---------------------------------------------------------------------------

-- | Look up a key in a CBOR map (list of pairs), using a 'Text' key.
lookupCBOR :: Text -> [(Term, Term)] -> Term
lookupCBOR key pairs =
  case lookup (TString key) pairs of
    Just v  -> v
    Nothing -> error $ "Key not found in CBOR map: " ++ T.unpack key

-- | Look up a key in a CBOR map, returning 'Nothing' if absent.
lookupCBORMaybe :: Text -> [(Term, Term)] -> Maybe Term
lookupCBORMaybe key = lookup (TString key)