#include "toolcall.h"

namespace humanus {

// Process current state and decide next actions using tools
bool ToolCallAgent::think() {
    // Get response with tool options
    auto response = llm->ask_tool(
        memory->messages,
        system_prompt,
        next_step_prompt,
        available_tools.to_params(),
        tool_choice
    );

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

    // Log response info
    logger->info("✨ " + name + "'s thoughts: " + response["content"].get<std::string>());
    logger->info(
        "🛠️  " + name + " selected " + std::to_string(tool_calls.size()) + " tool(s) to use"
    );
    if (tool_calls.size() > 0) {
        std::string tools_str;
        for (const auto& tool_call : tool_calls) {
            tools_str += tool_call.function.name + " ";
        }
        logger->info(
            "🧰 Tools being prepared: " + tools_str
        );
    }

    try {
        // Handle different tool_choices modes
        if (tool_choice == "none") {
            if (tool_calls.size() > 0) {
                logger->warn("🤔 Hmm, " + name + "tried to use tools when they weren't available!");
            }
            if (!response["content"].empty()) {
                memory->add_message(Message::assistant_message(response["content"]));
                return true;
            }
            return false;
        } 

        // Create and add assistant message
        auto assistant_msg = Message::assistant_message(response["content"], tool_calls);
        memory->add_message(assistant_msg);

        if (tool_choice == "required" && tool_calls.empty()) {
            return true; // Will be handled in act()
        }

        // For 'auto' mode, continue with content if no commands but content exists
        if (tool_choice == "auto" && !tool_calls.empty()) {
            return !response["content"].empty();
        }

        return !tool_calls.empty();
    } catch (const std::exception& e) {
        logger->error("🚨 Oops! The " + name + "'s thinking process hit a snag: " + std::string(e.what()));
        memory->add_message(Message::assistant_message(
            "Error encountered while processing: " + std::string(e.what())
        ));
        return false;
    }
}

// Execute tool calls and handle their results
std::string ToolCallAgent::act() {
    if (tool_calls.empty()) {
        if (tool_choice == "required") {
            throw std::runtime_error("Required tools but none selected");
        }

        // Return last message content if no tool calls
        return memory->messages.empty() || memory->messages.back().content.empty() ? "No content or commands to execute" : memory->messages.back().content.dump();
    }

    std::vector<std::string> results;
    for (const auto& tool_call : tool_calls) {
        auto result = execute_tool(tool_call);
        logger->info(
            "🎯 Tool '" + tool_call.function.name + "' completed its mission! Result: " + result.substr(0, 500) + (result.size() > 500 ? "..." : "")
        );

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

    std::string result_str;
    for (const auto& result : results) {
        result_str += result + "\n\n";
    }

    return result_str;
}

// Execute a single tool call with robust error handling
std::string ToolCallAgent::execute_tool(ToolCall tool_call) {
    if (tool_call.empty() || tool_call.function.empty() || tool_call.function.name.empty()) {
        return "Error: Invalid command format";
    }

    std::string name = tool_call.function.name;
    if (available_tools.tools_map.find(name) == available_tools.tools_map.end()) {
        return "Error: Unknown tool '" + name + "'. Please use one of the following tools: " + 
               std::accumulate(available_tools.tools_map.begin(), available_tools.tools_map.end(), std::string(),
                               [](const std::string& a, const auto& b) {
                                   return a + (a.empty() ? "" : ", ") + b.first;
                               });
    }

    try {
        // Parse arguments
        json args = tool_call.function.arguments;

        if (args.is_string()) {
            args = json::parse(args.get<std::string>());
        }

        // Execute the tool
        logger->info("🔧 Activating tool: '" + name + "'...");
        ToolResult result = available_tools.execute(name, args);

        // Format result for display
        auto observation = result.empty() ?
                            "Cmd `" + name + "` completed with no output" :
                            "Observed output of cmd `" + name + "` executed:\n" + result.to_string();

        // Handle special tools like `finish`
        _handle_special_tool(name, result);

        return observation;
    } catch (const json::exception& /* e */) {
        std::string error_msg = "Error parsing arguments for " + name + ": Invalid JSON format";
        logger->error(
            "📝 Oops! The arguments for '" + name + "' don't make sense - invalid JSON"
        );
        return "Error: " + error_msg;
    } catch (const std::exception& e) {
        std::string error_msg = "⚠️ Tool '" + name + "' encountered a problem: " + std::string(e.what());
        logger->error(error_msg);
        return "Error: " + error_msg;
    }
}

// Handle special tool execution and state changes
void ToolCallAgent::_handle_special_tool(const std::string& name, const ToolResult& result, const json& kwargs) {
    if (!_is_special_tool(name)) {
        return;
    }

    if (_should_finish_execution(name, result, kwargs)) {
        logger->info("🏁 Special tool '" + name + "' has completed the task!");
        state = AgentState::FINISHED;
    }
}

// Determine if tool execution should finish the agent
bool ToolCallAgent::_should_finish_execution(const std::string& name, const ToolResult& result, const json& kwargs) {
    return true; // Currently, all special tools (terminate) finish the agent
}

bool ToolCallAgent::_is_special_tool(const std::string& name) {
    return special_tool_names.find(name) != special_tool_names.end();
}
    
}