#ifndef HUMANUS_FLOW_PLANNING_H
#define HUMANUS_FLOW_PLANNING_H

#include "base.h"
#include "../agent/base.h"
#include "../llm.h"
#include "../logger.h"
#include "../schema.h"
#include "../tool/planning.h"
#include <regex>

namespace humanus {

// A flow that manages planning and execution of tasks using agents.
struct PlanningFlow : public BaseFlow {
    std::shared_ptr<LLM> llm;
    std::shared_ptr<PlanningTool> planning_tool;
    std::vector<std::string> executor_keys;
    std::string active_plan_id;
    int current_step_index = -1;

    PlanningFlow(
        const std::shared_ptr<LLM>& llm = nullptr,
        const std::shared_ptr<PlanningTool>& planning_tool = nullptr,
        const std::vector<std::string>& executor_keys = {},
        const std::string& active_plan_id = "",
        const std::map<std::string, std::shared_ptr<BaseAgent>>& agents = {},
        const std::vector<std::shared_ptr<BaseTool>>& tools = {}, 
        const std::string& primary_agent_key = ""
    ) : BaseFlow(agents, tools, primary_agent_key),
        llm(llm),
        planning_tool(planning_tool),
        executor_keys(executor_keys),
        active_plan_id(active_plan_id) {
        if (!llm) {
            this->llm = LLM::get_instance("default");
        }
        if (!planning_tool) {
            this->planning_tool = std::make_shared<PlanningTool>();
        }
        if (active_plan_id.empty()) {
            this->active_plan_id = "plan_" + std::to_string(std::chrono::system_clock::now().time_since_epoch().count());
        }
        if (executor_keys.empty()) {
            for (const auto& [key, agent] : agents) {
                this->executor_keys.push_back(key);
            }
        }
    }
    
    // Get an appropriate executor agent for the current step.
    // Can be extended to select agents based on step type/requirements.
    std::shared_ptr<BaseAgent> get_executor(const std::string& step_type = "") const;

    // Execute the planning flow with agents.
    std::string execute(const std::string& input) override;

    // Create an initial plan based on the request using the flow's LLM and PlanningTool.
    void _create_initial_plan(const std::string& request);

    // Parse the current plan to identify the first non-completed step's index and info.
    // Returns (None, None) if no active step is found.
    void _get_current_step_info(int& current_step_index, json& step_info);

    // Execute the current step with the specified agent using agent.run().
    std::string _execute_step(const std::shared_ptr<BaseAgent>& executor, const json& step_info);

    // Mark the current step as completed.
    void _mark_step_completed();

    // Get the current plan as formatted text.
    std::string _get_plan_text();
    
    // Generate plan text directly from storage if the planning tool fails.
    std::string _generate_plan_text_from_storage();

    // Finalize the plan and provide a summary using the flow's LLM directly
    std::string _finalize_plan();
};

}

#endif // HUMANUS_FLOW_PLANNING_H