{- |
Module      : Llama.Performance
Description : High level Performance interface for llama-cpp
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
-}
module Llama.Performance (
    printContextPerformance
, resetContextPerformance
, printSamplerPerformance
, resetSamplerPerformance
, getContextPerformance
, getSamplerPerformance
  ) where

import Llama.Internal.Types
import Foreign
import Llama.Internal.Foreign

-- | Print performance information for a context
printContextPerformance :: Context -> IO ()
printContextPerformance :: Context -> IO ()
printContextPerformance (Context ForeignPtr CLlamaContext
ctxFPtr) =
  ForeignPtr CLlamaContext -> (Ptr CLlamaContext -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr CLlamaContext
ctxFPtr ((Ptr CLlamaContext -> IO ()) -> IO ())
-> (Ptr CLlamaContext -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr CLlamaContext
ctxPtr ->
    CLlamaContext -> IO ()
c_llama_perf_context_print (Ptr CLlamaContext -> CLlamaContext
CLlamaContext Ptr CLlamaContext
ctxPtr)

-- | Reset performance information for a context
resetContextPerformance :: Context -> IO ()
resetContextPerformance :: Context -> IO ()
resetContextPerformance (Context ForeignPtr CLlamaContext
ctxFPtr) =
  ForeignPtr CLlamaContext -> (Ptr CLlamaContext -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr CLlamaContext
ctxFPtr ((Ptr CLlamaContext -> IO ()) -> IO ())
-> (Ptr CLlamaContext -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr CLlamaContext
ctxPtr ->
    CLlamaContext -> IO ()
c_llama_perf_context_reset (Ptr CLlamaContext -> CLlamaContext
CLlamaContext Ptr CLlamaContext
ctxPtr)

-- | Print performance information for a sampler chain
printSamplerPerformance :: Sampler -> IO ()
printSamplerPerformance :: Sampler -> IO ()
printSamplerPerformance (Sampler ForeignPtr LlamaSampler
samplerFPtr) =
  ForeignPtr LlamaSampler -> (Ptr LlamaSampler -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr LlamaSampler
samplerFPtr ((Ptr LlamaSampler -> IO ()) -> IO ())
-> (Ptr LlamaSampler -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr LlamaSampler
samplerPtr ->
    Ptr LlamaSampler -> IO ()
c_llama_perf_sampler_print Ptr LlamaSampler
samplerPtr

-- | Reset performance information for a sampler chain
resetSamplerPerformance :: Sampler -> IO ()
resetSamplerPerformance :: Sampler -> IO ()
resetSamplerPerformance (Sampler ForeignPtr LlamaSampler
samplerFPtr) =
  ForeignPtr LlamaSampler -> (Ptr LlamaSampler -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr LlamaSampler
samplerFPtr ((Ptr LlamaSampler -> IO ()) -> IO ())
-> (Ptr LlamaSampler -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr LlamaSampler
samplerPtr ->
    Ptr LlamaSampler -> IO ()
c_llama_perf_sampler_reset Ptr LlamaSampler
samplerPtr

-- | Get performance data for a context
getContextPerformance :: Context -> IO LlamaPerfContextData
getContextPerformance :: Context -> IO LlamaPerfContextData
getContextPerformance (Context ForeignPtr CLlamaContext
ctxFPtr) = do
  (Ptr LlamaPerfContextData -> IO LlamaPerfContextData)
-> IO LlamaPerfContextData
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr LlamaPerfContextData -> IO LlamaPerfContextData)
 -> IO LlamaPerfContextData)
-> (Ptr LlamaPerfContextData -> IO LlamaPerfContextData)
-> IO LlamaPerfContextData
forall a b. (a -> b) -> a -> b
$ \Ptr LlamaPerfContextData
perfDataPtr -> do
    ForeignPtr CLlamaContext -> (Ptr CLlamaContext -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr CLlamaContext
ctxFPtr ((Ptr CLlamaContext -> IO ()) -> IO ())
-> (Ptr CLlamaContext -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr CLlamaContext
ctxPtr -> do
      CLlamaContext -> Ptr LlamaPerfContextData -> IO ()
c_llama_perf_context (Ptr CLlamaContext -> CLlamaContext
CLlamaContext Ptr CLlamaContext
ctxPtr) Ptr LlamaPerfContextData
perfDataPtr
    Ptr LlamaPerfContextData -> IO LlamaPerfContextData
forall a. Storable a => Ptr a -> IO a
peek Ptr LlamaPerfContextData
perfDataPtr

-- | Get performance data for a sampler chain
getSamplerPerformance :: Sampler -> IO LlamaPerfSamplerData
getSamplerPerformance :: Sampler -> IO LlamaPerfSamplerData
getSamplerPerformance (Sampler ForeignPtr LlamaSampler
samplerFPtr) = do
  (Ptr LlamaPerfSamplerData -> IO LlamaPerfSamplerData)
-> IO LlamaPerfSamplerData
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr LlamaPerfSamplerData -> IO LlamaPerfSamplerData)
 -> IO LlamaPerfSamplerData)
-> (Ptr LlamaPerfSamplerData -> IO LlamaPerfSamplerData)
-> IO LlamaPerfSamplerData
forall a b. (a -> b) -> a -> b
$ \Ptr LlamaPerfSamplerData
perfDataPtr -> do
    ForeignPtr LlamaSampler -> (Ptr LlamaSampler -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr LlamaSampler
samplerFPtr ((Ptr LlamaSampler -> IO ()) -> IO ())
-> (Ptr LlamaSampler -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr LlamaSampler
samplerPtr -> do
      Ptr LlamaSampler -> Ptr LlamaPerfSamplerData -> IO ()
c_llama_perf_sampler Ptr LlamaSampler
samplerPtr Ptr LlamaPerfSamplerData
perfDataPtr
    Ptr LlamaPerfSamplerData -> IO LlamaPerfSamplerData
forall a. Storable a => Ptr a -> IO a
peek Ptr LlamaPerfSamplerData
perfDataPtr