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

module Data.Ollama.Pull
  ( -- * Downloaded Models
    pull
  , pullOps
  ) where

import Data.Aeson
import Data.ByteString.Lazy.Char8 qualified as BSL
import Data.Maybe (fromMaybe)
import Data.Ollama.Common.Utils as CU
import Data.Text (Text)
import Data.Text qualified as T
import GHC.Generics
import GHC.Int (Int64)
import Network.HTTP.Client

-- TODO: Add Options parameter

-- | Configuration options for pulling a model.
data PullOps = PullOps
  { PullOps -> Text
name :: Text
  -- ^ The name of the model to pull.
  , PullOps -> Maybe Bool
insecure :: Maybe Bool
  -- ^ Option to allow insecure connections.
  -- If set to 'Just True', the pull operation will allow insecure connections.
  , PullOps -> Maybe Bool
stream :: Maybe Bool
  -- ^ Option to enable streaming of the download.
  -- If set to 'Just True', the download will be 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 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 completed, 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)

{- |
Pull a model with additional options for insecure connections and streaming.
This function interacts directly with the Ollama API to download the specified model.

Example:

> pullOps "myModel" (Just True) (Just True)

This will attempt to pull "myModel" with insecure connections allowed and enable streaming.
-}
pullOps ::
  -- | Model Name
  Text ->
  -- | Insecure
  Maybe Bool ->
  -- | Stream
  Maybe Bool ->
  IO ()
pullOps :: Text -> Maybe Bool -> Maybe Bool -> IO ()
pullOps Text
modelName Maybe Bool
mInsecure Maybe Bool
mStream = do
  let url :: Text
url = OllamaClient -> Text
CU.host OllamaClient
defaultOllama
  Manager
manager <- ManagerSettings -> IO Manager
newManager ManagerSettings
defaultManagerSettings
  Request
initialRequest <- String -> IO Request
forall (m :: * -> *). MonadThrow m => String -> m Request
parseRequest (String -> IO Request) -> String -> IO Request
forall a b. (a -> b) -> a -> b
$ Text -> String
T.unpack (Text
url Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"/api/pull")
  let reqBody :: PullOps
reqBody =
        PullOps
          { $sel:name:PullOps :: Text
name = Text
modelName
          , $sel:insecure:PullOps :: Maybe Bool
insecure = Maybe Bool
mInsecure
          , $sel:stream:PullOps :: Maybe Bool
stream = Maybe Bool
mStream
          }
      request :: Request
request =
        Request
initialRequest
          { method = "POST"
          , requestBody = RequestBodyLBS $ encode reqBody
          }
  Request -> Manager -> (Response BodyReader -> IO ()) -> IO ()
forall a.
Request -> Manager -> (Response BodyReader -> IO a) -> IO a
withResponse Request
request Manager
manager ((Response BodyReader -> IO ()) -> IO ())
-> (Response BodyReader -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Response BodyReader
response -> do
    let go :: IO ()
go = do
          Method
bs <- BodyReader -> BodyReader
brRead (BodyReader -> BodyReader) -> BodyReader -> BodyReader
forall a b. (a -> b) -> a -> b
$ Response BodyReader -> BodyReader
forall body. Response body -> body
responseBody Response BodyReader
response
          let eRes :: Maybe PullResp
eRes = ByteString -> Maybe PullResp
forall a. FromJSON a => ByteString -> Maybe a
decode (Method -> ByteString
BSL.fromStrict Method
bs) :: Maybe PullResp
          case Maybe PullResp
eRes of
            Maybe PullResp
Nothing -> String -> IO ()
putStrLn String
"Something went wrong"
            Just PullResp
res -> do
              if PullResp -> Text
status PullResp
res Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
"success"
                then 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')
                  IO ()
go
                else do
                  String -> IO ()
putStrLn String
"Completed"
    IO ()
go

{- |
Pull a model using default options. This simplifies the pull operation by
not requiring additional options.

Example:

> pull "myModel"

This will pull "myModel" using default settings (no insecure connections and no streaming).
-}
pull ::
  -- | Model Name
  Text ->
  IO ()
pull :: Text -> IO ()
pull Text
modelName = Text -> Maybe Bool -> Maybe Bool -> IO ()
pullOps Text
modelName Maybe Bool
forall a. Maybe a
Nothing Maybe Bool
forall a. Maybe a
Nothing