package com.pavis.admin.aigc.service.impl;

import cn.hutool.core.bean.BeanUtil;
import com.alibaba.fastjson2.JSON;
import com.pavis.admin.aigc.core.agent.AbstractBaseAgent;
import com.pavis.admin.aigc.core.agent.PavisAgent;
import com.pavis.admin.aigc.core.agent.browser.ExecutionParam;
import com.pavis.admin.aigc.core.agent.planner.PlanContext;
import com.pavis.admin.aigc.core.agent.tool.PlanningTool;
import com.pavis.admin.aigc.core.llm.LlmService;
import com.pavis.admin.aigc.enums.AgentState;
import com.pavis.admin.aigc.mapper.AgentPlanMapper;
import com.pavis.admin.aigc.model.entity.AgentPlanDO;
import com.pavis.admin.aigc.model.query.AgentPlanQuery;
import com.pavis.admin.aigc.model.query.AgentPlanStepQuery;
import com.pavis.admin.aigc.model.req.AgentPlanReq;
import com.pavis.admin.aigc.model.resp.AgentPlanDetailResp;
import com.pavis.admin.aigc.model.resp.AgentPlanResp;
import com.pavis.admin.aigc.model.resp.AgentPlanStepResp;
import com.pavis.admin.aigc.model.resp.AgentResp;
import com.pavis.admin.aigc.service.AgentPlanService;
import com.pavis.admin.aigc.service.AgentPlanStepService;
import com.pavis.admin.aigc.service.AgentService;
import com.pavis.admin.aigc.service.ToolService;
import com.pavis.admin.common.context.UserContextHolder;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.stereotype.Service;
import top.continew.starter.extension.crud.model.query.SortQuery;
import top.continew.starter.extension.crud.service.BaseServiceImpl;

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

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;

/**
 * 智能体规划记录业务实现
 *
 * @author semi
 * @since 2025/05/28 16:59
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class AgentPlanServiceImpl extends BaseServiceImpl<AgentPlanMapper, AgentPlanDO, AgentPlanResp, AgentPlanDetailResp, AgentPlanQuery, AgentPlanReq> implements AgentPlanService {

    private final LlmService llmService;

    private final AgentService agentService;

    private final ToolService toolService;

    private final AgentPlanStepService agentPlanStepService;

    @Override
    public PlanContext createPlan(AgentPlanReq req) {
        // 获取userId
        Long userId = UserContextHolder.getUserId();
        req.setUserId(userId);
        // 先做持久化获取planId
        Long planId = super.create(req);
        if (planId == null) {
            throw new IllegalArgumentException("Plan ID cannot be null or empty");
        }
        // 创建规划上下文
        PlanContext context = new PlanContext();
        context.setPlanId(planId);
        context.setConversationId(req.getConversationId());
        context.setMessageId(req.getMessageId());
        context.setMessageContent(req.getMessageContent());
        AgentPlanResp currentPlan = this.getByPlanId(planId);
        try {
            // 构建所有智能体信息
            String agentsInfo = buildAgentsInfo();
            // 生成计划提示
            String planPrompt = generatePlanPrompt(context.getMessageContent(), agentsInfo, planId);
            currentPlan.setThinkInput(planPrompt);
            // 使用 LLM 生成计划
            PromptTemplate promptTemplate = new PromptTemplate(planPrompt);
            Prompt prompt = promptTemplate.create();

            // 创建 规划工具 PlanningTool
            PlanningTool planningTool = new PlanningTool();
            currentPlan.setStartTime(LocalDateTime.now());
            ChatClient.CallResponseSpec response = llmService.getPlanningChatClient()
                .prompt(prompt)
                .toolCallbacks(List.of(planningTool.getFunctionToolCallback()))
                .advisors(memoryAdvisor -> memoryAdvisor.param(CHAT_MEMORY_CONVERSATION_ID_KEY, context
                    .getConversationId()).param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .call();
            String outputText = response.chatResponse().getResult().getOutput().getText();
            // 检查计划是否创建成功
            if (planId.equals(planningTool.getCurrentPlanId())) {
                currentPlan.setThinkOutput(outputText);
                currentPlan.setEndTime(LocalDateTime.now());
                currentPlan.setDuration(Duration.between(currentPlan.getStartTime(), currentPlan.getEndTime())
                    .getSeconds());
                currentPlan.setTitle(planningTool.getCurrentPlan().getTitle());
                currentPlan.setSteps(planningTool.getCurrentPlan().getSteps());
            }
            // 更新或存储计划信息
            AgentPlanResp savedPlan = this.saveOrUpdate(currentPlan);
            context.setPlan(savedPlan);
        } catch (Exception e) {
            log.error("Error creating plan for request: {}", context.getMessageContent(), e);
            // 处理异常情况
            throw new RuntimeException("Failed to create plan", e);
        }
        return context;
    }

    /**
     * 当前需要执行的单个计划步骤
     * 这一步需要将概括的计划步骤生成可供前端直接执行的步骤
     * 1、 先根据模型指定的智能体编码获取执行该步骤计划的智能体信息和智能体工具
     *
     * @param context 计划上下文
     */
    @Override
    public PlanContext executePlan(PlanContext context) {
        // 先校验上一步执行结果 TODO
        // 每次执行步骤时需要前端传入当前浏览器环境，包括插件版本、浏览器厂商及版本、当前打开的tab页列表【包含已经编号的dom tree】
        // 获取数据库存储的最新规划信息
        AgentPlanResp currentPlan = this.getByPlanId(context.getPlanId());
        if (currentPlan == null) {
            throw new IllegalArgumentException("Plan ID cannot be null or empty");
        }
        // 获取计划步骤数据
        currentPlan.setSteps(agentPlanStepService.getByPlanId(context.getPlanId()));
        // 设置当前步骤ID
        currentPlan.setCurrentPlanStepId(context.getCurrentPlanStepId());
        // 重新设置当前计划信息
        context.setPlan(currentPlan);
        // 获取当前执行的步骤
        AgentPlanStepResp currentPlanStep = currentPlan.getSteps()
            .stream()
            .filter(resp -> resp.getId().equals(context.getCurrentPlanStepId()))
            .findFirst()
            .orElse(null);
        if (currentPlanStep == null) {
            throw new IllegalArgumentException("Current step ID cannot be null or empty");
        }
        // 开始按照计划步骤生成具体执行内容
        log.info("Current Plan step: {}", currentPlanStep);
        // 根据计划时大模型指定的智能体编码获取具体的智能体信息和智能体工具
        AgentResp agent = agentService.getByCode(currentPlanStep.getAgentCode());
        if (agent == null) {
            // 如果没找到可以使用的智能体则直接返回，计划执行失败
            log.error("No agent found for agent code: {}", currentPlanStep.getAgentCode());
            currentPlanStep.setResult("No agent found for step type: " + currentPlanStep.getAgentCode());
            throw new IllegalArgumentException("No agent found for agent code:" + currentPlanStep.getAgentCode());
        }
        agent.setTools(toolService.selectByAgentId(agent.getId()));
        // 创建派维斯智能体
        PavisAgent pavisAgent = new PavisAgent(llmService, context.getPlanId(), agent, toolService);
        // 获取当前执行的计划步骤的序号
        int currentPlanStepIndex = currentPlanStep.getStepIndex();
        // 设置执行当前计划步骤的智能体信息
        pavisAgent.setState(AgentState.IN_PROGRESS);
        currentPlanStep.setAgent(pavisAgent);
        try {
            // 当前计划的整体情况
            String currentPlanStatus = context.getPlan().getPlanExecutionStateStringFormat(true);
            // 当前需要执行的计划步骤的描述信息
            String currentPlanStepDescription = currentPlanStep.getDescription();

            ExecutionParam executionParam = new ExecutionParam();
            executionParam.setPlanStatus(currentPlanStatus);
            executionParam.setCurrentStepIndex(currentPlanStepIndex);
            executionParam.setStepText(currentPlanStepDescription);
            executionParam.setExtraParams(currentPlanStep.getExecutionParams() == null
                ? ""
                : currentPlanStep.getExecutionParams());
            executionParam.setEnvData(context.getEnvData());
            // 需要获取前端插件的环境状态信息
            executionParam.setEnvDataString(executionParam.getCurrentEnvString());

            log.info("executionParams:{}", JSON.toJSONString(executionParam));

            AbstractBaseAgent.AgentExecResult agentExecResult = pavisAgent.run(executionParam);
            // Execute the step
            currentPlanStep.setAgentExecResult(agentExecResult);

        } catch (Exception e) {
            log.error("Error executing step: {}", e.getMessage(), e);
            currentPlanStep.setResult("Execution failed: " + e.getMessage());
        } finally {
            log.info("current plan:{}", JSON.toJSONString(currentPlan));
            this.saveOrUpdate(currentPlan);
        }
        // 判断任务是否完成 TODO
        return context;
    }

    @Override
    public AgentPlanResp getByPlanId(Long planId) {
        return baseMapper.lambdaQuery()
            .eq(AgentPlanDO::getId, planId)
            .oneOpt()
            .map(agentPlanDO -> BeanUtil.copyProperties(agentPlanDO, AgentPlanResp.class))
            .orElse(null);
    }

    @Override
    public AgentPlanResp saveOrUpdate(AgentPlanResp resp) {
        AgentPlanDO agentPlanDO = baseMapper.selectById(resp.getId());
        BeanUtil.copyProperties(resp, agentPlanDO);
        baseMapper.updateById(agentPlanDO);
        // 保存计划步骤
        boolean result = agentPlanStepService.saveOrUpdateBatch(resp.getSteps());
        if (result) {
            AgentPlanStepQuery query = new AgentPlanStepQuery();
            query.setPlanId(resp.getId());
            SortQuery sortQuery = new SortQuery();
            List<AgentPlanStepResp> list = agentPlanStepService.list(query, sortQuery);
            resp.setSteps(list);
        }
        return resp;
    }

    /**
     * 构建智能体信息，用于生成提示词
     *
     * @return 智能体信息描述
     */
    public String buildAgentsInfo() {
        List<AgentResp> agents = agentService.getAll();
        StringBuilder agentsInfo = new StringBuilder("Available Agents:\n");
        if (agents != null && !agents.isEmpty()) {
            for (AgentResp agent : agents) {
                agentsInfo.append("- Agent Name: ")
                    .append(agent.getName())
                    .append("\n  Description: ")
                    .append(agent.getDescription())
                    .append("\n");
            }
        }
        return agentsInfo.toString();
    }

    /**
     * 生成计划提示
     *
     * @param userMessage 用户请求
     * @param agentsInfo  代理信息
     * @param planId      计划ID
     * @return 格式化的提示字符串
     */
    private String generatePlanPrompt(String userMessage, String agentsInfo, Long planId) {
        return """
            ## 介绍
            我是一个AI助手，旨在帮助用户完成各种任务。我的设计目标是提供帮助、信息和多方面的支持。

            ## 目标
            我的主要目标是通过提供信息、执行任务和提供指导来帮助用户实现他们的目标。我致力于成为问题解决和任务完成的可靠伙伴。

            ## 我的任务处理方法
            当面对任务时，我通常会：
            1. 分析请求以理解需求
            2. 将复杂问题分解为可管理的步骤
            3. 为每个步骤使用适当的AGENT
            4. 以有帮助和有组织的方式交付结果

            ## 当前主要目标：
            创建一个合理的计划，包含清晰的步骤来完成任务。

            ## 可用代理信息：
            %s

            # 需要完成的任务：
            %s

            你可以使用规划工具来帮助创建计划，使用 %s 作为计划ID。

            重要提示：计划中的每个步骤都必须以[AGENT]开头，代理名称必须是上述列出的可用代理之一。
            例如："[BROWSER_AGENT] 搜索相关信息" 或 "[DEFAULT_AGENT] 处理搜索结果"
            """.formatted(agentsInfo, userMessage, planId);
    }
}