{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module TypeChain.ChatModels.OpenAI (OpenAIChat(..), OpenAIChatModel(..), initOpenAIChat) where

import Control.Lens hiding ((.=))
import Control.Monad.Catch
import Control.Monad.IO.Class
import Control.Monad.State

import Data.Aeson
import Data.ByteString.Lazy (ByteString)
import Data.Functor (($>))
import Data.Maybe (fromMaybe)

import GHC.Generics (Generic)

import TypeChain.ChatModels.Types

import Network.HTTP.Simple
import Network.HTTP.Conduit

import qualified Data.ByteString.Lazy as BS

data OpenAIChatModel = GPT35Turbo | GPT4 | GPT4Turbo

instance Show OpenAIChatModel where 
    show GPT35Turbo = "gpt-3.5-turbo"
    show GPT4       = "gpt-4"
    show GPT4Turbo  = "gpt-4-turbo-preview"

instance ToJSON OpenAIChatModel where 
    toJSON = toJSON . show

data OpenAIChat = OpenAIChat { chatModel   :: OpenAIChatModel
                             , messages    :: Maybe [Message] -- ^ @Nothing@ = Do not remember messages
                             , temperature :: Float
                             , apiKey      :: ApiKey
                             } deriving Generic

instance ToJSON OpenAIChat where 
    toJSON model = object [ "model"       .= chatModel model
                          , "temperature" .= temperature model
                          , "messages"    .= messages model
                          ]

-- | A list of responses from OpenAI's GPT-3.5-Turbo model
data Choices = Choices { message       :: Message 
                       , finish_reason :: String 
                       , index         :: Int
                       } deriving Generic

instance FromJSON Choices 

-- | Minimal JSON response datatype from OpenAI's GPT-3.5-Turbo model
data OpenAIResponse = OpenAIResponse { model   :: String
                                     , choices :: [Choices]
                                     } deriving Generic


instance FromJSON OpenAIResponse where 
    parseJSON = withObject "OpenAIResponse" $ \o -> do 
        model   <- o .: "model"
        choices <- o .: "choices"
        return $ OpenAIResponse model choices

-- | Create an OpenAI chat model with default values 
--
-- Model: GPT-3.5-Turbo
--
-- Memorization: Enabled 
--
-- Temperature: 0.7 
--
-- ApiKey: <empty value>
initOpenAIChat :: OpenAIChat
initOpenAIChat = OpenAIChat { chatModel   = GPT35Turbo 
                            , messages    = Just []
                            , temperature = 0.7
                            , apiKey      = "MISSING-API-KEY"
                            }

mkOpenAIChatHeaders :: ApiKey -> RequestHeaders
mkOpenAIChatHeaders k = [("Content-Type", "application/json"), ("Authorization", "Bearer " <> k)]

mkGPT35TurboRequest :: MonadThrow m => OpenAIChat -> m Request
mkGPT35TurboRequest gpt = do 
    initReq <- parseRequest "https://api.openai.com/v1/chat/completions"
    return $ initReq { requestHeaders = mkOpenAIChatHeaders (apiKey gpt) 
                     , requestBody    = RequestBodyLBS (encode gpt)
                     , method         = "POST"
                     }

instance ChatModel OpenAIChat where 
    predicts model m = do 
        let msgs = toMsgList m
        model `memorizes` msgs

        gpt <- gets (view model)
        req <- mkGPT35TurboRequest gpt
        res <- httpLBS req

        case decode @OpenAIResponse (responseBody res) of 
            Nothing -> liftIO $ putStrLn "WARNING: Failed to decode OpenAIResponse" $> []
            Just r  -> let newMsgs = map message (choices r) in model `memorizes` newMsgs $> newMsgs

instance RememberingChatModel OpenAIChat where 
    setMemoryEnabledFor model status = do 
        m <- gets (view model)
        case (messages m, status) of 
            (Nothing, True) -> model %= \m -> m { messages = Just [] }
            (Just _, False) -> model %= \m -> m { messages = Nothing }
            _               -> return () -- Do nothing since the status is already matching

    forgetFor model = do 
        m <- gets (view model)
        case messages m of 
            Nothing -> return () -- Do nothing since we're already forgetting
            Just _  -> model %= \m -> m { messages = Just [] }

    memorizes model msgs = model %= \m -> m { messages = (++) <$> messages m <*> Just (toMsgList msgs) }
    rememberFor model = gets (fromMaybe [] . messages . view model)