package tech.easyflow.ai.entity;

import com.easyagents.core.message.AiMessage;
import com.easyagents.core.message.Message;
import com.easyagents.core.message.SystemMessage;
import com.easyagents.core.message.ToolCall;
import com.easyagents.core.message.ToolMessage;
import com.easyagents.core.message.UserMessage;
import com.alibaba.fastjson2.JSONObject;
import com.fasterxml.jackson.annotation.JsonProperty;

import java.io.Serializable;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;

/**
 * Public API 聊天请求参数。
 */
public class ChatRequestParams implements Serializable {
    private static final long serialVersionUID = 1L;

    private BigInteger botId;
    private String conversationId;
    private List<Message> messages;

    public ChatRequestParams() {
    }

    @JsonProperty("messages")
    public void setMessagesFromJson(List<Object> rawMessages) {
        if (rawMessages == null) {
            this.messages = new ArrayList<>();
            return;
        }

        this.messages = new ArrayList<>();
        for (Object raw : rawMessages) {
            Message message = convertToMessage(raw);
            if (message != null) {
                this.messages.add(message);
            }
        }
    }

    private Message convertToMessage(Object raw) {
        JSONObject jsonObj = JSONObject.from(raw);
        if (jsonObj == null) {
            return null;
        }

        String role = jsonObj.getString("role");
        String content = jsonObj.getString("content");
        if (role == null || role.isBlank()) {
            throw new IllegalArgumentException("Invalid role: " + role);
        }
        if ("user".equals(role)) {
            if (content == null || content.isBlank()) {
                throw new IllegalArgumentException("User message cannot be empty");
            }
        }

        return switch (role) {
            case "user" -> jsonObj.toJavaObject(UserMessage.class);
            case "system" -> jsonObj.toJavaObject(SystemMessage.class);
            case "assistant" -> toAiMessage(jsonObj);
            case "tool" -> toToolMessage(jsonObj);
            default -> {
                UserMessage defaultMsg = new UserMessage();
                defaultMsg.setContent(content);
                yield defaultMsg;
            }
        };
    }

    /**
     * 把 JSON 结构恢复为 assistant 消息，兼容 `reasoningContent` / `reasoning_content`
     * 以及 `toolCalls` / `tool_calls` 两种写法。
     *
     * @param jsonObj 原始 JSON
     * @return assistant 消息
     */
    private AiMessage toAiMessage(JSONObject jsonObj) {
        AiMessage aiMessage = jsonObj.toJavaObject(AiMessage.class);
        String reasoningContent = jsonObj.getString("reasoningContent");
        if (reasoningContent == null || reasoningContent.isBlank()) {
            reasoningContent = jsonObj.getString("reasoning_content");
        }
        if (reasoningContent != null && !reasoningContent.isBlank()) {
            aiMessage.setReasoningContent(reasoningContent);
            aiMessage.setFullReasoningContent(reasoningContent);
        }
        List<ToolCall> toolCalls = parseToolCalls(jsonObj);
        if (!toolCalls.isEmpty()) {
            aiMessage.setToolCalls(toolCalls);
        }
        return aiMessage;
    }

    /**
     * 把 JSON 结构恢复为 tool 消息，兼容 `toolCallId` / `tool_call_id` 两种写法。
     *
     * @param jsonObj 原始 JSON
     * @return tool 消息
     */
    private ToolMessage toToolMessage(JSONObject jsonObj) {
        ToolMessage toolMessage = jsonObj.toJavaObject(ToolMessage.class);
        if (toolMessage.getToolCallId() == null || toolMessage.getToolCallId().isBlank()) {
            toolMessage.setToolCallId(jsonObj.getString("tool_call_id"));
        }
        return toolMessage;
    }

    /**
     * 解析 assistant 上的 tool calls。
     *
     * @param jsonObj 原始 JSON
     * @return tool call 列表
     */
    private List<ToolCall> parseToolCalls(JSONObject jsonObj) {
        Object rawToolCalls = jsonObj.get("toolCalls");
        if (rawToolCalls == null) {
            rawToolCalls = jsonObj.get("tool_calls");
        }
        List<ToolCall> toolCalls = new ArrayList<>();
        if (!(rawToolCalls instanceof List<?> rawList)) {
            return toolCalls;
        }
        for (Object rawToolCall : rawList) {
            JSONObject toolCallJson = JSONObject.from(rawToolCall);
            if (toolCallJson == null) {
                continue;
            }
            ToolCall toolCall = new ToolCall();
            toolCall.setId(toolCallJson.getString("id"));
            String toolName = toolCallJson.getString("name");
            if (toolName == null || toolName.isBlank()) {
                JSONObject functionJson = toolCallJson.getJSONObject("function");
                if (functionJson != null) {
                    toolName = functionJson.getString("name");
                    if (toolCallJson.getString("arguments") == null || toolCallJson.getString("arguments").isBlank()) {
                        toolCall.setArguments(functionJson.getString("arguments"));
                    }
                }
            }
            toolCall.setName(toolName);
            if (toolCall.getArguments() == null || toolCall.getArguments().isBlank()) {
                toolCall.setArguments(toolCallJson.getString("arguments"));
            }
            toolCalls.add(toolCall);
        }
        return toolCalls;
    }

    public List<Message> getMessages() {
        return messages;
    }

    public void setMessages(List<Message> messages) {
        this.messages = messages;
    }

    public BigInteger getBotId() {
        return botId;
    }

    public void setBotId(BigInteger botId) {
        this.botId = botId;
    }

    public String getConversationId() {
        return conversationId;
    }

    public void setConversationId(String conversationId) {
        this.conversationId = conversationId;
    }
}
