package tech.easyflow.ai.invoke.mapper;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.Assert;
import org.junit.Test;
import tech.easyflow.ai.invoke.exception.ModelInvokeException;
import tech.easyflow.ai.invoke.model.UnifiedChatRequest;
import tech.easyflow.ai.invoke.model.UnifiedChatResponse;
import tech.easyflow.ai.invoke.model.UnifiedChoice;
import tech.easyflow.ai.invoke.model.UnifiedContentPart;
import tech.easyflow.ai.invoke.model.UnifiedImageUrl;
import tech.easyflow.ai.invoke.model.UnifiedMessage;
import tech.easyflow.ai.invoke.model.UnifiedResponseFormat;
import tech.easyflow.ai.invoke.model.UnifiedTool;
import tech.easyflow.ai.invoke.model.UnifiedToolCall;
import tech.easyflow.ai.invoke.model.UnifiedToolCallFunction;
import tech.easyflow.ai.invoke.model.UnifiedToolFunction;
import tech.easyflow.ai.invoke.model.UnifiedUsage;
import tech.easyflow.ai.invoke.protocol.openai.OpenAiChatCompletionRequest;
import tech.easyflow.ai.invoke.protocol.openai.OpenAiChatCompletionResponse;

import java.util.Collections;
import java.util.List;

public class OpenAiProtocolMapperTest {

    private final OpenAiProtocolMapper mapper = new OpenAiProtocolMapper(new ObjectMapper());

    @Test
    public void shouldParseTextAndImageRequest() {
        String rawBody = """
                {
                  "model": "gpt-4-1-prod",
                  "messages": [
                    {
                      "role": "user",
                      "content": [
                        {
                          "type": "text",
                          "text": "帮我看下这张图"
                        },
                        {
                          "type": "image_url",
                          "image_url": {
                            "url": "data:image/png;base64,AAAA",
                            "detail": "high"
                          }
                        }
                      ]
                    }
                  ],
                  "stream": false,
                  "temperature": 0.2,
                  "top_p": 0.8,
                  "max_tokens": 512,
                  "seed": 7,
                  "tools": [
                    {
                      "type": "function",
                      "function": {
                        "name": "query_weather",
                        "description": "query weather",
                        "parameters": {
                          "type": "object"
                        }
                      }
                    }
                  ],
                  "tool_choice": "auto",
                  "response_format": {
                    "type": "json_schema",
                    "json_schema": {
                      "name": "weather_schema"
                    }
                  }
                }
                """;

        OpenAiChatCompletionRequest request = mapper.readRequest(rawBody);
        UnifiedChatRequest unifiedRequest = mapper.toUnifiedRequest(request);

        Assert.assertEquals("gpt-4-1-prod", unifiedRequest.getModel());
        Assert.assertEquals(Long.valueOf(7), unifiedRequest.getSeed());
        Assert.assertEquals(Integer.valueOf(512), unifiedRequest.getMaxTokens());
        Assert.assertNotNull(unifiedRequest.getTools());
        Assert.assertEquals(1, unifiedRequest.getTools().size());
        Assert.assertEquals("query_weather", unifiedRequest.getTools().get(0).getFunction().getName());
        Assert.assertEquals("json_schema", unifiedRequest.getResponseFormat().getType());
        Assert.assertEquals("weather_schema", unifiedRequest.getResponseFormat().getJsonSchema().get("name").asText());

        UnifiedMessage message = unifiedRequest.getMessages().get(0);
        Assert.assertEquals("user", message.getRole());
        Assert.assertNotNull(message.getContentParts());
        Assert.assertEquals(2, message.getContentParts().size());
        UnifiedContentPart imagePart = message.getContentParts().get(1);
        Assert.assertEquals("image_url", imagePart.getType());
        Assert.assertEquals("data:image/png;base64,AAAA", imagePart.getImageUrl().getUrl());
        Assert.assertEquals("high", imagePart.getImageUrl().getDetail());
    }

    @Test
    public void shouldRejectUnsupportedRootField() {
        String rawBody = """
                {
                  "model": "gpt-4-1-prod",
                  "messages": [
                    {
                      "role": "user",
                      "content": "hello"
                    }
                  ],
                  "n": 2
                }
                """;

        ModelInvokeException exception = Assert.assertThrows(
                ModelInvokeException.class,
                () -> mapper.readRequest(rawBody)
        );

        Assert.assertEquals(400, exception.getStatus());
        Assert.assertEquals("unsupported_field", exception.getCode());
    }

    @Test
    public void shouldAllowMissingOptionalFields() {
        String rawBody = """
                {
                  "model": "deepseek-chat",
                  "messages": [
                    {
                      "role": "user",
                      "content": "你好，介绍一下你自己。"
                    }
                  ]
                }
                """;

        OpenAiChatCompletionRequest request = mapper.readRequest(rawBody);
        UnifiedChatRequest unifiedRequest = mapper.toUnifiedRequest(request);

        Assert.assertEquals("deepseek-chat", unifiedRequest.getModel());
        Assert.assertNotNull(unifiedRequest.getMessages());
        Assert.assertEquals(1, unifiedRequest.getMessages().size());
        Assert.assertNull(unifiedRequest.getTools());
        Assert.assertNull(unifiedRequest.getToolChoice());
        Assert.assertNull(unifiedRequest.getResponseFormat());
    }

    @Test
    public void shouldMapToolCallsAndUsageInResponse() {
        UnifiedChatResponse response = new UnifiedChatResponse();
        response.setId("chatcmpl-1");
        response.setObject("chat.completion");
        response.setCreated(123L);
        response.setModel("gpt-4-1-prod");

        UnifiedToolCallFunction toolCallFunction = new UnifiedToolCallFunction();
        toolCallFunction.setName("query_weather");
        toolCallFunction.setArguments("{\"city\":\"shanghai\"}");

        UnifiedToolCall toolCall = new UnifiedToolCall();
        toolCall.setId("call_1");
        toolCall.setType("function");
        toolCall.setFunction(toolCallFunction);

        UnifiedMessage message = new UnifiedMessage();
        message.setRole("assistant");
        message.setContent(null);
        message.setToolCalls(Collections.singletonList(toolCall));

        UnifiedChoice choice = new UnifiedChoice();
        choice.setIndex(0);
        choice.setMessage(message);
        choice.setFinishReason("tool_calls");

        UnifiedUsage usage = new UnifiedUsage();
        usage.setPromptTokens(12);
        usage.setCompletionTokens(34);
        usage.setTotalTokens(46);

        response.setChoices(List.of(choice));
        response.setUsage(usage);

        OpenAiChatCompletionResponse openAiResponse = mapper.toOpenAiResponse(response);

        Assert.assertEquals("chatcmpl-1", openAiResponse.getId());
        Assert.assertEquals("tool_calls", openAiResponse.getChoices().get(0).getFinishReason());
        Assert.assertEquals("query_weather", openAiResponse.getChoices().get(0).getMessage().getToolCalls().get(0).getFunction().getName());
        Assert.assertEquals(Integer.valueOf(46), openAiResponse.getUsage().getTotalTokens());
    }
}
