/**
 * @file custom_agent.cpp
 * @brief Implementation of custom MCP agents
 */

#include "custom_agent.h"
#include <algorithm>
#include <stdexcept>

namespace examples {

// echo_agent implementation
echo_agent::echo_agent(const std::string& name)
    : mcp::agent(name) {
}

void echo_agent::initialize(const mcp::json& config) {
    // Configure the agent from the provided config
    if (config.contains("uppercase") && config["uppercase"].is_boolean()) {
        uppercase_enabled_ = config["uppercase"];
    }
    
    if (config.contains("prefix") && config["prefix"].is_boolean()) {
        prefix_enabled_ = config["prefix"];
    }
    
    if (config.contains("prefix_text") && config["prefix_text"].is_string()) {
        prefix_ = config["prefix_text"];
    }
    
    // Initialize any resources or state
    request_count_ = 0;
}

mcp::json echo_agent::process(const mcp::json& input) {
    request_count_++;
    
    mcp::json output = {
        {"request_count", request_count_}
    };
    
    // Process the text if provided
    if (input.contains("text") && input["text"].is_string()) {
        std::string text = input["text"];
        
        // Apply uppercase if enabled
        if (uppercase_enabled_) {
            std::transform(text.begin(), text.end(), text.begin(), ::toupper);
        }
        
        // Apply prefix if enabled
        if (prefix_enabled_) {
            text = prefix_ + text;
        }
        
        output["processed_text"] = text;
    }
    
    // Echo back all input
    output["original_input"] = input;
    
    return output;
}

void echo_agent::set_uppercase(bool enable) {
    uppercase_enabled_ = enable;
}

void echo_agent::set_prefix(bool enable, const std::string& prefix) {
    prefix_enabled_ = enable;
    if (!prefix.empty()) {
        prefix_ = prefix;
    }
}

// workflow_agent implementation
workflow_agent::workflow_agent(const std::string& name)
    : mcp::agent(name) {
}

void workflow_agent::initialize(const mcp::json& config) {
    // Clear existing workflows
    workflows_.clear();
    
    // Load workflows from config if provided
    if (config.contains("workflows") && config["workflows"].is_array()) {
        for (const auto& wf_json : config["workflows"]) {
            if (wf_json.contains("name") && wf_json.contains("steps") && 
                wf_json["name"].is_string() && wf_json["steps"].is_array()) {
                
                std::string wf_name = wf_json["name"];
                std::string wf_description = wf_json.value("description", "");
                
                mcp::workflow wf(wf_name, wf_description);
                
                // Add steps
                for (const auto& step_json : wf_json["steps"]) {
                    if (step_json.contains("tool_name") && step_json["tool_name"].is_string()) {
                        std::string tool_name = step_json["tool_name"];
                        mcp::json parameters = step_json.value("parameters", mcp::json::object());
                        
                        wf.add_tool_call(tool_name, parameters);
                    }
                }
                
                workflows_[wf_name] = wf;
            }
        }
    }
}

mcp::json workflow_agent::process(const mcp::json& input) {
    // Check if workflow name is provided
    if (!input.contains("workflow") || !input["workflow"].is_string()) {
        throw std::runtime_error("Missing required parameter: workflow (string)");
    }
    
    std::string workflow_name = input["workflow"];
    
    // Check if the workflow exists
    auto it = workflows_.find(workflow_name);
    if (it == workflows_.end()) {
        throw std::runtime_error("Workflow not found: " + workflow_name);
    }
    
    // Get the context from the input
    mcp::json context = input.value("context", mcp::json::object());
    
    // Execute the workflow
    mcp::json result = execute_workflow(it->second, context);
    
    return result;
}

void workflow_agent::register_workflow(const mcp::workflow& workflow) {
    workflows_[workflow.name()] = workflow;
}

// chain_agent implementation
chain_agent::chain_agent(const std::string& name)
    : mcp::agent(name) {
}

void chain_agent::initialize(const mcp::json& config) {
    // Clear existing chain
    chain_.clear();
    
    // Load chain from config if provided
    if (config.contains("chain") && config["chain"].is_array()) {
        for (const auto& step_json : config["chain"]) {
            if (step_json.contains("tool") && step_json["tool"].is_string()) {
                std::string tool_name = step_json["tool"];
                
                std::map<std::string, std::string> output_to_input_map;
                
                // Load mappings if provided
                if (step_json.contains("mappings") && step_json["mappings"].is_object()) {
                    for (auto it = step_json["mappings"].begin(); it != step_json["mappings"].end(); ++it) {
                        if (it.value().is_string()) {
                            output_to_input_map[it.key()] = it.value();
                        }
                    }
                }
                
                add_tool(tool_name, output_to_input_map);
            }
        }
    }
}

mcp::json chain_agent::process(const mcp::json& input) {
    if (chain_.empty()) {
        throw std::runtime_error("Chain is empty");
    }
    
    mcp::json current_context = input;
    
    // Execute each step in the chain
    for (const auto& step : chain_) {
        // Prepare parameters for this step
        mcp::json parameters = mcp::json::object();
        
        // Apply output-to-input mappings
        for (const auto& [output_path, input_path] : step.output_to_input_map) {
            // Extract value from the current context
            if (current_context.contains(output_path)) {
                // Set the value in the parameters (at the mapped path)
                parameters[input_path] = current_context[output_path];
            }
        }
        
        // For fields not mapped, take them directly from the input if they exist
        for (auto it = input.begin(); it != input.end(); ++it) {
            if (!parameters.contains(it.key())) {
                parameters[it.key()] = it.value();
            }
        }
        
        // Call the tool
        mcp::json result = call_tool(step.tool_name, parameters);
        
        // Update the context with the result
        for (auto it = result.begin(); it != result.end(); ++it) {
            current_context[it.key()] = it.value();
        }
    }
    
    return current_context;
}

void chain_agent::add_tool(const std::string& tool_name,
                          const std::map<std::string, std::string>& output_to_input_map) {
    chain_step step;
    step.tool_name = tool_name;
    step.output_to_input_map = output_to_input_map;
    chain_.push_back(step);
}

void chain_agent::clear_chain() {
    chain_.clear();
}

} // namespace examples