#include "include/tools/planning.h"
#include "mcp/include/mcp_tool.h"
#include <chrono>
#include <sstream>
#include <iomanip>

namespace openmanus {

Planning::Planning()
    : ToolBase("planning", "创建和管理结构化计划") {
    // 使用tool_builder构建工具定义
    mcp::tool_builder builder("planning");
    builder.with_description("创建和管理结构化计划");
    
    // 添加参数
    builder.with_string_param("action", "要执行的操作", true);
    builder.with_string_param("plan_id", "计划ID", false);
    builder.with_string_param("title", "计划标题", false);
    builder.with_string_param("description", "计划描述", false);
    builder.with_array_param("steps", "计划步骤列表", "string", false);
    builder.with_string_param("step_id", "步骤ID", false);
    builder.with_string_param("step_description", "步骤描述", false);
    builder.with_string_param("status", "步骤状态", false);
    
    mcp::tool tool = builder.build();
    
    // 从构建的工具中获取参数定义
    parameters_ = tool.parameters_schema;
    
    // 手动添加enum约束
    if (parameters_.contains("properties") && parameters_["properties"].contains("action")) {
        parameters_["properties"]["action"]["enum"] = {"create_plan", "get_plan_status", "update_step_status", "add_step"};
    }
    
    // 手动添加status的enum约束
    if (parameters_.contains("properties") && parameters_["properties"].contains("status")) {
        parameters_["properties"]["status"]["enum"] = {"pending", "in_progress", "completed", "failed"};
    }
}

mcp::json Planning::execute(const mcp::json& params) {
    if (!params.contains("action")) {
        return {
            {"success", false},
            {"error", "缺少action参数"}
        };
    }

    std::string action = params["action"];

    if (action == "create_plan") {
        return createPlan(params);
    } else if (action == "get_plan_status") {
        return getPlanStatus(params);
    } else if (action == "update_step_status") {
        return updateStepStatus(params);
    } else if (action == "add_step") {
        return addStep(params);
    } else {
        return {
            {"success", false},
            {"error", "未知的action: " + action}
        };
    }
}

mcp::json Planning::createPlan(const mcp::json& params) {
    if (!params.contains("title") || !params.contains("description")) {
        return {
            {"success", false},
            {"error", "缺少title或description参数"}
        };
    }

    // 生成计划ID
    auto now = std::chrono::system_clock::now();
    auto now_time_t = std::chrono::system_clock::to_time_t(now);
    std::stringstream ss;
    ss << "plan_" << std::put_time(std::localtime(&now_time_t), "%Y%m%d%H%M%S");
    std::string plan_id = ss.str();

    // 创建计划
    Plan plan;
    plan.id = plan_id;
    plan.title = params["title"];
    plan.description = params["description"];
    plan.current_step_index = 0;

    // 添加步骤
    if (params.contains("steps") && params["steps"].is_array()) {
        for (size_t i = 0; i < params["steps"].size(); ++i) {
            std::string step_id = "step_" + std::to_string(i + 1);
            std::string step_description = params["steps"][i];
            
            plan.step_ids.push_back(step_id);
            plan.step_descriptions[step_id] = step_description;
            plan.step_statuses[step_id] = "pending";
        }
    }

    // 保存计划
    plans_[plan_id] = plan;

    return {
        {"success", true},
        {"plan_id", plan_id},
        {"message", "计划创建成功"}
    };
}

mcp::json Planning::getPlanStatus(const mcp::json& params) {
    if (!params.contains("plan_id")) {
        return {
            {"success", false},
            {"error", "缺少plan_id参数"}
        };
    }

    std::string plan_id = params["plan_id"];
    
    // 查找计划
    auto it = plans_.find(plan_id);
    if (it == plans_.end()) {
        return {
            {"success", false},
            {"error", "找不到计划: " + plan_id}
        };
    }

    const Plan& plan = it->second;
    
    // 构建步骤信息
    mcp::json steps = mcp::json::array();
    for (const auto& step_id : plan.step_ids) {
        steps.push_back({
            {"id", step_id},
            {"description", plan.step_descriptions.at(step_id)},
            {"status", plan.step_statuses.at(step_id)}
        });
    }

    return {
        {"success", true},
        {"plan", {
            {"id", plan.id},
            {"title", plan.title},
            {"description", plan.description},
            {"steps", steps},
            {"current_step_index", plan.current_step_index}
        }}
    };
}

mcp::json Planning::updateStepStatus(const mcp::json& params) {
    if (!params.contains("plan_id") || !params.contains("step_id") || !params.contains("status")) {
        return {
            {"success", false},
            {"error", "缺少plan_id、step_id或status参数"}
        };
    }

    std::string plan_id = params["plan_id"];
    std::string step_id = params["step_id"];
    std::string status = params["status"];
    
    // 查找计划
    auto it = plans_.find(plan_id);
    if (it == plans_.end()) {
        return {
            {"success", false},
            {"error", "找不到计划: " + plan_id}
        };
    }

    Plan& plan = it->second;
    
    // 查找步骤
    auto step_status_it = plan.step_statuses.find(step_id);
    if (step_status_it == plan.step_statuses.end()) {
        return {
            {"success", false},
            {"error", "找不到步骤: " + step_id}
        };
    }

    // 更新步骤状态
    step_status_it->second = status;
    
    // 如果步骤完成，更新当前步骤索引
    if (status == "completed") {
        for (size_t i = 0; i < plan.step_ids.size(); ++i) {
            if (plan.step_ids[i] == step_id && i == plan.current_step_index) {
                plan.current_step_index = i + 1;
                break;
            }
        }
    }

    return {
        {"success", true},
        {"message", "步骤状态更新成功"}
    };
}

mcp::json Planning::addStep(const mcp::json& params) {
    if (!params.contains("plan_id") || !params.contains("step_description")) {
        return {
            {"success", false},
            {"error", "缺少plan_id或step_description参数"}
        };
    }

    std::string plan_id = params["plan_id"];
    std::string step_description = params["step_description"];
    
    // 查找计划
    auto it = plans_.find(plan_id);
    if (it == plans_.end()) {
        return {
            {"success", false},
            {"error", "找不到计划: " + plan_id}
        };
    }

    Plan& plan = it->second;
    
    // 生成步骤ID
    std::string step_id = "step_" + std::to_string(plan.step_ids.size() + 1);
    
    // 添加步骤
    plan.step_ids.push_back(step_id);
    plan.step_descriptions[step_id] = step_description;
    plan.step_statuses[step_id] = "pending";

    return {
        {"success", true},
        {"step_id", step_id},
        {"message", "步骤添加成功"}
    };
}

} // namespace openmanus