#include "lemon/routing_classifier_services.h"

#include <cmath>
#include <memory>
#include <optional>
#include <stdexcept>
#include <utility>

namespace lemon {
namespace {

void throw_if_error_response(const json& response, const std::string& context) {
    if (!response.is_object() || !response.contains("error")) {
        return;
    }
    std::string message = context + " failed";
    const json& error = response["error"];
    if (error.is_object() && error.contains("message") && error["message"].is_string()) {
        message += ": " + error["message"].get<std::string>();
    }
    throw std::runtime_error(message);
}

void ensure_model(const EnsureClassifierModelLoaded& ensure_loaded,
                  const std::string& model) {
    if (ensure_loaded) {
        ensure_loaded(model);
    }
}

std::map<std::string, double> scores_from_object(const json& obj) {
    std::map<std::string, double> scores;
    if (!obj.is_object()) {
        return scores;
    }
    for (auto it = obj.begin(); it != obj.end(); ++it) {
        if (it.value().is_number()) {
            scores[it.key()] = it.value().get<double>();
        }
    }
    return scores;
}

std::map<std::string, double> scores_from_label_score_array(const json& arr) {
    std::map<std::string, double> scores;
    if (!arr.is_array()) {
        return scores;
    }
    for (const auto& item : arr) {
        if (!item.is_object() ||
            !item.contains("label") || !item["label"].is_string() ||
            !item.contains("score") || !item["score"].is_number()) {
            continue;
        }
        scores[item["label"].get<std::string>()] = item["score"].get<double>();
    }
    return scores;
}

std::map<std::string, double> parse_scores_payload(const json& payload) {
    if (payload.is_array()) {
        return scores_from_label_score_array(payload);
    }
    if (!payload.is_object()) {
        return {};
    }

    for (const char* key : {"labels", "scores", "classification", "classifications"}) {
        if (!payload.contains(key)) {
            continue;
        }
        if (payload[key].is_object()) {
            auto scores = scores_from_object(payload[key]);
            if (!scores.empty()) return scores;
        }
        if (payload[key].is_array()) {
            auto scores = scores_from_label_score_array(payload[key]);
            if (!scores.empty()) return scores;
        }
    }

    if (payload.contains("label") && payload["label"].is_string() &&
        payload.contains("score") && payload["score"].is_number()) {
        return {{payload["label"].get<std::string>(), payload["score"].get<double>()}};
    }

    if (payload.contains("data") && payload["data"].is_array()) {
        auto scores = scores_from_label_score_array(payload["data"]);
        if (!scores.empty()) return scores;
    }

    return scores_from_object(payload);
}

json parse_json_text(const std::string& text) {
    json parsed = json::parse(text, nullptr, /*allow_exceptions=*/false);
    return parsed.is_discarded() ? json(nullptr) : parsed;
}

void validate_score_range(const std::map<std::string, double>& scores) {
    for (const auto& [label, score] : scores) {
        if (!std::isfinite(score) || score < 0.0 || score > 1.0) {
            throw std::runtime_error(
                "classifier score for label '" + label + "' must be in [0, 1]");
        }
    }
}

std::optional<std::string> try_extract_chat_text(const json& response) {
    if (response.is_object() && response.contains("choices") &&
        response["choices"].is_array() && !response["choices"].empty()) {
        const json& choice = response["choices"].front();
        if (choice.is_object()) {
            if (choice.contains("message") && choice["message"].is_object()) {
                const json& message = choice["message"];
                if (message.contains("content") && message["content"].is_string()) {
                    return message["content"].get<std::string>();
                }
            }
            if (choice.contains("text") && choice["text"].is_string()) {
                return choice["text"].get<std::string>();
            }
        }
    }

    if (response.is_object() && response.contains("content") &&
        response["content"].is_string()) {
        return response["content"].get<std::string>();
    }

    return std::nullopt;
}

} // namespace

std::vector<float> parse_embedding_vector(const json& response) {
    throw_if_error_response(response, "embedding request");

    const json* embedding = nullptr;
    if (response.is_object() && response.contains("data") &&
        response["data"].is_array() && !response["data"].empty()) {
        const json& first = response["data"].front();
        if (first.is_object() && first.contains("embedding")) {
            embedding = &first["embedding"];
        }
    }
    if (!embedding && response.is_object() && response.contains("embedding")) {
        embedding = &response["embedding"];
    }
    if (!embedding || !embedding->is_array()) {
        throw std::runtime_error("embedding response did not contain an embedding array");
    }

    std::vector<float> result;
    result.reserve(embedding->size());
    for (const auto& value : *embedding) {
        if (!value.is_number()) {
            throw std::runtime_error("embedding response contained a non-numeric value");
        }
        result.push_back(value.get<float>());
    }
    if (result.empty()) {
        throw std::runtime_error("embedding response contained an empty embedding");
    }
    return result;
}

std::string extract_chat_text(const json& response) {
    throw_if_error_response(response, "chat completion request");

    if (auto text = try_extract_chat_text(response)) {
        return *text;
    }

    throw std::runtime_error("chat completion response did not contain text content");
}

std::map<std::string, double> parse_classifier_scores(const json& response) {
    throw_if_error_response(response, "classifier request");

    if (auto content = try_extract_chat_text(response)) {
        json parsed = parse_json_text(*content);
        auto scores = parse_scores_payload(parsed);
        if (!scores.empty()) {
            validate_score_range(scores);
            return scores;
        }
        throw std::runtime_error(
            "classifier chat content did not contain label scores");
    }

    auto scores = parse_scores_payload(response);
    if (!scores.empty()) {
        validate_score_range(scores);
        return scores;
    }

    throw std::runtime_error("classifier response did not contain label scores");
}

ClassifierServices make_classifier_services_from_router_calls(
    RouterJsonCall embeddings,
    RouterJsonCall chat_completion,
    EnsureClassifierModelLoaded ensure_loaded) {
    ClassifierServices services;
    auto embeddings_call = std::make_shared<RouterJsonCall>(std::move(embeddings));
    auto chat_completion_call =
        std::make_shared<RouterJsonCall>(std::move(chat_completion));

    services.embed = [embeddings_call,
                      ensure_loaded](const std::string& model,
                                     const std::string& text) -> std::vector<float> {
        if (!*embeddings_call) {
            throw std::runtime_error("Router embeddings call is not configured");
        }
        ensure_model(ensure_loaded, model);
        json request = {
            {"model", model},
            {"input", text},
        };
        return parse_embedding_vector((*embeddings_call)(request));
    };

    services.run_classifier = [chat_completion_call,
                               ensure_loaded](const std::string& model,
                                              const std::string& input)
        -> std::map<std::string, double> {
        if (!*chat_completion_call) {
            throw std::runtime_error("Router chat_completion call is not configured");
        }
        ensure_model(ensure_loaded, model);
        json request = {
            {"model", model},
            {"stream", false},
            {"temperature", 0.0},
            {"messages", json::array({
                {
                    {"role", "system"},
                    {"content", "Classify the user input. Return only JSON mapping label names to numeric scores."},
                },
                {
                    {"role", "user"},
                    {"content", input},
                },
            })},
        };
        return parse_classifier_scores((*chat_completion_call)(request));
    };

    services.chat = [chat_completion_call,
                     ensure_loaded](const std::string& model,
                                    const std::string& prompt,
                                    const std::string& input) -> std::string {
        if (!*chat_completion_call) {
            throw std::runtime_error("Router chat_completion call is not configured");
        }
        ensure_model(ensure_loaded, model);
        json request = {
            {"model", model},
            {"stream", false},
            {"temperature", 0.0},
            {"messages", json::array({
                {{"role", "system"}, {"content", prompt}},
                {{"role", "user"}, {"content", input}},
            })},
        };
        return extract_chat_text((*chat_completion_call)(request));
    };

    return services;
}

} // namespace lemon
