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
  , HasMockContext(..)
  , addRecorderLocal
  ,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, ask)
import Control.Exception (bracket_)
import Control.Concurrent.STM (newTVarIO, readTVarIO, atomically, modifyTVar')
import Control.Monad.State (get, put, modify)
import Data.Dynamic (Dynamic)
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(..)
  , MockName
  , 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
--   This implementation does NOT use threadWithMockStore (no unsafePerformIO)
--   because the context is passed via ReaderT environment.
withMock :: ReaderT WithMockContext IO a -> IO a
withMock action = do
  ctxVar <- newTVarIO []
  recordersVar <- newTVarIO []
  let ctx = WithMockContext ctxVar recordersVar
  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 []
  recordersVar <- newTVarIO []
  let ctx = WithMockContext ctxVar recordersVar
  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.
-- | Capability to access mock context
--   This type class separates safe (ReaderT-based) context access
--   from unsafe (global registry-based) context access.
--
--   Instances:
--   - 'ReaderT WithMockContext m': Safe, uses 'ask' (no unsafePerformIO)
--   - 'IO': Unsafe, uses global registry (unsafePerformIO dependency)
class Monad m => HasMockContext m where
  getMockContext :: m WithMockContext

-- | Safe instance: ReaderT-based context access
--   Uses 'ask' to obtain context from the ReaderT environment.
--   This does NOT use unsafePerformIO.
instance MonadIO m => HasMockContext (ReaderT WithMockContext m) where
  getMockContext = ask

-- | Unsafe instance: IO-based context access
--   Uses global registry (threadWithMockStore) which is initialized
--   with unsafePerformIO. This is kept for backward compatibility
--   with 'withMockIO'.
instance HasMockContext IO where
  getMockContext = askWithMockContext

-- | Get WithMockContext from thread-local storage
--   This is used by the IO instance of HasMockContext.
--   It relies on the global registry (threadWithMockStore) which
--   is initialized with unsafePerformIO.
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."

-- | Add a recorder to the local context (ReaderT environment)
--   This does NOT use unsafePerformIO or global registry.
--   Used by RegisterMock instance for safe registration.
addRecorderLocal :: (HasMockContext m, MonadIO m)
  => Maybe MockName -> Dynamic -> m ()
addRecorderLocal name dyn = do
  WithMockContext {contextRecorders = ref} <- getMockContext
  liftIO $ atomically $ modifyTVar' ref ((name, dyn):)

-- | 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)
--   Uses HasMockContext to obtain context, which automatically selects:
--   - ReaderT: ask (safe, no unsafePerformIO)
--   - IO: askWithMockContext (unsafe, uses global registry)
instance
  ( HasMockContext m
  , 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 {contextVerifications = ctxVar} <- getMockContext  -- Updated pattern
    -- 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
--   Uses HasMockContext for capability-based context access
instance
  ( HasMockContext m
  , 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 {contextVerifications = ctxVar} <- getMockContext  -- Updated pattern
    _ <- 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
--   Uses HasMockContext for capability-based context access
instance
  ( HasMockContext m
  , MonadIO m
  , BuildExpectations () (Expectations params ()) params
  , Show params
  , EqParams params
  ) =>
  ExpectsDispatchImpl 'True () (Expectations params ()) m
  where
  expectsDispatchImpl mockFnM exp = do
    WithMockContext {contextVerifications = ctxVar} <- getMockContext  -- Updated pattern
    _ <- 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))