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

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.function.FunctionToolCallback;
import org.springframework.ai.tool.metadata.ToolMetadata;
import org.springframework.core.ParameterizedTypeReference;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

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

    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() {
        // 创建保留泛型信息的类型引用
        ParameterizedTypeReference<Map<String, Object>> typeRef =
                new ParameterizedTypeReference<Map<String, Object>>() {
                };

        return FunctionToolCallback.builder(name, this)
                .description(description)
                .inputSchema(PARAMETERS)
                .inputType(typeRef) // 使用 ParameterizedTypeReference
                .toolMetadata(ToolMetadata.builder().returnDirect(true).build())
                .build();
    }

    public ToolExecuteResult run(Map<String, Object> toolInput) {
        try {
            log.info("Planning tool parsed input: {}", toolInput);
            String command = (String) toolInput.get("command");

            // 安全处理 plan_id
            Long planId = null;
            Object planIdObj = toolInput.get("plan_id");
            if (planIdObj != null) {
                if (planIdObj instanceof Number) {
                    planId = ((Number) planIdObj).longValue();
                } else {
                    try {
                        planId = Long.parseLong(planIdObj.toString());
                    } catch (NumberFormatException e) {
                        log.warn("Invalid plan_id format: {}", planIdObj);
                    }
                }
            }

            String title = (String) toolInput.get("title");

            // 安全处理 steps
            List<String> steps = new ArrayList<>();
            Object stepsObj = toolInput.get("steps");
            if (stepsObj instanceof List) {
                for (Object item : (List<?>) stepsObj) {
                    if (item != null) {
                        steps.add(item.toString());
                    }
                }
            }

            log.info("Executing command: {}, planId: {}, title: {}, steps: {}", command, planId, title, steps);
            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.setOrder(i);
            planStepResp.setDescription(steps.get(i));
            planSteps.add(planStepResp);
        }
        plan.setSteps(planSteps);
        this.currentPlan = plan;
        return new ToolExecuteResult("Plan created: " + planId + "\n" + plan.getPlanExecutionStateStringFormat(false));
    }

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

}
