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

module Lib where

import Control.Monad (void)
import Data.List.NonEmpty (NonEmpty((:|)))
import Data.Maybe (fromMaybe)
import Data.Ollama.Chat qualified as Chat
import Data.Text.IO qualified as T
import Ollama (GenerateOps(..), Role(..), chat, defaultChatOps, defaultGenerateOps, generate)
import Ollama qualified

main :: IO ()
main :: IO ()
main = do
  -- Example 1: Streamed Text Generation
  -- This example demonstrates how to generate text using a model and stream the output directly to the console.
  -- The `stream` option enables processing of each chunk of the response as it arrives.
  IO (Either String GenerateResponse) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either String GenerateResponse) -> IO ())
-> IO (Either String GenerateResponse) -> IO ()
forall a b. (a -> b) -> a -> b
$
    GenerateOps -> IO (Either String GenerateResponse)
generate
      GenerateOps
defaultGenerateOps
        { modelName = "llama3.2"
        , prompt = "what is functional programming?"
        , stream = Just (T.putStr . Ollama.response_, pure ())
        }

  -- Example 2: Non-streamed Text Generation
  -- This example shows how to generate text and handle the complete response.
  -- The result is either an error message or the generated text.
  Either String GenerateResponse
eRes <-
    GenerateOps -> IO (Either String GenerateResponse)
generate
      GenerateOps
defaultGenerateOps
        { modelName = "llama3.2"
        , prompt = "What is 2+2?"
        }
  case Either String GenerateResponse
eRes of
    Left String
e -> String -> IO ()
putStrLn String
e
    Right Ollama.GenerateResponse {Bool
Maybe Int64
Text
UTCTime
$sel:response_:GenerateResponse :: GenerateResponse -> Text
model :: Text
createdAt :: UTCTime
response_ :: Text
done :: Bool
totalDuration :: Maybe Int64
loadDuration :: Maybe Int64
promptEvalCount :: Maybe Int64
promptEvalDuration :: Maybe Int64
evalCount :: Maybe Int64
evalDuration :: Maybe Int64
$sel:model:GenerateResponse :: GenerateResponse -> Text
$sel:createdAt:GenerateResponse :: GenerateResponse -> UTCTime
$sel:done:GenerateResponse :: GenerateResponse -> Bool
$sel:totalDuration:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:loadDuration:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:promptEvalCount:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:promptEvalDuration:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:evalCount:GenerateResponse :: GenerateResponse -> Maybe Int64
$sel:evalDuration:GenerateResponse :: GenerateResponse -> Maybe Int64
..} -> Text -> IO ()
T.putStrLn Text
response_

  -- Example 3: Chat with Streaming
  -- This example demonstrates setting up a chat session with streaming enabled.
  -- As messages are received, they are printed to the console.
  let msg :: Message
msg = Role -> Text -> Maybe [Text] -> Message
Ollama.Message Role
User Text
"What is functional programming?" Maybe [Text]
forall a. Maybe a
Nothing
      defaultMsg :: Message
defaultMsg = Role -> Text -> Maybe [Text] -> Message
Ollama.Message Role
User Text
"" Maybe [Text]
forall a. Maybe a
Nothing
  IO (Either String ChatResponse) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either String ChatResponse) -> IO ())
-> IO (Either String ChatResponse) -> IO ()
forall a b. (a -> b) -> a -> b
$
    ChatOps -> IO (Either String ChatResponse)
chat
      ChatOps
defaultChatOps
        { Chat.chatModelName = "llama3.2"
        , Chat.messages = msg :| []
        , Chat.stream =
            Just (T.putStr . Chat.content . fromMaybe defaultMsg . Chat.message, pure ())
        }

  -- Example 4: Non-streamed Chat
  -- Here, we handle a complete chat response, checking for potential errors.
  Either String ChatResponse
eRes1 <-
    ChatOps -> IO (Either String ChatResponse)
chat
      ChatOps
defaultChatOps
        { Chat.chatModelName = "llama3.2"
        , Chat.messages = msg :| []
        }
  case Either String ChatResponse
eRes1 of
    Left String
e -> String -> IO ()
putStrLn String
e
    Right ChatResponse
r -> do
      let mMessage :: Maybe Message
mMessage = ChatResponse -> Maybe Message
Ollama.message ChatResponse
r
      case Maybe Message
mMessage of
        Maybe Message
Nothing -> String -> IO ()
putStrLn String
"Something went wrong"
        Just Message
res -> Text -> IO ()
T.putStrLn (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Message -> Text
Ollama.content Message
res

  -- Example 5: Check Model Status (ps)
  -- This example checks the status of models using the `ps` function.
  -- It outputs the status or details of the available models.
  Maybe RunningModels
res <- IO (Maybe RunningModels)
Ollama.ps
  Maybe RunningModels -> IO ()
forall a. Show a => a -> IO ()
print Maybe RunningModels
res

  -- Example 6: Simple Embedding
  -- This demonstrates how to request embeddings for a given text using a specific model.
  IO (Maybe EmbeddingResp) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe EmbeddingResp) -> IO ())
-> IO (Maybe EmbeddingResp) -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> Text -> IO (Maybe EmbeddingResp)
Ollama.embedding Text
"llama3.1" Text
"What is 5+2?"

  -- Example 7: Embedding with Options
  -- This example uses the `embeddingOps` function, allowing for additional configuration like options and streaming.
  IO (Maybe EmbeddingResp) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe EmbeddingResp) -> IO ())
-> IO (Maybe EmbeddingResp) -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
-> Text -> Maybe Bool -> Maybe Text -> IO (Maybe EmbeddingResp)
Ollama.embeddingOps Text
"llama3.1" Text
"What is 5+2?" Maybe Bool
forall a. Maybe a
Nothing Maybe Text
forall a. Maybe a
Nothing

{-
Scotty example:
{-# LANGUAGE OverloadedStrings #-}

module Main where

import Web.Scotty
import Control.Monad.IO.Class (liftIO)
import Data.Aeson (FromJSON, ToJSON)
import Data.Text (Text)
import Data.Text qualified as T
import Database.SQLite.Simple
import Ollama (GenerateOps(..), defaultGenerateOps, generate)
import Data.Maybe (fromRight)

data PromptInput = PromptInput
  { conversation_id :: Int
  , prompt :: Text
  } deriving (Show, Generic)

instance FromJSON PromptInput
instance ToJSON PromptInput

main :: IO ()
main = do
  conn <- open "chat.db"
  execute_ conn "CREATE TABLE IF NOT EXISTS conversation (convo_id INTEGER PRIMARY KEY, convo_title TEXT)"
  execute_ conn "CREATE TABLE IF NOT EXISTS chats (chat_id INTEGER PRIMARY KEY, convo_id INTEGER, role TEXT, message TEXT, FOREIGN KEY(convo_id) REFERENCES conversation(convo_id))"
  
  scotty 3000 $ do
    post "/chat" $ do
      p <- jsonData :: ActionM PromptInput
      let cId = conversation_id p
      let trimmedP = T.dropEnd 3 $ T.drop 3 $ prompt p
      newConvoId <- case cId of
        -1 -> do
          liftIO $ execute conn "INSERT INTO conversation (convo_title) VALUES (?)" (Only ("latest title" :: String))
          [Only convoId] <- liftIO $ query_ conn "SELECT last_insert_rowid()" :: ActionM [Only Int]
          pure convoId
        _ -> pure cId

      liftIO $ execute conn "INSERT INTO chats (convo_id, role, message) VALUES (?, 'user', ?)" (newConvoId, trimmedP)
      
      stream $ \sendChunk flush -> do
        eRes <- generate defaultGenerateOps
                { modelName = "llama3.2"
                , prompt = prompt p
                , stream = Just (sendChunk . T.pack, flush)
                }
        case eRes of
            Left e -> return ()
            Right r -> do
                let res = response_ r
                liftIO $ execute conn "INSERT INTO chats (convo_id, role, message) VALUES (?, 'ai', ?)" (newConvoId, res)
-}