#include "planning.h"
#include <iomanip> // 添加iomanip头文件，用于std::setprecision

namespace humanus {

// Initialize the agent with a default plan ID and validate required tools.
void PlanningAgent::initialize_plan_and_verify_tools() {
    active_plan_id = "plan_" + std::to_string(std::chrono::system_clock::now().time_since_epoch().count());

    if (available_tools.tools_map.find("planning") == available_tools.tools_map.end()) {
        available_tools.add_tool(std::make_shared<PlanningTool>());
    }
}

// Decide the next action based on plan status.
bool PlanningAgent::think() {
    std::string prompt;
    if (!active_plan_id.empty()) {
        prompt = "CURRENT PLAN STATUS:\n" + get_plan() + "\n\n" + next_step_prompt;
    } else {
        prompt = next_step_prompt;
    }
    memory->add_message(Message::user_message(prompt));
    
    // Get the current step index before thinking
    current_step_index = _get_current_step_index();

    bool result = ToolCallAgent::think(); // Will set tool_calls

    // After thinking, if we decided to execute a tool and it's not a planning tool or special tool,
    // associate it with the current step for tracking
    if (result && !tool_calls.empty()) {
        auto latest_tool_call = tool_calls.back(); // Get the most recent tool call
        if (latest_tool_call.function.name != "planning"
        && !_is_special_tool(latest_tool_call.function.name)
        && current_step_index >= 0) {
            step_execution_tracker[latest_tool_call.id] = {
                {"step_index", current_step_index},
                {"tool_name", latest_tool_call.function.name},
                {"status", "pending"} // Will be updated after execution
            };
        }
    }

    return result;
}

// Execute a step and track its completion status.
std::string PlanningAgent::act() {
    std::string result = ToolCallAgent::act();

    // After executing the tool, update the plan status
    if (!tool_calls.empty()) {
        auto latest_tool_call = tool_calls.back();

        // Update the execution status to completed
        if (step_execution_tracker.find(latest_tool_call.id) != step_execution_tracker.end()) {
            step_execution_tracker[latest_tool_call.id]["status"] = "completed";
            step_execution_tracker[latest_tool_call.id]["result"] = result;

            // Update the plan status if this was a non-planning, non-special tool
            if (
                latest_tool_call.function.name != "planning"
                && !_is_special_tool(latest_tool_call.function.name)
            ) {
                update_plan_status(latest_tool_call.id);
            }
        }
    }

    return result;
}

// Retrieve the current plan status.
std::string PlanningAgent::get_plan() {
    if (active_plan_id.empty()) {
        return "No active plan. Please create a plan first.";
    }

    ToolResult result = available_tools.execute(
        "planning",
        {{"command", "get"}, {"plan_id", active_plan_id}}
    );

    return result.to_string();
}

// Run the agent with an optional initial request.
std::string PlanningAgent::run(const std::string& request) {
    if (!request.empty()) {
        create_initial_plan(request);
    }

    return BaseAgent::run();
}

// Update the current plan progress based on completed tool execution.
// Only marks a step as completed if the associated tool has been successfully executed.
void PlanningAgent::update_plan_status(const std::string& tool_call_id) {
    if (active_plan_id.empty()) {
        return;
    }

    if (step_execution_tracker.find(tool_call_id) == step_execution_tracker.end()) {
        logger->warn("No step tracking found for tool call " + tool_call_id);
        return;
    }

    auto tracker = step_execution_tracker[tool_call_id];
    if (tracker["status"] != "completed") {
        logger->warn("Tool call " + tool_call_id + " has not completed successfully");
        return;
    }

    int step_index = tracker["step_index"];

    try {
        // Mark the step as completed
        available_tools.execute(
            "planning",
            {
                {"command", "mark_step"}, 
                {"plan_id", active_plan_id}, 
                {"step_index", step_index}, 
                {"step_status", "completed"}
            }
        );
        logger->info(
            "Marked step " + std::to_string(step_index) + " as completed in plan " + active_plan_id
        );
    } catch (const std::exception& e) {
        logger->warn("Failed to update plan status: " + std::string(e.what()));
    }
}

// Parse the current plan to identify the first non-completed step's index.
// Returns -1 if no active step is found.
int PlanningAgent::_get_current_step_index() {
    if (active_plan_id.empty()) {
        return -1;
    }

    auto plan = get_plan();

    auto splitlines = [](const std::string& s) {
        std::vector<std::string> lines;
        std::istringstream iss(s);
        std::string line;
        while (std::getline(iss, line)) {
            lines.push_back(line);
        }
        return lines;
    };

    try {
        std::vector<std::string> plan_lines = splitlines(plan);
        int steps_index = -1;

        // Find the index of the "Steps:" line
        for (size_t i = 0; i < plan_lines.size(); ++i) {
            if (plan_lines[i].find("Steps:") == 0) {
                steps_index = i;
                break;
            }
        }

        if (steps_index == -1) {
            return -1;
        }
        
        // Find the first non-completed step
        for (size_t i = steps_index + 1; i < plan_lines.size(); ++i) {
            std::string line = plan_lines[i];
            if (line.find("[ ]") != std::string::npos || line.find("[→]") != std::string::npos) { // not_started or in_progress
                // Mark current step as in_progress
                available_tools.execute(
                    "planning",
                    {
                        {"command", "mark_step"}, 
                        {"plan_id", active_plan_id}, 
                        {"step_index", i}, 
                        {"step_status", "in_progress"}
                    }
                );
                return i;
            }
        }

        return -1; // # No active step found
    } catch (const std::exception& e) {
        logger->warn("Error finding current step index: " + std::string(e.what()));
        return -1;
    }
}

// Create an initial plan based on the request.
void PlanningAgent::create_initial_plan(const std::string& request) {
    logger->info("Creating initial plan with ID: " + request);

    std::vector<Message> messages = {
        Message::user_message(
            "Analyze the request and create a plan with ID " + active_plan_id + ": " + request
        )
    };
    memory->add_messages(messages);
    json response = llm->ask_tool(
        messages,
        {Message::system_message(system_prompt)},
        available_tools.to_params(),
        tool_choice
    );

    tool_calls = ToolCall::from_json_list(response["tool_calls"]);

    Message assistant_msg = Message::from_tool_calls(
        tool_calls, response["content"]
    );

    memory->add_message(assistant_msg);

    bool plan_created = false;
    for (const ToolCall& tool_call : tool_calls) {
        if (tool_call.function.name == "planning") {
            std::string result = execute_tool(tool_call);
            logger->info("Executed tool " + tool_call.function.name + " with result: " + result);

            // Add tool response to memory
            Message tool_msg = Message::tool_message(
                result,
                tool_call.id,
                tool_call.function.name
            );
            memory->add_message(tool_msg);
            plan_created = true;
            break;
        }
    }

    if (!plan_created) {
        logger->warn("No plan created from initial request");
        Message tool_msg = Message::assistant_message(
            "Error: Parameter `plan_id` is required for command: create"
        );
        memory->add_message(tool_msg);
    }
}

} // namespace humanus