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

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

This module provides functions to create a new model in the Ollama API using either a model file
content or a file path. It includes both an IO-based function ('createModel') and a monadic version
('createModelM') for use in 'MonadIO' contexts. The create operation is performed via a POST request
to the "/api//pull" endpoint, with streaming support for progress updates.

Note: If both 'modelFile' and 'path' are provided, 'modelFile' takes precedence.

Example:

>>> createModel "myModel" (Just "FROM llama3\nPARAMETER temperature 0.8") (Just True) Nothing Nothing
Creating model...
Completed
-}
module Data.Ollama.Create
  ( -- * Create Model API
    createModel
  , createModelM
  ) 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)

-- | Configuration for creating a new model.
data CreateModelOps = CreateModelOps
  { CreateModelOps -> Text
name :: !Text
  -- ^ The name of the model to create.
  , CreateModelOps -> Maybe Text
modelFile :: !(Maybe Text)
  -- ^ Optional model file content (e.g., Modelfile text). Takes precedence over 'path'.
  , CreateModelOps -> Maybe Bool
stream :: !(Maybe Bool)
  -- ^ Optional flag to enable streaming progress updates.
  , CreateModelOps -> Maybe FilePath
path :: !(Maybe FilePath)
  -- ^ Optional file path to a Modelfile.
  }
  deriving (Int -> CreateModelOps -> ShowS
[CreateModelOps] -> ShowS
CreateModelOps -> FilePath
(Int -> CreateModelOps -> ShowS)
-> (CreateModelOps -> FilePath)
-> ([CreateModelOps] -> ShowS)
-> Show CreateModelOps
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CreateModelOps -> ShowS
showsPrec :: Int -> CreateModelOps -> ShowS
$cshow :: CreateModelOps -> FilePath
show :: CreateModelOps -> FilePath
$cshowList :: [CreateModelOps] -> ShowS
showList :: [CreateModelOps] -> ShowS
Show, CreateModelOps -> CreateModelOps -> Bool
(CreateModelOps -> CreateModelOps -> Bool)
-> (CreateModelOps -> CreateModelOps -> Bool) -> Eq CreateModelOps
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CreateModelOps -> CreateModelOps -> Bool
== :: CreateModelOps -> CreateModelOps -> Bool
$c/= :: CreateModelOps -> CreateModelOps -> Bool
/= :: CreateModelOps -> CreateModelOps -> Bool
Eq)

-- | Response type for model creation.
newtype CreateModelResp
  = -- | The status of the create operation (e.g., "success").
    CreateModelResp {CreateModelResp -> Text
status :: Text}
  deriving (Int -> CreateModelResp -> ShowS
[CreateModelResp] -> ShowS
CreateModelResp -> FilePath
(Int -> CreateModelResp -> ShowS)
-> (CreateModelResp -> FilePath)
-> ([CreateModelResp] -> ShowS)
-> Show CreateModelResp
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CreateModelResp -> ShowS
showsPrec :: Int -> CreateModelResp -> ShowS
$cshow :: CreateModelResp -> FilePath
show :: CreateModelResp -> FilePath
$cshowList :: [CreateModelResp] -> ShowS
showList :: [CreateModelResp] -> ShowS
Show, CreateModelResp -> CreateModelResp -> Bool
(CreateModelResp -> CreateModelResp -> Bool)
-> (CreateModelResp -> CreateModelResp -> Bool)
-> Eq CreateModelResp
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CreateModelResp -> CreateModelResp -> Bool
== :: CreateModelResp -> CreateModelResp -> Bool
$c/= :: CreateModelResp -> CreateModelResp -> Bool
/= :: CreateModelResp -> CreateModelResp -> Bool
Eq)

instance HasDone CreateModelResp where
  getDone :: CreateModelResp -> Bool
getDone CreateModelResp {Text
status :: CreateModelResp -> Text
status :: Text
..} = Text
status Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
"success"

instance ToJSON CreateModelOps where
  toJSON :: CreateModelOps -> Value
toJSON
    ( CreateModelOps
        Text
name_
        Maybe Text
modelFile_
        Maybe Bool
stream_
        Maybe FilePath
path_
      ) =
      [Pair] -> Value
object
        [ Key
"name" Key -> Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Text
name_
        , Key
"modelfile" Key -> Maybe Text -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Text
modelFile_
        , Key
"stream" Key -> Maybe Bool -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe Bool
stream_
        , Key
"path" Key -> Maybe FilePath -> Pair
forall v. ToJSON v => Key -> v -> Pair
forall e kv v. (KeyValue e kv, ToJSON v) => Key -> v -> kv
.= Maybe FilePath
path_
        ]

instance FromJSON CreateModelResp where
  parseJSON :: Value -> Parser CreateModelResp
parseJSON = FilePath
-> (Object -> Parser CreateModelResp)
-> Value
-> Parser CreateModelResp
forall a. FilePath -> (Object -> Parser a) -> Value -> Parser a
withObject FilePath
"CreateModelResp" ((Object -> Parser CreateModelResp)
 -> Value -> Parser CreateModelResp)
-> (Object -> Parser CreateModelResp)
-> Value
-> Parser CreateModelResp
forall a b. (a -> b) -> a -> b
$ \Object
v ->
    Text -> CreateModelResp
CreateModelResp
      (Text -> CreateModelResp) -> Parser Text -> Parser CreateModelResp
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
"status"

{- | Creates a new model using either model file content or a file path.

Sends a POST request to the "/api//pull" endpoint to create a model with the specified name.
The model can be defined either by 'modelFile' (Modelfile content as text) or 'path' (file path to a Modelfile).
If both are provided, 'modelFile' is used. Supports streaming progress updates if 'stream' is 'Just True'.
Prints progress messages to the console during creation.
-}
createModel ::
  -- | Model name
  Text ->
  -- | Optional model file content
  Maybe Text ->
  -- | Optional streaming flag
  Maybe Bool ->
  -- | Optional file path to a Modelfile
  Maybe FilePath ->
  -- | Optional 'OllamaConfig' (defaults to 'defaultOllamaConfig' if 'Nothing')
  Maybe OllamaConfig ->
  IO ()
createModel :: Text
-> Maybe Text
-> Maybe Bool
-> Maybe FilePath
-> Maybe OllamaConfig
-> IO ()
createModel
  Text
modelName
  Maybe Text
modelFile_
  Maybe Bool
stream_
  Maybe FilePath
path_
  Maybe OllamaConfig
mbConfig =
    IO (Either OllamaError CreateModelResp) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either OllamaError CreateModelResp) -> IO ())
-> IO (Either OllamaError CreateModelResp) -> IO ()
forall a b. (a -> b) -> a -> b
$
      Text
-> ByteString
-> Maybe CreateModelOps
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError CreateModelResp))
-> IO (Either OllamaError CreateModelResp)
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"
        ( CreateModelOps -> Maybe CreateModelOps
forall a. a -> Maybe a
Just (CreateModelOps -> Maybe CreateModelOps)
-> CreateModelOps -> Maybe CreateModelOps
forall a b. (a -> b) -> a -> b
$
            CreateModelOps
              { name :: Text
name = Text
modelName
              , modelFile :: Maybe Text
modelFile = Maybe Text
modelFile_
              , stream :: Maybe Bool
stream = Maybe Bool
stream_
              , path :: Maybe FilePath
path = Maybe FilePath
path_
              }
        )
        Maybe OllamaConfig
mbConfig
        ((CreateModelResp -> IO ())
-> IO ()
-> Response BodyReader
-> IO (Either OllamaError CreateModelResp)
forall a.
(HasDone a, FromJSON a) =>
(a -> IO ())
-> IO () -> Response BodyReader -> IO (Either OllamaError a)
commonStreamHandler CreateModelResp -> IO ()
onToken IO ()
onComplete)
    where
      onToken :: CreateModelResp -> IO ()
      onToken :: CreateModelResp -> IO ()
onToken CreateModelResp
_ = do
        FilePath -> IO ()
putStrLn FilePath
"Creating model..."

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

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

Lifts the 'createModel' function into a 'MonadIO' context, allowing it to be used in monadic computations.
-}
createModelM ::
  MonadIO m =>
  Text ->
  Maybe Text ->
  Maybe Bool ->
  Maybe FilePath ->
  Maybe OllamaConfig ->
  m ()
createModelM :: forall (m :: * -> *).
MonadIO m =>
Text
-> Maybe Text
-> Maybe Bool
-> Maybe FilePath
-> Maybe OllamaConfig
-> m ()
createModelM Text
m Maybe Text
mf Maybe Bool
s Maybe FilePath
p 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 Text
-> Maybe Bool
-> Maybe FilePath
-> Maybe OllamaConfig
-> IO ()
createModel Text
m Maybe Text
mf Maybe Bool
s Maybe FilePath
p Maybe OllamaConfig
mbCfg