#ifndef HUMANUS_SCHEMA_H
#define HUMANUS_SCHEMA_H

#include "mcp_message.h"
#include <string>
#include <vector>
#include <map>
#include <algorithm>

namespace humanus {

using json = mcp::json;

// Agent execution states
enum class AgentState {
    IDLE = 0,
    RUNNING = 1,
    FINISHED = 2,
    ERR = 3 // Don't use ERROR
};

extern std::map<AgentState, std::string> agent_state_map;

struct Function {
    std::string name;
    json arguments;

    json to_json() const {
        json function;
        function["name"] = name;
        function["arguments"] = arguments;
        return function;
    }
    
    bool empty() const {
        return name.empty() && arguments.empty();
    }
};

// Represents a tool/function call in a message
struct ToolCall {
    std::string id;
    std::string type;
    Function function;

    bool empty() const {
        return id.empty() && type.empty() && function.empty();
    }

    json to_json() const {
        json tool_call;
        tool_call["id"] = id;
        tool_call["type"] = type;
        tool_call["function"] = function.to_json();
        return tool_call;
    }

    static ToolCall from_json(const json& tool_call_json) {
        ToolCall tool_call;
        tool_call.id = tool_call_json["id"];
        tool_call.type = tool_call_json["type"];
        tool_call.function.name = tool_call_json["function"]["name"];
        tool_call.function.arguments = tool_call_json["function"]["arguments"];
        return tool_call;
    }

    static std::vector<ToolCall> from_json_list(const json& tool_calls_json) {
        std::vector<ToolCall> tool_calls;
        for (const auto& tool_call_json : tool_calls_json) {
            tool_calls.push_back(from_json(tool_call_json));
        }
        return tool_calls;
    }
};

// Represents a chat message in the conversation
struct Message {
    std::string role;
    json content;
    std::string name;
    std::string tool_call_id;
    std::vector<ToolCall> tool_calls;

    Message(const std::string& role, const json& content, const std::string& name = "", const std::string& tool_call_id = "", const std::vector<ToolCall> tool_calls = {})
    : role(role), content(content), name(name), tool_call_id(tool_call_id), tool_calls(tool_calls) {}

    std::vector<Message> operator+(const Message& other) const {
        return {*this, other};
    }

    std::vector<Message> operator+(const std::vector<Message>& other) const {
        std::vector<Message> result;
        result.reserve(1 + other.size());
        result.push_back(*this);
        result.insert(result.end(), other.begin(), other.end());
        return result;
    }

    friend std::vector<Message> operator+(const std::vector<Message>& lhs, const Message& rhs) {
        std::vector<Message> result;
        result.reserve(lhs.size() + 1);
        result.insert(result.end(), lhs.begin(), lhs.end());
        result.push_back(rhs);
        return result;
    }

    friend std::vector<Message> operator+(const std::vector<Message>& lhs, const std::vector<Message>& rhs) {
        std::vector<Message> result;
        result.reserve(lhs.size() + rhs.size());
        result.insert(result.end(), lhs.begin(), lhs.end());
        result.insert(result.end(), rhs.begin(), rhs.end());
        return result;
    }

    // Convert message to dictionary format
    json to_json() const {
        json message;
        message["role"] = role;
        if (!content.empty()) {
            message["content"] = content;
        }
        if (!tool_calls.empty()) {
            message["tool_calls"] = json::array();
            for (const auto& call : tool_calls) {
                message["tool_calls"].push_back(call.to_json());
            }
        }
        if (!name.empty()) {
            message["name"] = name;
        }
        if (!tool_call_id.empty()) {
            message["tool_call_id"] = tool_call_id;
        }
        return message;
    }

    // Convert message to dictionary format
    json to_dict() const {
        return to_json();
    }

    static Message user_message(const json& content) {
        return Message("user", content);
    }

    static Message system_message(const json& content) {
        return Message("system", content);
    }

    static Message tool_message(const json& content, const std::string& tool_call_id = "", const std::string& name =  "") {
        return Message("tool", content, name, tool_call_id);
    }

    static Message assistant_message(const json& content = "", const std::vector<ToolCall>& tool_calls = {}) {
        return Message("assistant", content, "", "", tool_calls);
    }
};

struct Memory {
    std::vector<Message> messages;
    int max_messages;

    Memory(int max_messages = 100) : max_messages(max_messages) {}

    // Add a message to the memory
    void add_message(const Message& message) {
        messages.push_back(message);
        while (!messages.empty() && (messages.size() > max_messages || messages.begin()->role == "assistant" || messages.begin()->role == "tool")) {
            // Ensure the first message is always a user or system message
            messages.erase(messages.begin());
        }
    }

    // Add multiple messages to the memory
    void add_messages(const std::vector<Message>& messages) {
        for (const auto& message : messages) {
            add_message(message);
        }
    }
    
    // Clear all messages
    void clear() {
        messages.clear();
    }

    // Get the last n messages
    std::vector<Message> get_recent_messages(int n) const {
        n = std::min(n, static_cast<int>(messages.size()));
        return std::vector<Message>(messages.end() - n, messages.end());
    }

    // Convert messages to list of dicts
    json to_json_list() const {
        json memory = json::array();
        for (const auto& message : messages) {
            memory.push_back(message.to_json());
        }
        return memory;
    }
};

} // namespace humanus

#endif // HUMANUS_SCHEMA_H