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

import com.pavis.admin.aigc.core.agent.browser.EnvData;
import com.pavis.admin.aigc.core.agent.browser.ExecutionParam;
import com.pavis.admin.aigc.core.llm.LlmService;
import com.pavis.admin.aigc.enums.AgentState;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.tool.ToolCallback;

import java.time.LocalDateTime;
import java.util.List;

/**
 * An abstract base class for implementing AI agents that can execute multi-step tasks.
 * This class provides the core functionality for managing agent state, conversation flow,
 * and step-by-step execution of tasks.
 *
 * <p>
 * The agent supports a finite number of execution steps and includes mechanisms for:
 * <ul>
 * <li>State management (idle, running, finished)</li>
 * <li>Conversation tracking</li>
 * <li>Step limitation and monitoring</li>
 * <li>Thread-safe execution</li>
 * <li>Stuck-state detection and handling</li>
 * </ul>
 */
@Slf4j
public abstract class AbstractBaseAgent {

    @Getter
    @Setter
    private AgentState state = AgentState.NOT_STARTED;

    protected LlmService llmService;

    private ExecutionParam executionParam;

    /**
     * 智能体开始运行时间
     */
    @Getter
    private LocalDateTime startRunTime;

    /**
     * 智能体结束运行时间
     */
    @Getter
    private LocalDateTime endRunTime;

    public AbstractBaseAgent(LlmService llmService) {
        this.llmService = llmService;
    }

    /**
     * 获取智能体的名称
     * <p>
     * 实现要求：
     * 1. 返回一个简短但具有描述性的名称
     * 2. 名称应该反映该智能体的主要功能或特性
     * 3. 名称应该是唯一的，便于日志和调试
     * <p>
     * 示例实现：
     * - ToolCallAgent 返回 "ToolCallAgent"
     * - BrowserAgent 返回 "BrowserAgent"
     *
     * @return 智能体的名称
     */
    public abstract String getName();

    /**
     * 获取智能体的详细描述
     * <p>
     * 实现要求：
     * 1. 返回对该智能体功能的详细描述
     * 2. 描述应包含智能体的主要职责和能力
     * 3. 应说明该智能体与其他智能体的区别
     * <p>
     * 示例实现：
     * - ToolCallAgent: "负责管理和执行工具调用的智能体，支持多工具组合调用"
     * - ReActAgent: "实现思考(Reasoning)和行动(Acting)交替执行的智能体"
     *
     * @return 智能体的详细描述文本
     */
    public abstract String getDescription();

    /**
     * 获取智能体允许运行的最大步骤数
     *
     * @return 智能体允许运行的最大步骤数
     */
    public abstract Integer getMaxStepCount();

    /**
     * 检查是否处于卡住状态
     */
    public abstract boolean isStucked();

    /**
     * 处理卡住的问题
     */
    public abstract void handleStuck();

    /**
     * 获取智能体的工具清单
     *
     * @return 智能体的工具清单
     */
    public abstract List<ToolCallback> getToolCallList();

    /**
     * 智能体下一步执行提示
     *
     * @return 下一步执行结果
     */
    protected abstract AgentExecResult step();

    /**
     * 获取下一步操作的提示消息
     * <p>
     * 实现要求：
     * 1. 生成引导智能体执行下一步操作的提示消息
     * 2. 提示内容应该基于当前执行状态和上下文
     * 3. 消息应该清晰指导智能体要执行什么任务
     * <p>
     * 子类实现参考：
     * 1. ToolCallAgent：返回工具选择和调用相关的提示
     * 2. ReActAgent：返回思考或行动决策相关的提示
     *
     * @return 下一步操作的提示消息对象
     */
    protected abstract Message getNextStepWithEnvMessage();

    /**
     * 添加思考提示到消息列表中，构建智能体的思考链
     * <p>
     * 实现要求：
     * 1. 根据当前上下文和状态生成合适的系统提示词
     * 2. 提示词应该指导智能体如何思考和决策
     * 3. 可以递归地构建提示链，形成层次化的思考过程
     * 4. 返回添加的系统提示消息对象
     * <p>
     * 子类实现参考：
     * 1. ReActAgent: 实现基础的思考-行动循环提示
     * 2. ToolCallAgent: 添加工具选择和执行相关的提示
     *
     * @param messages 当前的消息列表，用于构建上下文
     * @return 添加的系统提示消息对象
     */
    protected Message addThinkPrompt(List<Message> messages) {
        // 获取操作系统信息
        String osName = "";
        String osVersion = "";
        String osArch = "";

        EnvData envData = getExecutionParam().getEnvData();
        if (envData != null) {
            osName = envData.getOsName();
            osVersion = envData.getOsVersion();
            osArch = envData.getOsArch();
        }

        // 获取当前日期时间，格式为yyyy-MM-dd
        String currentDateTime = java.time.LocalDate.now().toString(); // 格式为yyyy-MM-dd

        String stepPrompt = """
            - SYSTEM INFORMATION:
            OS: %s %s (%s)

            - Current Date:
            %s
            - 全局计划信息:
            {planStatus}

            - 当前要做的步骤要求:
            STEP {currentStepIndex} : {stepText}

            - 当前步骤的上下文信息:
            {extraParams}

            重要说明：
            1. 使用工具调用时，不需要额外的任何解释说明！
            2. 不要在工具调用前提供推理或描述！
            3. 专注于立即行动而非解释！

            """.formatted(osName, osVersion, osArch, currentDateTime);

        SystemPromptTemplate promptTemplate = new SystemPromptTemplate(stepPrompt);

        Message systemMessage = promptTemplate.createMessage(getExecutionParam().getExecutionParamMap());

        messages.add(systemMessage);
        return systemMessage;
    }

    /**
     * 控制智能体运行的整体流程
     * 主要用于监控智能体失败次数是否超出失败上限
     *
     * @param executionParam 执行智能体的数据
     * @return 智能体此步骤的执行结果
     */
    public AgentExecResult run(ExecutionParam executionParam) {
        startRunTime = LocalDateTime.now();
        // 智能体运行的步数数量计数器
        int currentStepCount = 0;
        if (state != AgentState.IN_PROGRESS) {
            throw new IllegalStateException("Cannot run agent from state: " + state);
        }
        // 设置智能体运行数据
        setExecutionParam(executionParam);

        String resultStr = "";
        AgentExecResult stepResult = null;
        try {
            state = AgentState.IN_PROGRESS;
            while (currentStepCount < getMaxStepCount() && !state.equals(AgentState.COMPLETED)) {
                currentStepCount++;
                log.info("Executing round {}/{}", currentStepCount, getMaxStepCount());

                // 执行具体步骤内容
                stepResult = step();

                if (isStucked()) {
                    log.warn("Agent stuck detected - Missing tool calls");
                    handleStuck();
                } else {
                    // 更新全局状态以保持一致性
                    log.info("Agent state: {}", stepResult.getState());
                    state = stepResult.getState();
                }

                resultStr = String.format("Round %s: %s", currentStepCount, stepResult.getResult());

            }

            if (currentStepCount >= getMaxStepCount()) {
                resultStr = String.format("Terminated: Reached max rounds ( %d )", getMaxStepCount());
            }

            endRunTime = LocalDateTime.now();
        } catch (Exception e) {
            endRunTime = LocalDateTime.now();
            state = AgentState.FAILED;
            log.error("Agent execution failed", e);
            throw e; // 重新抛出异常，让上层调用者知道发生了错误
        } finally {
            state = AgentState.COMPLETED; // Reset state after execution
        }
        if (stepResult == null) {
            stepResult = new AgentExecResult(resultStr, state);
        }
        return stepResult;
    }

    /**
     * 获取智能体的数据上下文
     * <p>
     * 使用说明：
     * 1. 返回智能体在执行过程中需要的所有上下文数据
     * 2. 数据可包含： - 当前执行状态 - 步骤信息 - 中间结果 - 配置参数
     * 3. 数据在run()方法执行时通过setExecutionParam()设置
     * <p>
     * 不要修改这个方法的实现，如果你需要传递上下文，继承并修改setExecutionParam方法，这样可以提高setExecutionParam()的的效率。
     *
     * @return 包含智能体上下文数据的ExecutionParam对象
     */
    protected final ExecutionParam getExecutionParam() {
        return executionParam;
    }

    protected void setExecutionParam(ExecutionParam executionParam) {
        this.executionParam = executionParam;
    }

    public static class AgentExecResult {

        private String result;

        private AgentState state;

        private ReActRecord reActRecord;

        public AgentExecResult(String result, AgentState state) {
            this.result = result;
            this.state = state;
            this.reActRecord = null;
        }

        public AgentExecResult(String result, AgentState state, ReActRecord reActRecord) {
            this.result = result;
            this.state = state;
            this.reActRecord = reActRecord;
        }

        public String getResult() {
            return result;
        }

        public void setResult(String result) {
            this.result = result;
        }

        public AgentState getState() {
            return state;
        }

        public void setState(AgentState state) {
            this.state = state;
        }

        public ReActRecord getReActRecord() {
            return reActRecord;
        }

    }

}
