{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}

module TypeChain.ChatModels.PromptTemplate (ToPrompt(..), makeTemplate, user, assistant, system) where

import Data.Char (toLower)
import Data.List (nub)

import Language.Haskell.TH

import TypeChain.ChatModels.Types

data TemplateToken = ConstString String | Var String deriving Eq

type PromptTemplate = (Q Exp, [Name])

-- | Typeclass used to convert generated record types into a list of messages.
--
-- Instances of this typeclass are generated by the `makeTemplate` function and 
-- should only be used if you need to construct a prompt manually.
class ToPrompt a where 

    -- | Return the list of messages that should be used as the prompt.
    toPrompt :: a -> [Message]

-- | Convert a String into a compile prompt template for the `makeTemplate` function.
--
-- This particular function is for user messages.
user :: String -> Q PromptTemplate
user xs = toTemplate xs [| UserMessage |]

-- | Convert a String into a compile prompt template for the `makeTemplate` function.
--
-- This particular function is for assistant messages.
assistant :: String -> Q PromptTemplate
assistant xs = toTemplate xs [| AssistantMessage |]

-- | Convert a String into a compile prompt template for the `makeTemplate` function. 
--
-- This particular function is for system messages.
system :: String -> Q PromptTemplate
system xs = toTemplate xs [| SystemMessage |]

toTemplate :: String -> Q Exp -> Q PromptTemplate
toTemplate xs f = do 
    let tempParam = mkName "template"

    let tokens = parseTemplateTokens xs
        expr   = tokensToExpr tempParam tokens
        names  = map mkName $ nub $ getVarTokens tokens

    return (appE f expr, names)

parseTemplateTokens :: String -> [TemplateToken]
parseTemplateTokens [] = [] 
parseTemplateTokens ('{':xs) = Var first : parseTemplateTokens rest
    where (first, tail -> rest) = break (== '}') xs
parseTemplateTokens (x:xs) = ConstString (x : first) : parseTemplateTokens rest
    where (first, rest) = break (== '{') xs

getVarTokens :: [TemplateToken] -> [String]
getVarTokens [] = [] 
getVarTokens (Var x : xs) = x : getVarTokens xs
getVarTokens (_     : xs) = getVarTokens xs

tokensToExpr :: Name -> [TemplateToken] -> Q Exp
tokensToExpr _ [] = [| "" |]
tokensToExpr name (ConstString x : xs) = [| x ++ $(tokensToExpr name xs) |]
tokensToExpr name (Var x : xs) = appE [| (++) |] (appE (varE $ mkName x) (varE name)) `appE` tokensToExpr name xs


-- | Given a typename and a list of messages, generate a data type and a function to construct it.
--
-- Example: `makeTemplate "Translate" [system "translate {a} to {b}.", user "{text}"]`
--
-- This generates a record named @Translate@ with fields @a@, @b@, and @text@. 
-- It also generates a function @mkTranslate :: String -> String -> String -> [Message]@.
-- To allow for quick and easy construction of the prompt if needed. Otherwise, you can use the 
-- generated data type in conjunction with the `toPrompt` function to be more explicit.
--
-- See the example on the repo's README.md for an example of what the generated code looks like.
makeTemplate :: String -> [Q PromptTemplate] -> Q [Dec] 
makeTemplate name xs = do 
    let typeName = mkName name
        funcName = mkName ("mk" ++ name)

    (exps, concat -> nub -> names) <- unzip <$> sequence xs

    exps'   <- sequence exps
    varbang <- bang sourceNoUnpack sourceStrict

    let recordFields      = map (, varbang, ConT ''String) names :: [VarBangType]
        promptFunc        = FunD 'toPrompt [Clause [VarP $ mkName "template"] (NormalB $ ListE exps') []]
        filledConstructor = foldl AppE (ConE typeName) (map VarE names)
        mkFuncClause      = Clause (map VarP names) (NormalB $ AppE (VarE 'toPrompt) filledConstructor) []

    return [ DataD     [] typeName [] Nothing [RecC typeName recordFields] []
           , InstanceD Nothing [] (AppT (ConT ''ToPrompt) (ConT typeName)) [promptFunc]
           , FunD      funcName [mkFuncClause]
           ]