#pragma once

#include <cstdint>
#include <optional>
#include <string>
#include <variant>
#include <vector>

#include <nlohmann/json.hpp>

#include "error.hpp"

namespace agent_smith
{

/// The supported OpenAI-compatible conversation roles.
enum class Role
{
    SYSTEM,
    USER,
    ASSISTANT,
    TOOL,
};

/// Server-reported token counts for one completed LLM response.
struct TokenUsage
{
    std::uint64_t prompt = 0;
    std::uint64_t completion = 0;
    std::uint64_t total = 0;
};

/// A system instruction supplied to the LLM.
struct SystemMessage
{
    std::string content;
};

/// A user-authored text message.
struct UserMessage
{
    std::string content;
};

/// One fully parsed function invocation requested by the LLM.
struct ToolCall
{
    std::string id;
    std::string name;
    nlohmann::json arguments;
};

/// An LLM response that may contain text, tool calls, and token usage.
struct AssistantMessage
{
    std::string content;
    std::vector<ToolCall> tool_calls;
    std::optional<TokenUsage> usage;
};

/// The result produced for one LLM tool invocation.
struct ToolMessage
{
    std::string tool_call_id;
    std::string name;
    std::string content;
    bool is_error = false;
};

/// One complete message in an Agent Smith conversation.
using Message =
    std::variant<SystemMessage, UserMessage, AssistantMessage, ToolMessage>;

/// Convert a conversation message to its OpenAI-compatible wire representation.
nlohmann::json toOpenAIJson(const Message& message);

/// Parse an untrusted OpenAI-compatible assistant message payload.
mw::E<Message> assistantMessageFromOpenAIJson(const nlohmann::json& json);

} // namespace agent_smith