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

{- |
Module      : Data.Ollama.Push
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental
Description : Functionality for pushing models to the Ollama server.

This module provides functions to push (upload) a model to the Ollama server. It includes
both an IO-based function ('push') and a monadic version ('pushM') for use in 'MonadIO'
contexts. The push operation is performed via a POST request to the "/api//pull" endpoint,
with support for streaming progress updates and insecure connections.

The 'PushOps' type configures the push request, and 'PushResp' represents the response
containing the status and progress details. Streaming mode, when enabled, provides
real-time progress updates by printing to the console.

Example:

>>> push "gemma3" Nothing (Just True) Nothing
Pushing...
Completed
-}
module Data.Ollama.Push
  ( -- * Push Model API
    push
  , pushM
  ) where

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

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

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

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

{- | Pushes a model to the Ollama server with specified options.

Sends a POST request to the "/api//pull" endpoint to upload the specified model. Supports
streaming progress updates (if 'stream' is 'Just True') and insecure connections (if
'insecure' is 'Just True'). Prints "Pushing..." during streaming and "Completed" when
finished. Returns '()' on completion.
-}
push ::
  -- | Model name
  Text ->
  -- | Optional insecure connection flag
  Maybe Bool ->
  -- | Optional streaming flag
  Maybe Bool ->
  -- | Optional 'OllamaConfig' (defaults to 'defaultOllamaConfig' if 'Nothing')
  Maybe OllamaConfig ->
  IO ()
push :: Text -> Maybe Bool -> Maybe Bool -> Maybe OllamaConfig -> IO ()
push Text
modelName Maybe Bool
mInsecure Maybe Bool
mStream Maybe OllamaConfig
mbConfig = do
  IO (Either OllamaError PushResp) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either OllamaError PushResp) -> IO ())
-> IO (Either OllamaError PushResp) -> IO ()
forall a b. (a -> b) -> a -> b
$
    Text
-> ByteString
-> Maybe PushOps
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError PushResp))
-> IO (Either OllamaError PushResp)
forall payload response.
ToJSON payload =>
Text
-> ByteString
-> Maybe payload
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError response))
-> IO (Either OllamaError response)
withOllamaRequest
      Text
"/api//push"
      ByteString
"POST"
      (PushOps -> Maybe PushOps
forall a. a -> Maybe a
Just (PushOps -> Maybe PushOps) -> PushOps -> Maybe PushOps
forall a b. (a -> b) -> a -> b
$ PushOps {name :: Text
name = Text
modelName, insecure :: Maybe Bool
insecure = Maybe Bool
mInsecure, stream :: Maybe Bool
stream = Maybe Bool
mStream})
      Maybe OllamaConfig
mbConfig
      ((PushResp -> IO ())
-> IO () -> Response BodyReader -> IO (Either OllamaError PushResp)
forall a.
(HasDone a, FromJSON a) =>
(a -> IO ())
-> IO () -> Response BodyReader -> IO (Either OllamaError a)
commonStreamHandler PushResp -> IO ()
onToken IO ()
onComplete)
  where
    onToken :: PushResp -> IO ()
    onToken :: PushResp -> IO ()
onToken PushResp
_ = String -> IO ()
putStrLn String
"Pushing... "

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

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

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

Example:

>>> import Control.Monad.IO.Class
>>> runReaderT (pushM "gemma3" Nothing (Just True) Nothing) someContext
Pushing...
Completed
-}
pushM :: MonadIO m => Text -> Maybe Bool -> Maybe Bool -> Maybe OllamaConfig -> m ()
pushM :: forall (m :: * -> *).
MonadIO m =>
Text -> Maybe Bool -> Maybe Bool -> Maybe OllamaConfig -> m ()
pushM Text
t Maybe Bool
insec Maybe Bool
s Maybe OllamaConfig
mbCfg = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Bool -> Maybe Bool -> Maybe OllamaConfig -> IO ()
push Text
t Maybe Bool
insec Maybe Bool
s Maybe OllamaConfig
mbCfg