{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

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

This module provides functions to pull (download) models from the Ollama server. It includes both
high-level ('pull', 'pullM') and low-level ('pullOps', 'pullOpsM') APIs for pulling models, with
support for streaming progress updates and insecure connections. The 'PullOps' type configures the
pull request, and 'PullResp' represents the response containing the status and progress details.

The pull operation is performed via a POST request to the "/api//pull" endpoint. Streaming mode,
when enabled, provides real-time progress updates by printing the remaining bytes to the console.

Example:

>>> pull "gemma3"
Remaining bytes: 123456789
...
Completed
Right (PullResp {status = "success", ...})
-}
module Data.Ollama.Pull
  ( -- * Pull Model API
    pull
  , pullOps
  , pullM
  , pullOpsM
  ) where

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

-- | Configuration options for pulling a model.
data PullOps = PullOps
  { PullOps -> Text
name :: !Text
  -- ^ The name of the model to pull (e.g., "gemma3").
  , PullOps -> Maybe Bool
insecure :: !(Maybe Bool)
  -- ^ Optional flag to allow insecure connections. If 'Just True', insecure connections are permitted.
  , PullOps -> Maybe Bool
stream :: !(Maybe Bool)
  -- ^ Optional flag to enable streaming of the download. If 'Just True', progress updates are streamed.
  }
  deriving (Int -> PullOps -> ShowS
[PullOps] -> ShowS
PullOps -> String
(Int -> PullOps -> ShowS)
-> (PullOps -> String) -> ([PullOps] -> ShowS) -> Show PullOps
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PullOps -> ShowS
showsPrec :: Int -> PullOps -> ShowS
$cshow :: PullOps -> String
show :: PullOps -> String
$cshowList :: [PullOps] -> ShowS
showList :: [PullOps] -> ShowS
Show, PullOps -> PullOps -> Bool
(PullOps -> PullOps -> Bool)
-> (PullOps -> PullOps -> Bool) -> Eq PullOps
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PullOps -> PullOps -> Bool
== :: PullOps -> PullOps -> Bool
$c/= :: PullOps -> PullOps -> Bool
/= :: PullOps -> PullOps -> Bool
Eq, (forall x. PullOps -> Rep PullOps x)
-> (forall x. Rep PullOps x -> PullOps) -> Generic PullOps
forall x. Rep PullOps x -> PullOps
forall x. PullOps -> Rep PullOps x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. PullOps -> Rep PullOps x
from :: forall x. PullOps -> Rep PullOps x
$cto :: forall x. Rep PullOps x -> PullOps
to :: forall x. Rep PullOps x -> PullOps
Generic, [PullOps] -> Value
[PullOps] -> Encoding
PullOps -> Bool
PullOps -> Value
PullOps -> Encoding
(PullOps -> Value)
-> (PullOps -> Encoding)
-> ([PullOps] -> Value)
-> ([PullOps] -> Encoding)
-> (PullOps -> Bool)
-> ToJSON PullOps
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
$ctoJSON :: PullOps -> Value
toJSON :: PullOps -> Value
$ctoEncoding :: PullOps -> Encoding
toEncoding :: PullOps -> Encoding
$ctoJSONList :: [PullOps] -> Value
toJSONList :: [PullOps] -> Value
$ctoEncodingList :: [PullOps] -> Encoding
toEncodingList :: [PullOps] -> Encoding
$comitField :: PullOps -> Bool
omitField :: PullOps -> Bool
ToJSON)

-- | Response data from a pull operation.
data PullResp = PullResp
  { PullResp -> Text
status :: !Text
  -- ^ The status of the pull operation (e.g., "success" or "failure").
  , PullResp -> Maybe Text
digest :: !(Maybe Text)
  -- ^ The digest (hash) of the model, if available.
  , PullResp -> Maybe Int64
total :: !(Maybe Int64)
  -- ^ The total size of the model in bytes, if available.
  , PullResp -> Maybe Int64
completed :: !(Maybe Int64)
  -- ^ The number of bytes downloaded, if available.
  }
  deriving (Int -> PullResp -> ShowS
[PullResp] -> ShowS
PullResp -> String
(Int -> PullResp -> ShowS)
-> (PullResp -> String) -> ([PullResp] -> ShowS) -> Show PullResp
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PullResp -> ShowS
showsPrec :: Int -> PullResp -> ShowS
$cshow :: PullResp -> String
show :: PullResp -> String
$cshowList :: [PullResp] -> ShowS
showList :: [PullResp] -> ShowS
Show, PullResp -> PullResp -> Bool
(PullResp -> PullResp -> Bool)
-> (PullResp -> PullResp -> Bool) -> Eq PullResp
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PullResp -> PullResp -> Bool
== :: PullResp -> PullResp -> Bool
$c/= :: PullResp -> PullResp -> Bool
/= :: PullResp -> PullResp -> Bool
Eq, (forall x. PullResp -> Rep PullResp x)
-> (forall x. Rep PullResp x -> PullResp) -> Generic PullResp
forall x. Rep PullResp x -> PullResp
forall x. PullResp -> Rep PullResp x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. PullResp -> Rep PullResp x
from :: forall x. PullResp -> Rep PullResp x
$cto :: forall x. Rep PullResp x -> PullResp
to :: forall x. Rep PullResp x -> PullResp
Generic, Maybe PullResp
Value -> Parser [PullResp]
Value -> Parser PullResp
(Value -> Parser PullResp)
-> (Value -> Parser [PullResp])
-> Maybe PullResp
-> FromJSON PullResp
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
$cparseJSON :: Value -> Parser PullResp
parseJSON :: Value -> Parser PullResp
$cparseJSONList :: Value -> Parser [PullResp]
parseJSONList :: Value -> Parser [PullResp]
$comittedField :: Maybe PullResp
omittedField :: Maybe PullResp
FromJSON)

instance HasDone PullResp where
  getDone :: PullResp -> Bool
getDone PullResp {Maybe Int64
Maybe Text
Text
status :: PullResp -> Text
digest :: PullResp -> Maybe Text
total :: PullResp -> Maybe Int64
completed :: PullResp -> Maybe Int64
status :: Text
digest :: Maybe Text
total :: Maybe Int64
completed :: Maybe Int64
..} = Text
status Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
"success"

{- | Pulls a model with full configuration.

Sends a POST request to the "/api//pull" endpoint to download the specified model. Supports
streaming progress updates (if 'stream' is 'Just True') and insecure connections (if 'insecure'
is 'Just True'). Prints remaining bytes during streaming and "Completed" when finished.
Returns 'Right' with a 'PullResp' on success or 'Left' with an 'OllamaError' on failure.
-}
pullOps ::
  -- | Model name
  Text ->
  -- | Optional insecure connection flag
  Maybe Bool ->
  -- | Optional streaming flag
  Maybe Bool ->
  -- | Optional 'OllamaConfig' (defaults to 'defaultOllamaConfig' if 'Nothing')
  Maybe OllamaConfig ->
  IO (Either OllamaError PullResp)
pullOps :: Text
-> Maybe Bool
-> Maybe Bool
-> Maybe OllamaConfig
-> IO (Either OllamaError PullResp)
pullOps Text
modelName Maybe Bool
mInsecure Maybe Bool
mStream Maybe OllamaConfig
mbConfig = do
  Text
-> ByteString
-> Maybe PullOps
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError PullResp))
-> IO (Either OllamaError PullResp)
forall payload response.
ToJSON payload =>
Text
-> ByteString
-> Maybe payload
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError response))
-> IO (Either OllamaError response)
withOllamaRequest
    Text
"/api//pull"
    ByteString
"POST"
    (PullOps -> Maybe PullOps
forall a. a -> Maybe a
Just (PullOps -> Maybe PullOps) -> PullOps -> Maybe PullOps
forall a b. (a -> b) -> a -> b
$ PullOps {name :: Text
name = Text
modelName, insecure :: Maybe Bool
insecure = Maybe Bool
mInsecure, stream :: Maybe Bool
stream = Maybe Bool
mStream})
    Maybe OllamaConfig
mbConfig
    ((PullResp -> IO ())
-> IO () -> Response BodyReader -> IO (Either OllamaError PullResp)
forall a.
(HasDone a, FromJSON a) =>
(a -> IO ())
-> IO () -> Response BodyReader -> IO (Either OllamaError a)
commonStreamHandler PullResp -> IO ()
onToken IO ()
onComplete)
  where
    onToken :: PullResp -> IO ()
    onToken :: PullResp -> IO ()
onToken PullResp
res = do
      let completed' :: Int64
completed' = Int64 -> Maybe Int64 -> Int64
forall a. a -> Maybe a -> a
fromMaybe Int64
0 (PullResp -> Maybe Int64
completed PullResp
res)
      let total' :: Int64
total' = Int64 -> Maybe Int64 -> Int64
forall a. a -> Maybe a -> a
fromMaybe Int64
0 (PullResp -> Maybe Int64
total PullResp
res)
      String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Remaining bytes: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int64 -> String
forall a. Show a => a -> String
show (Int64
total' Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
completed')

    onComplete :: IO ()
    onComplete :: IO ()
onComplete = String -> IO ()
putStrLn String
"Completed"

{- | Simplified API for pulling a model.

A higher-level function that pulls a model using default settings for insecure connections,
streaming, and Ollama configuration. Suitable for basic use cases.
-}
pull ::
  -- | Model name
  Text ->
  IO (Either OllamaError PullResp)
pull :: Text -> IO (Either OllamaError PullResp)
pull Text
modelName = Text
-> Maybe Bool
-> Maybe Bool
-> Maybe OllamaConfig
-> IO (Either OllamaError PullResp)
pullOps Text
modelName Maybe Bool
forall a. Maybe a
Nothing Maybe Bool
forall a. Maybe a
Nothing Maybe OllamaConfig
forall a. Maybe a
Nothing

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

Lifts the 'pull' function into a 'MonadIO' context, allowing it to be used in monadic computations.

Example:

>>> import Control.Monad.IO.Class
>>> runReaderT (pullM "gemma3") someContext
Right (PullResp {status = "success", ...})
-}
pullM :: MonadIO m => Text -> m (Either OllamaError PullResp)
pullM :: forall (m :: * -> *).
MonadIO m =>
Text -> m (Either OllamaError PullResp)
pullM Text
t = IO (Either OllamaError PullResp) -> m (Either OllamaError PullResp)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either OllamaError PullResp)
 -> m (Either OllamaError PullResp))
-> IO (Either OllamaError PullResp)
-> m (Either OllamaError PullResp)
forall a b. (a -> b) -> a -> b
$ Text -> IO (Either OllamaError PullResp)
pull Text
t

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

Lifts the 'pullOps' function into a 'MonadIO' context, allowing it to be used in monadic computations
with full configuration options.

Example:

>>> import Control.Monad.IO.Class
>>> runReaderT (pullOpsM "gemma3" Nothing (Just True) Nothing) someContext
Remaining bytes: 123456789
...
Completed
Right (PullResp {status = "success", ...})
-}
pullOpsM ::
  MonadIO m =>
  Text ->
  Maybe Bool ->
  Maybe Bool ->
  Maybe OllamaConfig ->
  m (Either OllamaError PullResp)
pullOpsM :: forall (m :: * -> *).
MonadIO m =>
Text
-> Maybe Bool
-> Maybe Bool
-> Maybe OllamaConfig
-> m (Either OllamaError PullResp)
pullOpsM Text
t Maybe Bool
mbInsecure Maybe Bool
mbStream Maybe OllamaConfig
mbCfg = IO (Either OllamaError PullResp) -> m (Either OllamaError PullResp)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either OllamaError PullResp)
 -> m (Either OllamaError PullResp))
-> IO (Either OllamaError PullResp)
-> m (Either OllamaError PullResp)
forall a b. (a -> b) -> a -> b
$ Text
-> Maybe Bool
-> Maybe Bool
-> Maybe OllamaConfig
-> IO (Either OllamaError PullResp)
pullOps Text
t Maybe Bool
mbInsecure Maybe Bool
mbStream Maybe OllamaConfig
mbCfg