package dev.mokksy.aimocks.openai.official.completions

import assertk.assertThat
import assertk.assertions.isEqualTo
import com.openai.core.JsonValue
import com.openai.models.chat.completions.ChatCompletionChunk
import com.openai.models.chat.completions.ChatCompletionCreateParams
import com.openai.models.chat.completions.ChatCompletionMessageParam
import com.openai.models.chat.completions.ChatCompletionStreamOptions
import com.openai.models.chat.completions.ChatCompletionUserMessageParam
import dev.mokksy.aimocks.openai.official.AbstractOpenaiTest
import dev.mokksy.aimocks.openai.openai
import io.kotest.matchers.shouldBe
import kotlinx.coroutines.flow.flow
import org.junit.jupiter.api.Test
import kotlin.time.Duration.Companion.milliseconds

internal class StreamingChatCompletionOpenaiTest : AbstractOpenaiTest() {
    @Test
    fun `Should respond to Streaming Chat Completion`() {
        openai.completion("openai-completions-list") {
            temperature = temperatureValue
            seed = seedValue
            model = modelName
            topP = topPValue
            maxTokens = maxCompletionTokensValue
            userMessageContains("OpenAI streaming list test")
        } respondsStream {
            responseChunks = listOf("All", " we", " need", " is", " Love")
            delay = 50.milliseconds
            delayBetweenChunks = 10.milliseconds
            finishReason = "stop"
        }

        verifyStreamingCall("Please run the OpenAI streaming list test")
    }

    @Test
    fun `Should respond to Streaming Chat Completion with Flow`() {
        openai.completion("openai-completions-flow") {
            temperature = temperatureValue
            seed = seedValue
            model = modelName
            maxTokens = maxCompletionTokensValue
            userMessageContains("OpenAI streaming flow test")
        } respondsStream {
            responseFlow =
                flow {
                    emit("All")
                    emit(" we")
                    emit(" need")
                    emit(" is")
                    emit(" Love")
                }
            delay = 60.milliseconds
            delayBetweenChunks = 15.milliseconds
            finishReason = "stop"
        }

        verifyStreamingCall("Please run the OpenAI streaming flow test")
    }

    private fun verifyStreamingCall(userMessage: String) {
        val params =
            @Suppress("DEPRECATION")
            ChatCompletionCreateParams
                .builder()
                .streamOptions(
                    ChatCompletionStreamOptions.builder().includeUsage(true).build(),
                ).temperature(temperatureValue)
                .topP(topPValue)
                .maxCompletionTokens(maxCompletionTokensValue)
                .seed(seedValue.toLong())
                .messages(
                    listOf(
                        ChatCompletionMessageParam.ofUser(
                            ChatCompletionUserMessageParam
                                .builder()
                                .role(JsonValue.from("user"))
                                .content(
                                    ChatCompletionUserMessageParam.Content.ofText(
                                        userMessage,
                                    ),
                                ).build(),
                        ),
                    ),
                ).model(modelName)
                .build()

        // when
        val result = StringBuffer()
        client
            .chat()
            .completions()
            .createStreaming(params)
            .use { messageStreamResponse ->
                messageStreamResponse
                    .stream()
                    .peek { println("Received: $it") }
                    .peek {
                        // validation
                        it.model() shouldBe modelName
                    }.flatMap { completion: ChatCompletionChunk ->
                        completion.choices().stream()
                    }.flatMap { choice: ChatCompletionChunk.Choice ->
                        choice.delta().content().stream()
                    }.forEach {
                        result.append(it)
                        print(it)
                    }
            }

        // then
        assertThat(result.toString()).isEqualTo("All we need is Love")
    }
}
