package dev.ai4j.openai4j.chat;

import dev.ai4j.openai4j.OpenAiClient;
import dev.ai4j.openai4j.RateLimitAwareTest;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.EnumSource;

import java.net.URL;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

import static dev.ai4j.openai4j.chat.ChatCompletionModel.GPT_4O;
import static dev.ai4j.openai4j.chat.ChatCompletionTest.*;
import static dev.ai4j.openai4j.chat.FunctionCallUtil.argument;
import static dev.ai4j.openai4j.chat.FunctionCallUtil.argumentsAsMap;
import static dev.ai4j.openai4j.chat.ResponseFormatType.JSON_OBJECT;
import static dev.ai4j.openai4j.chat.ResponseFormatType.TEXT;
import static dev.ai4j.openai4j.chat.ToolType.FUNCTION;
import static java.util.Collections.singletonMap;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.params.provider.EnumSource.Mode.EXCLUDE;

class ChatCompletionAsyncTest extends RateLimitAwareTest {

    private final OpenAiClient client = OpenAiClient.builder()
            .baseUrl(System.getenv("OPENAI_BASE_URL"))
            .openAiApiKey(System.getenv("OPENAI_API_KEY"))
            .logRequests()
            .logResponses()
            .build();

    @Test
    void testSimpleApi() throws Exception {

        // when
        CompletableFuture<String> future = new CompletableFuture<>();

        client.chatCompletion(USER_MESSAGE)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        String response = future.get(30, SECONDS);

        // then
        assertThat(response).containsIgnoringCase("hello world");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testCustomizableApi(ChatCompletionModel model) throws Exception {

        // given
        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .addSystemMessage(SYSTEM_MESSAGE)
                .addUserMessage(USER_MESSAGE)
                .temperature(1.0)
                .topP(0.1)
                .n(1)
                .stream(false)
                .stop("one", "two")
                .maxCompletionTokens(3)
                .presencePenalty(0.0)
                .frequencyPenalty(0.0)
                .logitBias(singletonMap("50256", -100))
                .user("Klaus")
                .responseFormat(TEXT)
                .seed(42)
                .build();

        CompletableFuture<ChatCompletionResponse> future = new CompletableFuture<>();

        // when
        client.chatCompletion(request)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        ChatCompletionResponse response = future.get(30, SECONDS);

        // then
        assertThat(response.choices()).hasSize(1);
        assertThat(response.choices().get(0).message().content()).containsIgnoringCase("hello world");

        assertThat(response.content()).containsIgnoringCase("hello world");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testTools(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather in Boston in Celsius?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .tools(WEATHER_TOOL)
                .build();

        CompletableFuture<ChatCompletionResponse> future = new CompletableFuture<>();

        // when
        client.chatCompletion(request)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        ChatCompletionResponse response = future.get(30, SECONDS);

        // then
        AssistantMessage assistantMessage = response.choices().get(0).message();
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.functionCall()).isNull();
        assertThat(assistantMessage.toolCalls()).isNotNull().hasSize(1);

        ToolCall toolCall = assistantMessage.toolCalls().get(0);
        assertThat(toolCall.id()).isNotBlank();
        assertThat(toolCall.type()).isEqualTo(FUNCTION);
        assertThat(toolCall.function()).isNotNull();

        FunctionCall functionCall = toolCall.function();
        assertThat(functionCall.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall.arguments()).isNotBlank();

        Map<String, Object> arguments = argumentsAsMap(functionCall.arguments());
        assertThat(arguments).hasSizeBetween(1, 2);
        assertThat(arguments.get("location").toString()).contains("Boston");

        // given
        String location = argument("location", functionCall);
        String unit = argument("unit", functionCall);
        String currentWeather = currentWeather(location, unit);
        ToolMessage toolMessage = ToolMessage.from(toolCall.id(), currentWeather);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, toolMessage)
                .build();

        // when
        ChatCompletionResponse secondResponse = client.chatCompletion(secondRequest).execute();

        // then
        assertThat(secondResponse.content()).contains("11");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testFunctions(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather in Boston in Celsius?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .functions(WEATHER_FUNCTION)
                .build();

        CompletableFuture<ChatCompletionResponse> future = new CompletableFuture<>();

        // when
        client.chatCompletion(request)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        ChatCompletionResponse response = future.get(30, SECONDS);

        // then
        AssistantMessage assistantMessage = response.choices().get(0).message();
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.toolCalls()).isNull();

        FunctionCall functionCall = assistantMessage.functionCall();
        assertThat(functionCall.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall.arguments()).isNotBlank();

        Map<String, Object> arguments = argumentsAsMap(functionCall.arguments());
        assertThat(arguments).hasSizeBetween(1, 2);
        assertThat(arguments.get("location").toString()).contains("Boston");

        // given
        String location = argument("location", functionCall);
        String unit = argument("unit", functionCall);
        String currentWeather = currentWeather(location, unit);
        FunctionMessage functionMessage = FunctionMessage.from(functionCall.name(), currentWeather);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, functionMessage)
                .build();

        // when
        ChatCompletionResponse secondResponse = client.chatCompletion(secondRequest).execute();

        // then
        assertThat(secondResponse.content()).contains("11");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testToolChoice(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather in Boston in Celsius?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .tools(WEATHER_TOOL)
                .toolChoice(WEATHER_TOOL_NAME)
                .build();

        CompletableFuture<ChatCompletionResponse> future = new CompletableFuture<>();

        // when
        client.chatCompletion(request)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        ChatCompletionResponse response = future.get(30, SECONDS);

        // then
        AssistantMessage assistantMessage = response.choices().get(0).message();
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.toolCalls()).isNotNull().hasSize(1);

        ToolCall toolCall = assistantMessage.toolCalls().get(0);
        assertThat(toolCall.id()).isNotBlank();
        assertThat(toolCall.type()).isEqualTo(FUNCTION);
        assertThat(toolCall.function()).isNotNull();

        FunctionCall functionCall = toolCall.function();
        assertThat(functionCall.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall.arguments()).isNotBlank();

        Map<String, Object> arguments = argumentsAsMap(functionCall.arguments());
        assertThat(arguments).hasSizeBetween(1, 2);
        assertThat(arguments.get("location").toString()).contains("Boston");

        // given
        String location = argument("location", functionCall);
        String unit = argument("unit", functionCall);
        String currentWeather = currentWeather(location, unit);
        ToolMessage toolMessage = ToolMessage.from(toolCall.id(), currentWeather);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, toolMessage)
                .build();

        // when
        ChatCompletionResponse secondResponse = client.chatCompletion(secondRequest).execute();

        // then
        assertThat(secondResponse.content()).contains("11");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testFunctionChoice(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather in Boston in Celsius?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .functions(WEATHER_FUNCTION)
                .functionCall(WEATHER_TOOL_NAME)
                .build();

        CompletableFuture<ChatCompletionResponse> future = new CompletableFuture<>();

        // when
        client.chatCompletion(request)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        ChatCompletionResponse response = future.get(30, SECONDS);

        // then
        AssistantMessage assistantMessage = response.choices().get(0).message();
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.toolCalls()).isNull();

        FunctionCall functionCall = assistantMessage.functionCall();
        assertThat(functionCall.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall.arguments()).isNotBlank();

        Map<String, Object> arguments = argumentsAsMap(functionCall.arguments());
        assertThat(arguments).hasSizeBetween(1, 2);
        assertThat(arguments.get("location").toString()).contains("Boston");

        // given
        String location = argument("location", functionCall);
        String unit = argument("unit", functionCall);
        String currentWeather = currentWeather(location, unit);
        FunctionMessage functionMessage = FunctionMessage.from(functionCall.name(), currentWeather);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, functionMessage)
                .build();

        // when
        ChatCompletionResponse secondResponse = client.chatCompletion(secondRequest).execute();

        // then
        assertThat(secondResponse.content()).contains("11");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
            "GPT_4", "GPT_4_0613", // Does not support parallel tools
    })
    void testParallelTools(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather like in Boston and Madrid?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .tools(WEATHER_TOOL)
                .build();

        CompletableFuture<ChatCompletionResponse> future = new CompletableFuture<>();

        // when
        client.chatCompletion(request)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        ChatCompletionResponse response = future.get(30, SECONDS);

        // then
        AssistantMessage assistantMessage = response.choices().get(0).message();
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.toolCalls()).isNotNull().hasSize(2);

        ToolCall toolCall1 = assistantMessage.toolCalls().get(0);
        assertThat(toolCall1.id()).isNotBlank();
        assertThat(toolCall1.type()).isEqualTo(FUNCTION);
        assertThat(toolCall1.function()).isNotNull();

        FunctionCall functionCall1 = toolCall1.function();
        assertThat(functionCall1.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall1.arguments()).isNotBlank();

        Map<String, Object> arguments1 = argumentsAsMap(functionCall1.arguments());
        assertThat(arguments1).hasSizeBetween(1, 2);
        assertThat(arguments1.get("location").toString()).contains("Boston");

        ToolCall toolCall2 = assistantMessage.toolCalls().get(1);
        assertThat(toolCall2.id()).isNotBlank();
        assertThat(toolCall2.type()).isEqualTo(FUNCTION);
        assertThat(toolCall2.function()).isNotNull();

        FunctionCall functionCall2 = toolCall2.function();
        assertThat(functionCall2.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall2.arguments()).isNotBlank();

        Map<String, Object> arguments2 = argumentsAsMap(functionCall2.arguments());
        assertThat(arguments2).hasSizeBetween(1, 2);
        assertThat(arguments2.get("location").toString()).contains("Madrid");

        // given
        String location1 = argument("location", functionCall1);
        String unit1 = argument("unit", functionCall1);
        String currentWeather1 = currentWeather(location1, unit1);
        ToolMessage toolMessage1 = ToolMessage.from(toolCall1.id(), currentWeather1);

        String location2 = argument("location", functionCall2);
        String unit2 = argument("unit", functionCall2);
        String currentWeather2 = currentWeather(location2, unit2);
        ToolMessage toolMessage2 = ToolMessage.from(toolCall2.id(), currentWeather2);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, toolMessage1, toolMessage2)
                .build();

        // when
        ChatCompletionResponse secondResponse = client.chatCompletion(secondRequest).execute();

        // then
        assertThat(secondResponse.content()).contains("11", "22");
    }

    private static String currentWeather(String location, String unit) {
        if (unit != null) {
            ChatCompletionTest.Unit.valueOf(unit.toUpperCase());
        }

        if (location.contains("Boston")) {
            return "{\"temperature\": 11,\"unit\":\"celsius\",\"description\":\"Foggy\"}";
        } else if (location.contains("Madrid")) {
            return "{\"temperature\": 22,\"unit\":\"celsius\",\"description\":\"Sunny\"}";
        } else {
            throw new IllegalArgumentException("Unexpected location: " + location);
        }
    }

    enum Unit {
        CELSIUS, FAHRENHEIT
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
            "GPT_4", "GPT_4_0613", // Does not support response_format
    })
    void testJsonResponseFormat(ChatCompletionModel model) throws Exception {

        // given
        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .addUserMessage("Extract information from the following text:\n" +
                        "Who is Klaus Heisler?\n" +
                        "Respond in JSON format with two fields: 'name' and 'surname'")
                .responseFormat(JSON_OBJECT)
                .build();

        CompletableFuture<ChatCompletionResponse> future = new CompletableFuture<>();

        // when
        client.chatCompletion(request)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        ChatCompletionResponse response = future.get(30, SECONDS);

        // then
        assertThat(response.content()).isEqualToIgnoringWhitespace("{\"name\":\"Klaus\",\"surname\":\"Heisler\"}");
    }

    @Test
    void testGpt4Vision() throws Exception {

        // given
        String imageUrl = "https://upload.wikimedia.org/wikipedia/commons/e/e9/Felis_silvestris_silvestris_small_gradual_decrease_of_quality.png";

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(GPT_4O)
                .messages(UserMessage.from("What is in this image?", imageUrl))
                .maxCompletionTokens(100)
                .build();

        CompletableFuture<ChatCompletionResponse> future = new CompletableFuture<>();

        // when
        client.chatCompletion(request)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        ChatCompletionResponse response = future.get(30, SECONDS);

        // then
        assertThat(response.content()).containsIgnoringCase("cat");
    }

    @Test
    void testGpt4Audio() throws Exception {

        // given
        URL resource = getClass().getClassLoader().getResource("sample.b64");
        final byte[] bytes = Files.readAllBytes(Paths.get(resource.toURI()));

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model("gpt-4o-audio-preview")
                .messages(UserMessage.builder()
                        .addText("What is on the audio?")
                        .addInputAudio(InputAudio.builder()
                                .format("wav")
                                .data(new String(bytes))
                                .build())
                        .build())
                .maxCompletionTokens(100)
                .temperature(0.0)
                .build();

        CompletableFuture<ChatCompletionResponse> future = new CompletableFuture<>();

        // when
        client.chatCompletion(request)
                .onResponse(future::complete)
                .onError(future::completeExceptionally)
                .execute();

        ChatCompletionResponse response = future.get(30, SECONDS);

        // then
        assertThat(response.content()).containsIgnoringCase("hello");
    }
}
