WithMock.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

{- | withMock: Declarative mock expectations DSL
-}
module Test.MockCat.WithMock
  ( withMock
  , withMockIO
  , askWithMockContext
  , expects
  , MockResult(..)
  , called
  , with
  , calledInOrder
  , calledInSequence
  , times
  , once
  , never
  , atLeast
  , atMost
  , greaterThan
  , lessThan
  , anything
  , WithMockContext(..)
  , Expectation(..)
  , Expectations(..)
  , verifyExpectationDirect
  ) where

import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Reader (ReaderT(..), runReaderT)
import Control.Exception (bracket_)
import Control.Concurrent.STM (newTVarIO, readTVarIO, atomically, modifyTVar')
import Control.Monad.State (get, put, modify)
import Test.MockCat.Verify (TimesSpec(..), times, once, never, atLeast, atMost, greaterThan, lessThan, anything, ResolvableMock, ResolvableParamsOf)
import Test.MockCat.Internal.Verify (verifyExpectationDirect)
import Test.MockCat.Internal.Types
  ( VerifyOrderMethod(..)
  , WithMockContext(..)
  , Expectation(..)
  , Expectations(..)
  , runExpectations
  , addExpectation
  , InvocationRecorder(..)
  , ResolvedMock(..)
  )
import qualified Test.MockCat.Internal.Registry.Core as MockRegistry
import Test.MockCat.Internal.Registry.Core (getThreadWithMockContext, setThreadWithMockContext, clearThreadWithMockContext)
import Unsafe.Coerce (unsafeCoerce)

import Test.MockCat.Param (Param(..), param, EqParams(..))
import Data.Kind (Type)
import Data.Proxy (Proxy(..))

-- | A specialized Unit type that carries the parameter type information.
--   This is used to improve type inference for unit-returning mock helpers.
newtype MockResult params = MockResult ()
  deriving (Show, Eq)



-- | Run a block with mock expectations that are automatically verified
withMock :: ReaderT WithMockContext IO a -> IO a
withMock action = do
  ctxVar <- newTVarIO []
  let ctx = WithMockContext ctxVar
  bracket_ (setThreadWithMockContext ctx) clearThreadWithMockContext $ do
    result <- runReaderT action ctx
    -- Verify all registered verification actions
    actions <- readTVarIO ctxVar
    sequence_ actions
    pure result

-- | IO version of withMock
withMockIO :: IO a -> IO a
withMockIO action = do
  ctxVar <- newTVarIO []
  let ctx = WithMockContext ctxVar
  bracket_ (setThreadWithMockContext ctx) clearThreadWithMockContext $ do
    result <- action
    -- Verify all registered verification actions
    actions <- readTVarIO ctxVar
    sequence_ actions
    pure result

-- | Retrieve the current mock context from thread-local storage.
--   Throws an error if no context is found.
askWithMockContext :: IO WithMockContext
askWithMockContext = do
  mCtx <- getThreadWithMockContext
  case mCtx of
    Just ctx -> pure ctx
    Nothing -> errorWithoutStackTrace "askWithMockContext: No WithMockContext found in current thread. Use withMock or withMockIO."

-- | Attach expectations to a mock function
--   Supports both single expectation and multiple expectations in a do block
infixl 0 `expects`

-- | Type class to extract params type from an expectation expression
class ExtractParams exp where
  type ExpParams exp :: Type
  extractParams :: exp -> Proxy (ExpParams exp)

instance ExtractParams (Expectations params ()) where
  type ExpParams (Expectations params ()) = params
  extractParams _ = Proxy

instance ExtractParams (fn -> Expectations params ()) where
  type ExpParams (fn -> Expectations params ()) = params
  extractParams _ = Proxy

-- | Register expectations for a mock function
--   Accepts an Expectations builder
--   The params type is usually inferred from the mock function, but can be overridden
class BuildExpectations fn exp params | fn exp -> params where
  buildExpectations :: fn -> exp -> [Expectation params]

-- | Instance for direct Expectations value
instance {-# OVERLAPPABLE #-} forall fn params. (ResolvableParamsOf fn ~ params) => BuildExpectations fn (Expectations params ()) params where
  buildExpectations _ = runExpectations

-- | Instance for function form when fn is MockResult
instance {-# OVERLAPPING #-} forall params. BuildExpectations (MockResult params) (MockResult params -> Expectations params ()) params where
  buildExpectations _ f = runExpectations (f (MockResult ()))

-- | Instance for direct Expectations value when fn is MockResult
instance {-# OVERLAPPING #-} forall params. BuildExpectations (MockResult params) (Expectations params ()) params where
  buildExpectations _ = runExpectations

-- | Instance for direct Expectations value when fn is ()
instance {-# OVERLAPPING #-} forall params. BuildExpectations () (Expectations params ()) params where
  buildExpectations _ = runExpectations

-- | Instance for function form (fn -> Expectations params ())
instance {-# OVERLAPPABLE #-} forall fn params. (ResolvableParamsOf fn ~ params) => BuildExpectations fn (fn -> Expectations params ()) params where
  buildExpectations fn f = runExpectations (f fn)

-- | Type class for dispatching expectations based on mock function type
class ExpectsDispatch fn exp m where
  expects :: m fn -> exp -> m fn

-- | Specialized instance for MockResult helpers (handles type inference)
instance {-# OVERLAPPING #-} (exp ~ Expectations params (), ExpectsDispatchImpl 'True (MockResult params) (Expectations params ()) m) => ExpectsDispatch (MockResult params) exp m where
  expects = expectsDispatchImpl @'True

-- | Specialized instance for Unit helpers (fallback)
instance {-# OVERLAPPING #-} (exp ~ Expectations params (), ExpectsDispatchImpl 'True () (Expectations params ()) m) => ExpectsDispatch () exp m where
  expects = expectsDispatchImpl @'True

-- | Generic instance for normal mocks
instance {-# OVERLAPPABLE #-} (ExpectsDispatchImpl 'False fn exp m) => ExpectsDispatch fn exp m where
  expects = expectsDispatchImpl @'False

-- | Internal class for implementation dispatch
class ExpectsDispatchImpl (flag :: Bool) fn exp m where
  expectsDispatchImpl :: m fn -> exp -> m fn

-- | Instance for normal mocks (flag ~ 'False)
--   Strict matching of params
instance
  ( MonadIO m
  , ResolvableMock fn
  , ResolvableParamsOf fn ~ params
  , ExtractParams exp
  , ExpParams exp ~ params
  , BuildExpectations fn exp params
  , Show params
  , EqParams params
  ) =>
  ExpectsDispatchImpl 'False fn exp m
  where
  expectsDispatchImpl mockFnM exp = do
    WithMockContext ctxVar <- liftIO askWithMockContext
    -- Try to help type inference by using exp first
    let _ = extractParams exp :: Proxy params
    mockFn <- mockFnM
    -- Get the recorder from the thread-local store (set by mock/register)
    -- This avoids StableName lookup and is HPC-safe
    (mockName, mRecorder) <- liftIO $ MockRegistry.getLastRecorder @(InvocationRecorder params)

    let resolved = case mRecorder of
          Just recorder -> ResolvedMock mockName recorder
          Nothing -> errorWithoutStackTrace "expects: mock recorder not found. Use mock inside withMock/runMockT."

    let expectations = buildExpectations mockFn exp
    let actions = map (verifyExpectationDirect resolved) expectations
    liftIO $ atomically $ modifyTVar' ctxVar (++ actions)
    pure mockFn

-- | Instance for MockResult mocks (flag ~ 'True)
--   Dynamic resolution using expectation params
instance
  ( MonadIO m
  , BuildExpectations (MockResult params) (Expectations params ()) params
  , Show params
  , EqParams params
  ) =>
  ExpectsDispatchImpl 'True (MockResult params) (Expectations params ()) m
  where
  expectsDispatchImpl mockFnM exp = do
    WithMockContext ctxVar <- liftIO askWithMockContext
    _ <- mockFnM
    (mockName, mRecorder) <- liftIO MockRegistry.getLastRecorderRaw
    resolved <- case mRecorder of
      Just raw -> do
         let recorder = unsafeCoerce raw :: InvocationRecorder params
         pure $ ResolvedMock mockName recorder
      Nothing -> errorWithoutStackTrace "expects: mock recorder not found (Dynamic Resolution Failed). Ensure the mock helper function was called."
    -- Use the expectations directly since we know the context
    let expectations = runExpectations exp
    let actions = map (verifyExpectationDirect resolved) expectations
    liftIO $ atomically $ modifyTVar' ctxVar (++ actions)
    pure (MockResult ())

-- | Instance for Unit mocks (flag ~ 'True)
--   Dynamic resolution using expectation params
instance
  ( MonadIO m
  , BuildExpectations () (Expectations params ()) params
  , Show params
  , EqParams params
  ) =>
  ExpectsDispatchImpl 'True () (Expectations params ()) m
  where
  expectsDispatchImpl mockFnM exp = do
    WithMockContext ctxVar <- liftIO askWithMockContext
    _ <- mockFnM
    (mockName, mRecorder) <- liftIO MockRegistry.getLastRecorderRaw
    resolved <- case mRecorder of
      Just raw -> do
         let recorder = unsafeCoerce raw :: InvocationRecorder params
         pure $ ResolvedMock mockName recorder
      Nothing -> errorWithoutStackTrace "expects: mock recorder not found (Dynamic Resolution Failed). Ensure the mock helper function was called."
    let expectations = buildExpectations () exp
    let actions = map (verifyExpectationDirect resolved) expectations
    liftIO $ atomically $ modifyTVar' ctxVar (++ actions)
    pure ()



-- | Create a count expectation builder
--   The params type is inferred from the mock function in expects
--   Use type application to specify params when needed: called @(Param String) once
-- | Class-based called builder so that the `params` type can be resolved
--   via instance selection in the `expects` context.
--   This version uses a type class to help type inference by allowing
--   the params type to be inferred from the context where it's used.
class Called params where
  called :: TimesSpec -> Expectations params ()

-- | Default instance that works for any params type
instance {-# OVERLAPPABLE #-} Called params where
  called (TimesSpec method) = do
    addExpectation (CountAnyExpectation method)

-- | Combine expectations with arguments
--   Accepts both raw values (like "a") and Param values (like param "a")
class WithArgs exp args params | exp args -> params where
  with :: exp -> args -> Expectations params ()

instance {-# OVERLAPPING #-}
  WithArgs (Expectations params ()) params params
  where
  with expM args = do
    expM
    -- Extract the last expectation (last in list, since addExpectation appends) and modify it to include args
    Expectations $ do
      exps <- get
      case reverse exps of
        [] -> error "with: no expectation to add arguments to"
        (CountAnyExpectation method : rest) -> do
          put (reverse rest)
          modify (++ [CountExpectation method args])
        _ -> error "with: can only add arguments to count-only expectations"

instance {-# OVERLAPPABLE #-}
  (params ~ Param a, Show a, Eq a) =>
  WithArgs (Expectations params ()) a params
  where
  with expM rawValue = do
    expM
    -- Extract the last expectation (last in list, since addExpectation appends) and modify it to include args
    Expectations $ do
      exps <- get
      case reverse exps of
        [] -> error "with: no expectation to add arguments to"
        (CountAnyExpectation method : rest) -> do
          put (reverse rest)
          modify (++ [CountExpectation method (param rawValue)])
        _ -> error "with: can only add arguments to count-only expectations"




-- | Create an order expectation
--   Accepts both Param values and raw values
class CalledInOrder args params | args -> params where
  calledInOrder :: args -> Expectations params ()

-- | Convenience instance: infer params from function argument type @a@
instance
  (params ~ Param a, Show a, Eq a) =>
  CalledInOrder [a] params
  where
  calledInOrder args =
    addExpectation (OrderExpectation ExactlySequence (map param args))

-- | Create a partial order expectation
--   Accepts both Param values and raw values
class CalledInSequence args params | args -> params where
  calledInSequence :: args -> Expectations params ()

-- | Convenience instance: infer params from function argument type @a@
instance
  (params ~ Param a, Show a, Eq a) =>
  CalledInSequence [a] params
  where
  calledInSequence args =
    addExpectation (OrderExpectation PartiallySequence (map param args))