package com.pavis.admin.aigc.core.agent;

import com.pavis.admin.aigc.core.agent.tool.*;
import com.pavis.admin.aigc.core.llm.LlmService;
import com.pavis.admin.aigc.enums.AgentState;
import com.pavis.admin.aigc.model.entity.ToolDO;
import com.pavis.admin.aigc.model.resp.AgentResp;
import com.pavis.admin.aigc.service.ToolService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.metadata.ToolMetadata;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * 智能体实例，实现了ReAct（Reasoning + Acting）模式的智能体
 * 用于管理系统Agent的全生命周期
 */
@Slf4j
public class PavisAgent extends ReActAgent {

    /**
     * 本次计划id
     */
    private Long planId;

    /**
     * 智能体复合类
     * 包括智能体的基本信息和智能体可以使用的工具
     */
    private AgentResp agent;

    /**
     * 用户提示词
     */
    private Prompt userPrompt;

    /**
     * 大模型相应
     */
    private ChatResponse response;

    /**
     * 推理-执行记录器
     */
    private ReActRecord reActRecord;

    /**
     * 用于获取所有工具集
     */
    private ToolService toolService;;

    /**
     * 构造函数
     *
     * @param llmService LLM服务实例，用于处理自然语言交互
     */
    public PavisAgent(LlmService llmService, Long planId, AgentResp agent, ToolService toolService) {
        super(llmService);
        this.planId = planId;
        this.agent = agent;
        this.toolService = toolService;
    }

    @Override
    protected boolean think() {
        reActRecord = new ReActRecord();
        try {
            List<Message> messages = new ArrayList<>();
            addThinkPrompt(messages);

            ChatOptions chatOptions = ToolCallingChatOptions.builder().internalToolExecutionEnabled(false).build();
            Message nextStepMessage = getNextStepWithEnvMessage();
            messages.add(nextStepMessage);
            reActRecord.startThinking(messages.toString());

            log.debug("Messages prepared for the prompt: {}", messages);

            userPrompt = new Prompt(messages, chatOptions);

            response = llmService.getAgentChatClient(String.valueOf(planId))
                .getChatClient()
                .prompt(userPrompt)
                .advisors(memoryAdvisor -> memoryAdvisor.param(CHAT_MEMORY_CONVERSATION_ID_KEY, String.valueOf(planId))
                    .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .toolCallbacks(getToolCallList())
                .call()
                .chatResponse();

            String responseByLLm = response.getResult().getOutput().getText();
            if (responseByLLm != null && !responseByLLm.isEmpty()) {
                log.info(String.format("💬 %s's response: %s", getName(), responseByLLm));
            }

            List<AssistantMessage.ToolCall> toolCalls = response.getResult().getOutput().getToolCalls();

            log.info(String.format("✨ %s's thoughts: %s", getName(), responseByLLm));
            log.info(String.format("🛠️ %s selected %d tools to use", getName(), toolCalls.size()));

            if (!toolCalls.isEmpty()) {
                log.info(String.format("🧰 Tools being prepared: %s", toolCalls.stream()
                    .map(AssistantMessage.ToolCall::name)
                    .collect(Collectors.toList())));
                reActRecord.setExecutionNeeded(true);
                reActRecord.setToolName(toolCalls.get(0).name());
                reActRecord.setToolParameters(toolCalls.get(0).arguments());
                responseByLLm = String.format("Tools being prepared. Executing tool: tool %s,arguments %s.", toolCalls
                    .get(0)
                    .name(), toolCalls.get(0).arguments());
            }
            reActRecord.setStatus("SUCCESS");
            reActRecord.finishThinking(responseByLLm);
            return !toolCalls.isEmpty();
        } catch (Exception e) {
            log.error(String.format("🚨 Oops! The %s's thinking process hit a snag: %s", getName(), e.getMessage()));
            reActRecord.recordError(e.getMessage());
            return false;
        }
    }

    @Override
    protected AgentExecResult act() {
        // 这里是返回run方法的返回值，会被获取result
        AgentExecResult result = new AgentExecResult(reActRecord.getThinkOutput(), AgentState.COMPLETED, reActRecord);
        log.info("Agent Think Result: {}", result.getResult());
        return result;
    }

    @Override
    public String getName() {
        return this.agent.getName();
    }

    @Override
    public String getDescription() {
        return this.agent.getDescription();
    }

    @Override
    public Integer getMaxStepCount() {
        return this.agent.getMaxSteps();
    }

    @Override
    public boolean isStucked() {
        return false;
    }

    @Override
    public void handleStuck() {

    }

    @Override
    public List<ToolCallback> getToolCallList() {
        List<ToolCallback> toolCallbacks = new ArrayList<>();
        Map<String, ToolCallBackContext> toolCallbackMap = new HashMap<>();
        List<ToolCallBiFunctionDef> toolDefinitions = new ArrayList<>();
        @SuppressWarnings("unchecked") Map<String, Object> envData = getExecutionParam().getEnvData().getEnvDataMap();
        // 得先获取所有工具
        List<ToolDO> tools = toolService.getAll();
        if (tools != null) {
            for (ToolDO tool : tools) {
                toolDefinitions.add(new BrowserUseTool(planId, tool, envData));
            }
        }
        // 添加终止工具定义
        toolDefinitions.add(new TerminateTool(planId));

        // 为每个工具创建 FunctionToolCallback
        for (ToolCallBiFunctionDef toolDefinition : toolDefinitions) {
            FunctionToolCallback functionToolcallback = FunctionToolCallback.builder(toolDefinition
                .getName(), toolDefinition)
                .description(toolDefinition.getDescription())
                .inputSchema(toolDefinition.getParameters())
                .inputType(toolDefinition.getInputType())
                .toolMetadata(ToolMetadata.builder().returnDirect(toolDefinition.isReturnDirect()).build())
                .build();
            toolDefinition.setPlanId(planId);
            ToolCallBackContext functionToolcallbackContext = new ToolCallBackContext(functionToolcallback, toolDefinition);
            toolCallbackMap.put(toolDefinition.getName(), functionToolcallbackContext);
        }

        for (String toolKey : getAvailableToolKeys()) {
            if (toolCallbackMap.containsKey(toolKey)) {
                ToolCallBackContext toolCallback = toolCallbackMap.get(toolKey);
                if (toolCallback != null) {
                    toolCallbacks.add(toolCallback.getToolCallback());
                }
            } else {
                log.warn("Tool callback for {} not found in the map.", toolKey);
            }
        }
        return toolCallbacks;
    }

    @Override
    protected Message getNextStepWithEnvMessage() {
        String nextStepPrompt = """

            CURRENT STEP ENVIRONMENT STATUS:
            {current_step_env_data}

            """;
        nextStepPrompt += this.agent.getNextStepPrompt();
        PromptTemplate promptTemplate = new PromptTemplate(nextStepPrompt);
        return promptTemplate.createMessage(getExecutionParam().getExecutionParamMap());
    }

    @Override
    protected Message addThinkPrompt(List<Message> messages) {
        super.addThinkPrompt(messages);
        SystemPromptTemplate promptTemplate = new SystemPromptTemplate(this.agent.getSystemPrompt());
        Message systemMessage = promptTemplate.createMessage(getExecutionParam().getExecutionParamMap());
        messages.add(systemMessage);
        return systemMessage;
    }

    /**
     * 获取当前智能体的可用工具key值
     *
     * @return 智能体可用工具的列表
     */
    public List<String> getAvailableToolKeys() {
        List<ToolDO> agentToolList = this.agent.getTools();
        if (agentToolList == null || agentToolList.isEmpty()) {
            return null;
        }
        return agentToolList.stream().map(ToolDO::getToolKey).collect(Collectors.toList());
    }
}
