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

module Data.Ollama.Push
  ( -- * Push API
    push
  , pushOps
  ) 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
data PushOps = PushOps
  { PushOps -> Text
name :: Text
  , PushOps -> Maybe Bool
insecure :: Maybe Bool
  , PushOps -> Maybe Bool
stream :: Maybe Bool
  }
  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)

data PushResp = PushResp
  { PushResp -> Text
status :: Text
  , PushResp -> Maybe Text
digest :: Maybe Text
  , PushResp -> Maybe Int64
total :: Maybe Int64
  }
  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)

-- | Push a model with options
pushOps ::
  -- | Model name
  Text ->
  -- | Insecure
  Maybe Bool ->
  -- | Stream
  Maybe Bool ->
  IO ()
pushOps :: Text -> Maybe Bool -> Maybe Bool -> IO ()
pushOps 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/push")
  let reqBody :: PushOps
reqBody =
        PushOps
          { $sel:name:PushOps :: Text
name = Text
modelName
          , $sel:insecure:PushOps :: Maybe Bool
insecure = Maybe Bool
mInsecure
          , $sel:stream:PushOps :: 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 PushResp
eRes = ByteString -> Maybe PushResp
forall a. FromJSON a => ByteString -> Maybe a
decode (Method -> ByteString
BSL.fromStrict Method
bs) :: Maybe PushResp
          case Maybe PushResp
eRes of
            Maybe PushResp
Nothing -> String -> IO ()
putStrLn String
"Something went wrong"
            Just PushResp
res -> do
              if PushResp -> Text
status PushResp
res Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
"success"
                then do
                  let total' :: Int64
total' = Int64 -> Maybe Int64 -> Int64
forall a. a -> Maybe a -> a
fromMaybe Int64
0 (PushResp -> Maybe Int64
total PushResp
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'
                  IO ()
go
                else do
                  String -> IO ()
putStrLn String
"Completed"
    IO ()
go

-- Higher level API for Pull
-- This API is untested. Will test soon!

-- | Push a model
push ::
  -- | Model name
  Text ->
  IO ()
push :: Text -> IO ()
push Text
modelName = Text -> Maybe Bool -> Maybe Bool -> IO ()
pushOps Text
modelName Maybe Bool
forall a. Maybe a
Nothing Maybe Bool
forall a. Maybe a
Nothing