{-# LANGUAGE OverloadedStrings #-}

{- |
Module      : Data.Ollama.Ps
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental
Description : Functionality for listing running models in the Ollama client.

This module provides functions to retrieve a list of models currently running on the Ollama server.
It includes both an IO-based function ('ps') and a monadic version ('psM') for use in 'MonadIO'
contexts. The operation is performed via a GET request to the "/api//ps" endpoint, returning a
'RunningModels' type containing a list of 'RunningModel' records with details about each running model.

Example:

>>> ps Nothing
Right (RunningModels [RunningModel ...])
-}
module Data.Ollama.Ps
  ( -- * List Running Models API
    ps
  , psM

    -- * Model Types
  , RunningModels (..)
  , RunningModel (..)
  ) where

import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.Aeson
import Data.Ollama.Common.Config (OllamaConfig)
import Data.Ollama.Common.Error (OllamaError)
import Data.Ollama.Common.Types as CT
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import Data.Time
import GHC.Int (Int64)

-- | A wrapper type containing a list of running models.
newtype RunningModels
  = -- | List of 'RunningModel' records describing currently running models.
    RunningModels [RunningModel]
  deriving (RunningModels -> RunningModels -> Bool
(RunningModels -> RunningModels -> Bool)
-> (RunningModels -> RunningModels -> Bool) -> Eq RunningModels
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RunningModels -> RunningModels -> Bool
== :: RunningModels -> RunningModels -> Bool
$c/= :: RunningModels -> RunningModels -> Bool
/= :: RunningModels -> RunningModels -> Bool
Eq, Int -> RunningModels -> ShowS
[RunningModels] -> ShowS
RunningModels -> String
(Int -> RunningModels -> ShowS)
-> (RunningModels -> String)
-> ([RunningModels] -> ShowS)
-> Show RunningModels
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RunningModels -> ShowS
showsPrec :: Int -> RunningModels -> ShowS
$cshow :: RunningModels -> String
show :: RunningModels -> String
$cshowList :: [RunningModels] -> ShowS
showList :: [RunningModels] -> ShowS
Show)

-- | Details about a specific running model.
data RunningModel = RunningModel
  { RunningModel -> Text
name_ :: !Text
  -- ^ The name of the running model instance.
  , RunningModel -> Text
modelName :: !Text
  -- ^ The base model name (e.g., "gemma3").
  , RunningModel -> Int64
size_ :: !Int64
  -- ^ The size of the model in bytes.
  , RunningModel -> Text
modelDigest :: !Text
  -- ^ The digest (hash) of the model.
  , RunningModel -> ModelDetails
modelDetails :: !ModelDetails
  -- ^ Additional details about the model (e.g., format, family, parameters).
  , RunningModel -> UTCTime
expiresAt :: !UTCTime
  -- ^ The timestamp when the model's memory allocation expires.
  , RunningModel -> Int64
sizeVRam :: !Int64
  -- ^ The size of the model's VRAM usage in bytes.
  }
  deriving (RunningModel -> RunningModel -> Bool
(RunningModel -> RunningModel -> Bool)
-> (RunningModel -> RunningModel -> Bool) -> Eq RunningModel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RunningModel -> RunningModel -> Bool
== :: RunningModel -> RunningModel -> Bool
$c/= :: RunningModel -> RunningModel -> Bool
/= :: RunningModel -> RunningModel -> Bool
Eq, Int -> RunningModel -> ShowS
[RunningModel] -> ShowS
RunningModel -> String
(Int -> RunningModel -> ShowS)
-> (RunningModel -> String)
-> ([RunningModel] -> ShowS)
-> Show RunningModel
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RunningModel -> ShowS
showsPrec :: Int -> RunningModel -> ShowS
$cshow :: RunningModel -> String
show :: RunningModel -> String
$cshowList :: [RunningModel] -> ShowS
showList :: [RunningModel] -> ShowS
Show)

-- | JSON parsing instance for 'RunningModels'.
instance FromJSON RunningModels where
  parseJSON :: Value -> Parser RunningModels
parseJSON = String
-> (Object -> Parser RunningModels)
-> Value
-> Parser RunningModels
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"Models" ((Object -> Parser RunningModels) -> Value -> Parser RunningModels)
-> (Object -> Parser RunningModels)
-> Value
-> Parser RunningModels
forall a b. (a -> b) -> a -> b
$ \Object
v -> [RunningModel] -> RunningModels
RunningModels ([RunningModel] -> RunningModels)
-> Parser [RunningModel] -> Parser RunningModels
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
v Object -> Key -> Parser [RunningModel]
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"models"

-- | JSON parsing instance for 'RunningModel'.
instance FromJSON RunningModel where
  parseJSON :: Value -> Parser RunningModel
parseJSON = String
-> (Object -> Parser RunningModel) -> Value -> Parser RunningModel
forall a. String -> (Object -> Parser a) -> Value -> Parser a
withObject String
"RunningModel" ((Object -> Parser RunningModel) -> Value -> Parser RunningModel)
-> (Object -> Parser RunningModel) -> Value -> Parser RunningModel
forall a b. (a -> b) -> a -> b
$ \Object
v ->
    Text
-> Text
-> Int64
-> Text
-> ModelDetails
-> UTCTime
-> Int64
-> RunningModel
RunningModel
      (Text
 -> Text
 -> Int64
 -> Text
 -> ModelDetails
 -> UTCTime
 -> Int64
 -> RunningModel)
-> Parser Text
-> Parser
     (Text
      -> Int64
      -> Text
      -> ModelDetails
      -> UTCTime
      -> Int64
      -> RunningModel)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Object
v Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"name"
      Parser
  (Text
   -> Int64
   -> Text
   -> ModelDetails
   -> UTCTime
   -> Int64
   -> RunningModel)
-> Parser Text
-> Parser
     (Int64 -> Text -> ModelDetails -> UTCTime -> Int64 -> RunningModel)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"model"
      Parser
  (Int64 -> Text -> ModelDetails -> UTCTime -> Int64 -> RunningModel)
-> Parser Int64
-> Parser
     (Text -> ModelDetails -> UTCTime -> Int64 -> RunningModel)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser Int64
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"size"
      Parser (Text -> ModelDetails -> UTCTime -> Int64 -> RunningModel)
-> Parser Text
-> Parser (ModelDetails -> UTCTime -> Int64 -> RunningModel)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser Text
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"digest"
      Parser (ModelDetails -> UTCTime -> Int64 -> RunningModel)
-> Parser ModelDetails -> Parser (UTCTime -> Int64 -> RunningModel)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser ModelDetails
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"details"
      Parser (UTCTime -> Int64 -> RunningModel)
-> Parser UTCTime -> Parser (Int64 -> RunningModel)
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser UTCTime
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"expires_at"
      Parser (Int64 -> RunningModel)
-> Parser Int64 -> Parser RunningModel
forall a b. Parser (a -> b) -> Parser a -> Parser b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Object
v Object -> Key -> Parser Int64
forall a. FromJSON a => Object -> Key -> Parser a
.: Key
"size_vram"

{- | Retrieves a list of currently running models from the Ollama server.

Sends a GET request to the "/api//ps" endpoint to fetch the list of running models.
Returns 'Right' with a 'RunningModels' containing the list of 'RunningModel' on success,
or 'Left' with an 'OllamaError' on failure.
Example:

>>> ps Nothing
Right (RunningModels [RunningModel {name_ = "gemma3:instance1", modelName = "gemma3", ...}])
-}
ps ::
  -- | Optional 'OllamaConfig' (defaults to 'defaultOllamaConfig' if 'Nothing')
  Maybe OllamaConfig ->
  IO (Either OllamaError RunningModels)
ps :: Maybe OllamaConfig -> IO (Either OllamaError RunningModels)
ps Maybe OllamaConfig
mbConfig = do
  Text
-> ByteString
-> Maybe Value
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError RunningModels))
-> IO (Either OllamaError RunningModels)
forall payload response.
ToJSON payload =>
Text
-> ByteString
-> Maybe payload
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError response))
-> IO (Either OllamaError response)
withOllamaRequest
    Text
"/api//ps"
    ByteString
"GET"
    (Maybe Value
forall a. Maybe a
Nothing :: Maybe Value)
    Maybe OllamaConfig
mbConfig
    Response BodyReader -> IO (Either OllamaError RunningModels)
forall a.
FromJSON a =>
Response BodyReader -> IO (Either OllamaError a)
commonNonStreamingHandler

{- | MonadIO version of 'ps' for use in monadic contexts.

Lifts the 'ps' function into a 'MonadIO' context, allowing it to be used in monadic computations.
-}
psM :: MonadIO m => Maybe OllamaConfig -> m (Either OllamaError RunningModels)
psM :: forall (m :: * -> *).
MonadIO m =>
Maybe OllamaConfig -> m (Either OllamaError RunningModels)
psM Maybe OllamaConfig
mbCfg = IO (Either OllamaError RunningModels)
-> m (Either OllamaError RunningModels)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either OllamaError RunningModels)
 -> m (Either OllamaError RunningModels))
-> IO (Either OllamaError RunningModels)
-> m (Either OllamaError RunningModels)
forall a b. (a -> b) -> a -> b
$ Maybe OllamaConfig -> IO (Either OllamaError RunningModels)
ps Maybe OllamaConfig
mbCfg