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

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.TypeReference;
import com.pavis.admin.aigc.model.resp.AgentPlanResp;
import com.pavis.admin.aigc.model.resp.AgentPlanStepResp;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.openai.api.OpenAiApi.FunctionTool;
import org.springframework.ai.tool.metadata.ToolMetadata;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 计划工具类
 */
@Slf4j
public class PlanningTool implements Function<String, ToolExecuteResult> {

    // 匹配字符串开头的方括号，支持中文和其他字符
    Pattern pattern = Pattern.compile("^\\s*\\[([^\\]]+)\\]");

    private AgentPlanResp currentPlan;

    public Long getCurrentPlanId() {
        return currentPlan != null ? currentPlan.getId() : null;
    }

    public AgentPlanResp getCurrentPlan() {
        return currentPlan;
    }

    private static final String PARAMETERS = """
        {
            "type": "object",
            "properties": {
                "command": {
                    "description": "create a execution plan , Available commands: create",
                    "enum": [
                        "create"
                    ],
                    "type": "string"
                },
                "plan_id": {
                    "description": "Unique identifier for the plan",
                    "type": "integer"
                },
                "title": {
                    "description": "Title for the plan",
                    "type": "string"
                },
                "steps": {
                    "description": "List of plan steps",
                    "type": "array",
                    "items": {
                        "type": "string"
                    }
                },
                "step_index": {
                    "description": "Index of step to update",
                    "type": "integer"
                },
                "step_status": {
                    "description": "Status to set for step",
                    "enum": ["not_started", "in_progress", "completed", "blocked"],
                    "type": "string"
                },
                "step_notes": {
                    "description": "Additional notes for step",
                    "type": "string"
                }
            },
            "required": ["command"]
        }
        """;

    private static final String name = "planning";

    private static final String description = "Planning tool for managing tasks ";

    public FunctionTool getToolDefinition() {
        return new FunctionTool(new FunctionTool.Function(description, name, PARAMETERS));
    }

    public FunctionToolCallback getFunctionToolCallback() {
        return FunctionToolCallback.builder(name, this)
            .description(description)
            .inputSchema(PARAMETERS)
            .inputType(String.class)
            .toolMetadata(ToolMetadata.builder().returnDirect(true).build())
            .build();
    }

    public ToolExecuteResult run(String toolInput) {
        try {
            Map<String, Object> input = JSON.parseObject(toolInput, new TypeReference<Map<String, Object>>() {
            });
            String command = (String)input.get("command");
            Long planId = (Long)input.get("plan_id");
            String title = (String)input.get("title");
            List<String> steps = JSON.parseObject(JSON.toJSONString(input
                .get("steps")), new TypeReference<List<String>>() {
                });

            return switch (command) {
                case "create" -> createPlan(planId, title, steps);
                // case "update" -> updatePlan(planId, title, steps);
                // case "get" -> getPlan(planId);
                // case "mark_step" -> markStep(planId, stepIndex, stepStatus, stepNotes);
                // case "delete" -> deletePlan(planId);
                default -> {
                    log.info("收到无效的命令: {}", command);
                    throw new IllegalArgumentException("Invalid command: " + command);
                }
            };
        } catch (Exception e) {
            log.info("执行计划工具时发生错误", e);
            return new ToolExecuteResult("Error executing planning tool: " + e.getMessage());
        }
    }

    public ToolExecuteResult createPlan(Long planId, String title, List<String> steps) {
        if (planId == null || title == null || steps == null || steps.isEmpty()) {
            log.info("创建计划时缺少必要参数: planId={}, title={}, steps={}", planId, title, steps);
            return new ToolExecuteResult("Required parameters missing");
        }
        AgentPlanResp plan = new AgentPlanResp();
        plan.setId(planId);
        plan.setTitle(title);
        List<AgentPlanStepResp> planSteps = new ArrayList<>();
        // 使用创建并添加步骤
        for (int i = 0; i < steps.size(); i++) {
            AgentPlanStepResp planStepResp = new AgentPlanStepResp();
            planStepResp.setPlanId(planId);
            planStepResp.setStepIndex(i);
            planStepResp.setDescription(steps.get(i));
            planStepResp.setAgentCode(getAgentCodeFromStepDesc(planStepResp.getDescription()));
            planSteps.add(planStepResp);
        }
        plan.setSteps(planSteps);
        this.currentPlan = plan;
        return new ToolExecuteResult("Plan created: " + planId + "\n" + plan.getPlanExecutionStateStringFormat(false));
    }

    @Override
    public ToolExecuteResult apply(String toolInput) {
        log.info("Planning tool raw input: {}", toolInput);
        return run(toolInput);
    }

    /**
     * 根据步骤描述获取 智能体编码
     *
     * @param stepDesc 步骤描述
     * @return 智能体编码
     */
    private String getAgentCodeFromStepDesc(String stepDesc) {
        Matcher matcher = pattern.matcher(stepDesc);
        if (matcher.find()) {
            // 对匹配到的内容进行trim和转小写处理
            return matcher.group(1).trim().toUpperCase();
        }
        return "DEFAULT_AGENT"; // Default agent if no match found
    }

}
