package dev.mokksy.aimocks.ollama.lc4j

import assertk.assertThat
import assertk.assertions.isGreaterThanOrEqualTo
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.ollama.OllamaStreamingChatModel
import dev.mokksy.aimocks.ollama.AbstractMockOllamaTest
import dev.mokksy.aimocks.ollama.mockOllama
import io.kotest.matchers.shouldBe
import kotlinx.coroutines.flow.asFlow
import kotlinx.coroutines.flow.map
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test
import kotlin.test.fail
import kotlin.time.Duration.Companion.milliseconds
import kotlin.time.measureTime

internal class OllamaStreamingChatCompletionLc4jTest : AbstractMockOllamaTest() {
    private lateinit var model: OllamaStreamingChatModel

    @BeforeEach
    fun setupModel() {
        model =
            OllamaStreamingChatModel
                .builder()
                .customHeaders(
                    mapOf(
                        "Content-Type" to "application/json", // add lost header
                    ),
                ).baseUrl(mockOllama.baseUrl())
                .modelName(modelName)
                .temperature(temperatureValue)
                .logRequests(true)
                .logResponses(true)
                .topP(topPValue)
                .seed(seedValue)
                .topK(topKValue.toInt())
                .build()
    }

    @Test
    suspend fun `Should respond to Streaming Chat Completion`() {
        // Configure mock response
        val expectedResponse = "Hello, how can I help you today?"

        val delayBetweenChunks = 100.milliseconds
        val initialDelay = 500.milliseconds

        mockOllama.chat("ollama-lc4j-streaming-chat-$seedValue") {
            model = modelName
            seed = seedValue
            temperature = temperatureValue
            userMessageContains("Lc4j streaming test")
            topP = topPValue
            topK = topKValue
            stream = true
        } respondsStream {
            responseFlow =
                expectedResponse
                    .splitToSequence(" ")
                    .asFlow()
                    .map { "$it " }
            this.delay = initialDelay
            this.delayBetweenChunks = delayBetweenChunks
        }

        // Use langchain4j Ollama client to send a request
        val tokens = mutableListOf<String>()
        val executionTime =
            measureTime {
                model
                    .chatFlow {
                        messages += userMessage("Run the Lc4j streaming test")
                    }.collect { reply ->
                        when (reply) {
                            is StreamingChatModelReply.PartialResponse -> {
                                tokens += reply.partialResponse
                            }

                            is StreamingChatModelReply.CompleteResponse -> {
                                reply.response.modelName() shouldBe modelName
                            }

                            is StreamingChatModelReply.Error -> {
                                fail("Error: $reply", reply.cause)
                            }
                        }
                    }
                tokens
                    .joinToString("")
                    .removeSuffix(" ") shouldBe expectedResponse
            }

        assertThat(executionTime).isGreaterThanOrEqualTo(
            initialDelay + delayBetweenChunks * tokens.size,
        )
    }
}
