package dev.mokksy.aimocks.anthropic.lc4j

import dev.langchain4j.data.message.UserMessage.userMessage
import dev.langchain4j.kotlin.model.chat.StreamingChatModelReply
import dev.langchain4j.kotlin.model.chat.chatFlow
import dev.langchain4j.model.anthropic.AnthropicStreamingChatModel
import dev.langchain4j.model.chat.request.ChatRequest
import dev.langchain4j.model.chat.response.ChatResponse
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler
import dev.langchain4j.model.output.FinishReason
import dev.mokksy.aimocks.anthropic.AbstractAnthropicIntegrationTest
import dev.mokksy.aimocks.anthropic.anthropic
import io.kotest.matchers.equals.shouldBeEqual
import io.kotest.matchers.nulls.shouldNotBeNull
import io.kotest.matchers.shouldBe
import kotlinx.coroutines.flow.buffer
import kotlinx.coroutines.flow.flow
import org.awaitility.kotlin.await
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.fail
import java.util.concurrent.ConcurrentLinkedQueue
import java.util.concurrent.atomic.AtomicReference

private const val FLOW_BUFFER_SIZE = 8096

internal class StreamingChatCompletionLc4jTest : AbstractAnthropicIntegrationTest() {
    private lateinit var model: AnthropicStreamingChatModel

    @BeforeEach
    fun setupModel() {
        model =
            AnthropicStreamingChatModel
                .builder()
                .apiKey("foo")
                .baseUrl(anthropic.baseUrl() + "/v1")
                .modelName(modelName)
                .build()
    }

    @Test
    suspend fun `Should respond to Streaming Chat Completion`() {
        val userMessage = "Please tell me what we need"
        anthropic.messages {
            model = modelName
            userMessageContains("what we need")
        } respondsStream {
            responseChunks = listOf("All", " we", " need", " is", " Love")
        }

        verifyStreamingKotlinFlow(userMessage, "All we need is Love")
    }

    @Test
    fun `Should respond to Streaming Chat Completion with Flow`() {
        val userMessage = "Please tell me what's you see in the sea"
        anthropic.messages {
            model = modelName
            userMessageContains("what's you see")
        } respondsStream {
            responseFlow =
                flow {
                    emit("Yellow")
                    emit(" submarine")
                }
        }

        verifyLC4JStreamingCall(userMessage, "Yellow submarine")
    }

    private fun verifyLC4JStreamingCall(
        userMessage: String,
        expectedResponse: String,
    ) {
        val partialResults = ConcurrentLinkedQueue<String>()
        val errorRef = AtomicReference<Throwable?>()
        val responseRef = AtomicReference<ChatResponse>()

        model.chat(
            ChatRequest
                .builder()
                .messages(userMessage(userMessage))
                .build(),
            object : StreamingChatResponseHandler {
                override fun onCompleteResponse(completeResponse: ChatResponse) {
                    logger.info { "Received CompleteResponse: $completeResponse" }
                    responseRef.set(completeResponse)
                }

                override fun onPartialResponse(partialResponse: String) {
                    logger.info { "Received partial response: $partialResponse" }
                    partialResults.add(partialResponse)
                }

                override fun onError(error: Throwable) {
                    logger.info { "Received error: $error" }
                    fail("Unexpected error", error)
                }
            },
        )

        await.untilAsserted {
            errorRef.get()?.let { throw AssertionError("Unexpected error", it) }
            responseRef.get().shouldNotBeNull()
        }
        val chatResponse = responseRef.get()
        chatResponse.finishReason() shouldBe FinishReason.STOP
        chatResponse.aiMessage().text() shouldBeEqual expectedResponse
    }

    private suspend fun verifyStreamingKotlinFlow(
        userMessage: String,
        expectedResponse: String,
    ) {
        val result = ConcurrentLinkedQueue<String>()
        val finishReason = AtomicReference<FinishReason>()
        model
            .chatFlow {
                messages += userMessage(userMessage)
            }.buffer(capacity = FLOW_BUFFER_SIZE)
            .collect {
                when (it) {
                    is StreamingChatModelReply.PartialResponse -> {
                        result.add(it.partialResponse)
                        logger.info { "token = ${it.partialResponse}" }
                    }

                    is StreamingChatModelReply.CompleteResponse -> {
                        logger.info { "Completed: $it" }
                        finishReason.set(it.response.finishReason())
                    }

                    is StreamingChatModelReply.Error -> {
                        logger.info { "Error: $it" }
                        it.cause.printStackTrace()
                    }
                }
            }

        await.untilAsserted {
            result.joinToString("") shouldBeEqual expectedResponse
        }
        finishReason.get() shouldBe FinishReason.STOP
    }
}
