{-# LANGUAGE OverloadedStrings #-}

module Test.Langchain.Memory.Core (tests) where

import Test.Tasty
import Test.Tasty.HUnit

import Langchain.LLM.Core (Message (..), Role (..), defaultMessageData)
import Langchain.Memory.Core
import Langchain.Runnable.Core

import qualified Data.List.NonEmpty as NE
import Data.Text (Text)
import Langchain.Error (toString)

systemMsg :: Text -> Message
systemMsg text = Message System text defaultMessageData

userMsg :: Text -> Message
userMsg text = Message User text defaultMessageData

aiMsg :: Text -> Message
aiMsg text = Message Assistant text defaultMessageData

utilityTests :: TestTree
utilityTests =
  testGroup
    "Utility Functions Tests"
    [ testCase "initialChatMessage should create chat with system message" $ do
        let prompt = "You are a helpful assistant"
            result = initialChatMessage prompt
        NE.length result @?= 1
        NE.head result @?= systemMsg prompt
    , testCase "trimChatMessage should keep specified number of messages" $ do
        let msgs =
              NE.fromList
                [ systemMsg "System"
                , userMsg "User1"
                , aiMsg "AI1"
                , userMsg "User2"
                ]
            trimmed = trimChatMessage 2 msgs
        NE.length trimmed @?= 2
        NE.toList trimmed @?= [aiMsg "AI1", userMsg "User2"]
    , testCase "trimChatMessage should keep all messages if n >= length" $ do
        let msgs = NE.fromList [systemMsg "System", userMsg "User1"]
            trimmed = trimChatMessage 3 msgs
        NE.length trimmed @?= 2
        NE.toList trimmed @?= [systemMsg "System", userMsg "User1"]
    , testCase "trimChatMessage should handle minimum size of 1" $ do
        let msgs = NE.fromList [systemMsg "System", userMsg "User1", aiMsg "AI1"]
            trimmed = trimChatMessage 1 msgs
        NE.length trimmed @?= 1
        NE.toList trimmed @?= [aiMsg "AI1"]
    , testCase "addAndTrim should add message and trim history" $ do
        let msgs = NE.fromList [systemMsg "System", userMsg "User1", aiMsg "AI1"]
            newMsg = userMsg "User2"
            result = addAndTrim 2 newMsg msgs
        NE.length result @?= 2
        NE.toList result @?= [aiMsg "AI1", userMsg "User2"]
    ]

windowBufferMemoryTests :: TestTree
windowBufferMemoryTests =
  testGroup
    "WindowBufferMemory Tests"
    [ testCase "messages should return current messages" $ do
        let initialMsgs = NE.fromList [systemMsg "System"]
            memory = WindowBufferMemory 3 initialMsgs
        result <- messages memory
        case result of
          Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
          Right msgs -> msgs @?= initialMsgs
    , testCase "addMessage should add message when under capacity" $ do
        let initialMsgs = NE.fromList [systemMsg "System"]
            memory = WindowBufferMemory 3 initialMsgs
            newMsg = userMsg "User1"
        result <- addMessage memory newMsg
        case result of
          Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
          Right newMemory -> do
            msgsResult <- messages newMemory
            case msgsResult of
              Left err ->
                assertFailure $
                  "Expected Right but got Left: " ++ toString err
              Right msgs ->
                NE.toList msgs
                  @?= [ systemMsg "System"
                      , userMsg "User1"
                      ]
    , testCase "addMessage should maintain max window size" $ do
        let initialMsgs =
              NE.fromList
                [ systemMsg "System"
                , userMsg "User1"
                , aiMsg "AI1"
                ]
            memory = WindowBufferMemory 3 initialMsgs
            newMsg = userMsg "User2"
        result <- addMessage memory newMsg
        case result of
          Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
          Right newMemory -> do
            msgsResult <- messages newMemory
            case msgsResult of
              Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
              Right msgs -> do
                NE.length msgs @?= 3
                NE.toList msgs
                  @?= [ systemMsg "System"
                      , aiMsg "AI1"
                      , userMsg "User2"
                      ]
    , testCase "addUserMessage should add message with User role" $ do
        let initialMsgs = NE.fromList [systemMsg "System"]
            memory = WindowBufferMemory 3 initialMsgs
        result <- addUserMessage memory "Hello"
        case result of
          Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
          Right newMemory -> do
            msgsResult <- messages newMemory
            case msgsResult of
              Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
              Right msgs -> do
                NE.length msgs @?= 2
                NE.toList msgs @?= [systemMsg "System", userMsg "Hello"]
    , testCase "addAiMessage should add message with Assistant role" $ do
        let initialMsgs = NE.fromList [systemMsg "System"]
            memory = WindowBufferMemory 3 initialMsgs
        result <- addAiMessage memory "I can help"
        case result of
          Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
          Right newMemory -> do
            msgsResult <- messages newMemory
            case msgsResult of
              Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
              Right msgs -> do
                NE.length msgs @?= 2
                NE.toList msgs
                  @?= [ systemMsg "System"
                      , aiMsg "I can help"
                      ]
    , testCase "clear should reset to just system message" $ do
        let initialMsgs =
              NE.fromList
                [ systemMsg "System"
                , userMsg "User1"
                , aiMsg "AI1"
                ]
            memory = WindowBufferMemory 3 initialMsgs
        result <- clear memory
        case result of
          Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
          Right newMemory -> do
            msgsResult <- messages newMemory
            case msgsResult of
              Left err ->
                assertFailure $
                  "Expected Right but got Left: "
                    ++ toString err
              Right msgs -> do
                NE.length msgs @?= 1
                NE.head msgs @?= systemMsg "You are an AI model"
    ]

runnableTests :: TestTree
runnableTests =
  testGroup
    "Runnable Instance Tests"
    [ testCase "invoke should add user message" $ do
        let initialMsgs = NE.fromList [systemMsg "System"]
            memory = WindowBufferMemory 3 initialMsgs
        result <- invoke memory "Test input"
        case result of
          Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
          Right newMemory -> do
            msgsResult <- messages newMemory
            case msgsResult of
              Left err -> assertFailure $ "Expected Right but got Left: " ++ toString err
              Right msgs -> do
                NE.length msgs @?= 2
                NE.toList msgs @?= [systemMsg "System", userMsg "Test input"]
    ]

tests :: TestTree
tests =
  testGroup
    "Langchain.Memory.Core Tests"
    [ utilityTests
    , windowBufferMemoryTests
    , runnableTests
    ]