#ifndef HUMANUS_AGENT_PLANNING_H
#define HUMANUS_AGENT_PLANNING_H

#include "toolcall.h"
#include "../tool/planning.h"
#include "../prompt.h"

namespace humanus {

/**
 * An agent that creates and manages plans to solve tasks.

 * This agent uses a planning tool to create and manage structured plans,
 * and tracks progress through individual steps until task completion.
*/
struct PlanningAgent : ToolCallAgent {
    std::string active_plan_id;

    // Add a dictionary to track the step status for each tool call
    std::map<std::string, json> step_execution_tracker;
    int current_step_index;

    PlanningAgent(
        const ToolCollection& available_tools = ToolCollection(
            {
                std::make_shared<PlanningTool>(),
                std::make_shared<Terminate>()
            }
        ),
        const std::string& tool_choice = "auto",
        const std::set<std::string>& special_tool_names = {"terminate"},
        const std::string& name = "planning", 
        const std::string& description = "An agent that creates and manages plans to solve tasks", 
        const std::string& system_prompt = prompt::planning::PLANNING_SYSTEM_PROMPT,
        const std::string& next_step_prompt = prompt::planning::NEXT_STEP_PROMPT,
        const std::shared_ptr<LLM>& llm = nullptr,
        const std::shared_ptr<Memory>& memory = nullptr,
        AgentState state = AgentState::IDLE,
        int max_steps = 20,
        int current_step = 0,
        int duplicate_threshold = 2
    ) : ToolCallAgent(
            available_tools, 
            tool_choice, 
            special_tool_names, 
            name, 
            description, 
            system_prompt, 
            next_step_prompt, 
            llm, 
            memory, 
            state, 
            max_steps, 
            current_step, 
            duplicate_threshold
        ) {
        current_step_index = -1; // will be set in think()
        initialize_plan_and_verify_tools();
    }

    // Initialize the agent with a default plan ID and validate required tools.
    void initialize_plan_and_verify_tools();

    // Decide the next action based on plan status.
    bool think() override;

    // Execute a step and track its completion status.
    std::string act() override;

    // Retrieve the current plan status.
    std::string get_plan();

    // Run the agent with an optional initial request.
    std::string run(const std::string& request = "") override;

    // 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 update_plan_status(const std::string& tool_call_id);

    // Parse the current plan to identify the first non-completed step's index.
    // Returns None if no active step is found.
    int _get_current_step_index();

    // Create an initial plan based on the request.
    void create_initial_plan(const std::string& request);
};


} // namespace humanus

#endif // HUMANUS_AGENT_PLANNING_H