#include "OpenAIProvider.hpp"
#include <spdlog/sinks/stdout_color_sinks.h>
#include <cstdlib>

namespace aissia {

OpenAIProvider::OpenAIProvider(const json& config) {
    m_logger = spdlog::get("OpenAIProvider");
    if (!m_logger) {
        m_logger = spdlog::stdout_color_mt("OpenAIProvider");
    }

    // Get API key from environment
    std::string apiKeyEnv = config.value("api_key_env", "OPENAI_API_KEY");
    const char* apiKey = std::getenv(apiKeyEnv.c_str());
    if (!apiKey) {
        m_logger->error("API key not found in environment: {}", apiKeyEnv);
        throw std::runtime_error("Missing API key: " + apiKeyEnv);
    }
    m_apiKey = apiKey;

    m_model = config.value("model", "gpt-4o");
    m_maxTokens = config.value("max_tokens", 4096);
    m_baseUrl = config.value("base_url", "https://api.openai.com");

    m_client = std::make_unique<HttpClient>(m_baseUrl, 60);
    m_client->setBearerToken(m_apiKey);
    m_client->setHeader("Content-Type", "application/json");

    m_logger->info("OpenAIProvider initialized: model={}", m_model);
}

LLMResponse OpenAIProvider::chat(const std::string& systemPrompt,
                                  const json& messages,
                                  const json& tools) {
    // Build messages array with system prompt
    json allMessages = json::array();
    allMessages.push_back({
        {"role", "system"},
        {"content", systemPrompt}
    });
    for (const auto& msg : messages) {
        allMessages.push_back(msg);
    }

    json request = {
        {"model", m_model},
        {"max_tokens", m_maxTokens},
        {"messages", allMessages}
    };

    if (!tools.empty()) {
        request["tools"] = convertTools(tools);
    }

    m_logger->debug("Sending request to OpenAI: {} messages", allMessages.size());

    auto response = m_client->post("/v1/chat/completions", request);

    if (!response.success) {
        m_logger->error("OpenAI API error: {}", response.error);
        LLMResponse errorResp;
        errorResp.text = "Error: " + response.error;
        errorResp.is_end_turn = true;
        return errorResp;
    }

    try {
        json jsonResponse = json::parse(response.body);
        m_lastRawResponse = jsonResponse;
        return parseResponse(jsonResponse);
    } catch (const json::exception& e) {
        m_logger->error("Failed to parse OpenAI response: {}", e.what());
        LLMResponse errorResp;
        errorResp.text = "Parse error: " + std::string(e.what());
        errorResp.is_end_turn = true;
        return errorResp;
    }
}

json OpenAIProvider::convertTools(const json& tools) {
    // Convert to OpenAI function format
    json openaiTools = json::array();

    for (const auto& tool : tools) {
        json openaiTool = {
            {"type", "function"},
            {"function", {
                {"name", tool.value("name", "")},
                {"description", tool.value("description", "")},
                {"parameters", tool.value("input_schema", json::object())}
            }}
        };
        openaiTools.push_back(openaiTool);
    }

    return openaiTools;
}

LLMResponse OpenAIProvider::parseResponse(const json& response) {
    LLMResponse result;

    // Parse usage
    if (response.contains("usage")) {
        result.input_tokens = response["usage"].value("prompt_tokens", 0);
        result.output_tokens = response["usage"].value("completion_tokens", 0);
    }

    result.model = response.value("model", m_model);

    // Parse choices
    if (response.contains("choices") && !response["choices"].empty()) {
        const auto& choice = response["choices"][0];
        result.stop_reason = choice.value("finish_reason", "");

        if (choice.contains("message")) {
            const auto& message = choice["message"];

            // Get text content
            if (message.contains("content") && !message["content"].is_null()) {
                result.text = message["content"].get<std::string>();
            }

            // Get tool calls
            if (message.contains("tool_calls") && message["tool_calls"].is_array()) {
                for (const auto& tc : message["tool_calls"]) {
                    ToolCall call;
                    call.id = tc.value("id", "");
                    if (tc.contains("function")) {
                        call.name = tc["function"].value("name", "");
                        std::string args = tc["function"].value("arguments", "{}");
                        try {
                            call.input = json::parse(args);
                        } catch (...) {
                            call.input = json::object();
                        }
                    }
                    result.tool_calls.push_back(call);
                }
            }
        }
    }

    // Determine if this is end turn
    result.is_end_turn = result.tool_calls.empty() ||
                         result.stop_reason == "stop";

    m_logger->debug("OpenAI response: text={} chars, tools={}, stop={}",
                   result.text.size(), result.tool_calls.size(), result.stop_reason);

    return result;
}

json OpenAIProvider::formatToolResults(const std::vector<ToolResult>& results) {
    // OpenAI format: separate messages with role "tool"
    json messages = json::array();

    for (const auto& result : results) {
        messages.push_back({
            {"role", "tool"},
            {"tool_call_id", result.tool_call_id},
            {"content", result.content}
        });
    }

    // Return array of messages (caller should append each)
    return messages;
}

void OpenAIProvider::appendAssistantMessage(json& messages, const LLMResponse& response) {
    // Build assistant message
    json assistantMsg = {
        {"role", "assistant"}
    };

    if (!response.text.empty()) {
        assistantMsg["content"] = response.text;
    } else {
        assistantMsg["content"] = nullptr;
    }

    // Add tool_calls if present
    if (!response.tool_calls.empty()) {
        json toolCalls = json::array();
        for (const auto& call : response.tool_calls) {
            toolCalls.push_back({
                {"id", call.id},
                {"type", "function"},
                {"function", {
                    {"name", call.name},
                    {"arguments", call.input.dump()}
                }}
            });
        }
        assistantMsg["tool_calls"] = toolCalls;
    }

    messages.push_back(assistantMsg);
}

} // namespace aissia