#include "database.hpp"
#include <chrono>
#include <iostream>
#include <syncstream>

namespace {
json nullableString(const pqxx::field& field) {
    return field.is_null() ? json(nullptr) : json(field.as<std::string>());
}
}

Database::Database(const DatabaseConfig& cfg, int workers, bool debug)
    : connStr(
        "host=" + cfg.host +
        " port=" + std::to_string(cfg.port) +
        " dbname=" + cfg.dbname +
        " user=" + cfg.user +
        " password=" + cfg.password),
      debug(debug) {
    logDebug("Creating schema for random_data");
    {
        pqxx::connection probe(connStr);
        pqxx::work txn(probe);
        txn.exec(
            "CREATE TABLE IF NOT EXISTS random_data ("
            "    id SERIAL PRIMARY KEY,"
            "    name VARCHAR(100) NOT NULL,"
            "    text TEXT NOT NULL,"
            "    description TEXT,"
            "    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP"
            ")"
        );
        txn.commit();
    }
    logDebug("Schema is ready; starting " + std::to_string(workers) + " database worker(s)");
    for (int i = 0; i < workers; ++i) {
        threadPool.emplace_back([this, i]() { worker(i); });
    }
}

Database::~Database() {
    stop();
}

void Database::stop() {
    {
        std::lock_guard<std::mutex> lock(mtx);
        stopFlag = true;
    }
    logDebug("Stopping database workers");
    cv.notify_all();
    for (auto& t : threadPool) {
        if (t.joinable()) {
            t.join();
        }
    }
    logDebug("All database workers stopped");
}

void Database::worker(int id) {
    logDebug("Worker " + std::to_string(id) + " connecting to database");
    auto conn = std::make_unique<pqxx::connection>(connStr);
    logDebug("Worker " + std::to_string(id) + " connected to database");
    std::mt19937 rng(std::random_device{}());

    for (;;) {
        Task task;
        {
            std::unique_lock<std::mutex> lock(mtx);
            cv.wait(lock, [&] { return stopFlag || !queue.empty(); });
            if (stopFlag && queue.empty()) {
                logDebug("Worker " + std::to_string(id) + " stopped");
                return;
            }
            task = std::move(queue.front());
            queue.pop();
        }

        json res;
        try {
            if (task.op == Op::RandomRecord) {
                logDebug("Worker " + std::to_string(id) + " fetching one random record");
                res = getRandomRecord(*conn, rng);
            } else {
                logDebug("Worker " + std::to_string(id) + " fetching " + std::to_string(task.count)
                         + " random record(s)");
                res = getRandomRecords(*conn, rng, task.count, task.med, task.max);
            }
        } catch (const std::exception& e) {
            logDebug("Worker " + std::to_string(id) + " database error: " + e.what()
                     + "; reconnecting");
            conn = std::make_unique<pqxx::connection>(connStr);
            logDebug("Worker " + std::to_string(id) + " reconnected to database");
            res = json{{"error", "Database error"}};
        }
        logDebug("Worker " + std::to_string(id) + " completed request"
                 + (res.contains("error") ? ": " + res["error"].get<std::string>() : ""));
        task.cb(std::move(res));
    }
}

void Database::enqueue(Task task) {
    std::size_t queueSize;
    const char* operation = task.op == Op::RandomRecord ? "one random record" : "random records";
    {
        std::lock_guard<std::mutex> lock(mtx);
        queue.push(std::move(task));
        queueSize = queue.size();
    }
    logDebug("Queued request for " + std::string(operation)
             + "; pending requests: " + std::to_string(queueSize));
    cv.notify_one();
}

void Database::getRandomRecord(Callback cb) {
    enqueue(Task{Op::RandomRecord, 0, 0, 0, std::move(cb)});
}

void Database::getRandomRecords(int count, int med, int max, Callback cb) {
    enqueue(Task{Op::RandomRecords, count, med, max, std::move(cb)});
}

json Database::getRandomRecord(pqxx::connection& conn, std::mt19937& rng) {
    int totalCount = getTotalCount(conn);
    logDebug("random_data contains " + std::to_string(totalCount) + " record(s)");

    if (totalCount == 0) {
        return json{
            {"error", "No records found in database"}
        };
    }

    std::uniform_int_distribution<int> idDist(1, totalCount);
    int randomId = idDist(rng);
    logDebug("Selected random record ID " + std::to_string(randomId));

    pqxx::work txn(conn);

    auto result = txn.exec(
        "SELECT id, name, text, description, created_at "
        "FROM random_data "
        "WHERE id = $1",
        pqxx::params{randomId}
    );

    if (result.empty()) {
        logDebug("Selected record no longer exists; falling back to the first available record");
        result = txn.exec("SELECT id, name, text, description, created_at FROM random_data LIMIT 1");
    }

    if (result.empty()) {
        txn.abort();
        return json{
            {"error", "No records found in database"}
        };
    }

    const auto& row = result[0];

    json response = {
        {"id", row[0].as<int>()},
        {"name", row[1].as<std::string>()},
        {"text", row[2].as<std::string>()},
        {"description", nullableString(row[3])},
        {"created_at", row[4].as<std::string>()},
        {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
    };

    txn.commit();
    return response;
}

json Database::getRandomRecordsFast(pqxx::connection& conn, const int count) {
    logDebug("Using recursive CTE strategy for " + std::to_string(count) + " record(s)");
    pqxx::work txn(conn);

    std::string query =
        "WITH RECURSIVE random_selection AS ("
        "    WITH bounds AS ("
        "        SELECT"
        "            MIN(t.id) AS min_id,"
        "            MAX(t.id) AS max_id"
        "        FROM random_data AS t"
        "    ),"
        "    first_row AS ("
        "        SELECT"
        "            t.id,"
        "            b.min_id,"
        "            b.max_id,"
        "            ARRAY[t.id] AS selected_ids,"
        "            1 AS depth"
        "        FROM random_data AS t, bounds AS b"
        "        WHERE"
        "            t.id >= b.min_id + ((b.max_id - b.min_id) * random())::INT"
        "        ORDER BY t.id"
        "        LIMIT 1"
        "    ),"
        "    recursive_rows AS ("
        "        SELECT"
        "            t.id,"
        "            r.min_id,"
        "            r.max_id,"
        "            r.selected_ids || t.id AS selected_ids,"
        "            r.depth + 1 AS depth"
        "        FROM random_data AS t, first_row AS r"
        "        WHERE"
        "            t.id >= r.min_id + ((r.max_id - r.min_id) * random())::INT"
        "            AND NOT (t.id = ANY(r.selected_ids))"
        "            AND r.depth + 1 < " + std::to_string(count + 1) +
        "        ORDER BY t.id"
        "        LIMIT 1"
        "    )"
        "    SELECT * FROM first_row"
        "    UNION ALL"
        "    SELECT * FROM recursive_rows"
        ")"
        "SELECT"
        "    rd.id,"
        "    rd.name,"
        "    rd.text,"
        "    rd.description,"
        "    rd.created_at"
        "FROM random_selection AS rs"
        "JOIN random_data AS rd ON rd.id = rs.id"
        "ORDER BY rs.depth";

    auto result = txn.exec(query);

    if (result.empty()) {
        return json{
            {"error", "No records found in database"},
            {"count", 0}
        };
    }

    json response = json::array();
    for (const auto& row : result) {
        response.push_back({
            {"id", row[0].as<int>()},
            {"name", row[1].as<std::string>()},
            {"text", row[2].as<std::string>()},
            {"description", nullableString(row[3])},
            {"created_at", row[4].as<std::string>()}
        });
    }

    txn.commit();
    logDebug("Recursive CTE returned " + std::to_string(response.size()) + " record(s)");

    return {
        {"data", response},
        {"count", static_cast<int>(response.size())},
        {"strategy", "fast_recursive_cte"},
        {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
    };
}

json Database::getRandomRecordsSlow(pqxx::connection& conn, std::mt19937& rng, const int count, const int totalCount) {
    logDebug("Using individual-ID strategy for " + std::to_string(count) + " of "
             + std::to_string(totalCount) + " record(s)");
    std::set<int> ids;
    {
        std::uniform_int_distribution<int> idDist(1, totalCount);
        while (ids.size() < static_cast<size_t>(count)) {
            ids.insert(idDist(rng));
        }
    }

    pqxx::work txn(conn);

    if (ids.empty()) {
        return json{
            {"data", json::array()},
            {"count", 0},
            {"strategy", "slow"}
        };
    }

    std::string query = getQueryRandomIds(static_cast<int>(ids.size()));

    pqxx::params param_list;
    for (int id : ids) {
        param_list.append(id);
    }

    auto result = txn.exec(query, param_list);

    json response = json::array();
    for (const auto& row : result) {
        response.push_back({
            {"id", row[0].as<int>()},
            {"name", row[1].as<std::string>()},
            {"text", row[2].as<std::string>()},
            {"description", nullableString(row[3])},
            {"created_at", row[4].as<std::string>()}
        });
    }

    txn.commit();
    logDebug("Individual-ID query returned " + std::to_string(response.size()) + " record(s)");

    return {
        {"data", response},
        {"count", static_cast<int>(response.size())},
        {"strategy", "slow"},
        {"timestamp", std::chrono::system_clock::now().time_since_epoch().count()}
    };
}

json Database::getRandomRecords(pqxx::connection& conn, std::mt19937& rng, int count, int med, int max) {
    if (count <= 0) {
        return json{
            {"error", "Count must be positive"},
            {"count", 0}
        };
    }

    int totalCount = getTotalCount(conn);
    logDebug("Requested " + std::to_string(count) + " random record(s); "
             + std::to_string(totalCount) + " available");

    if (totalCount == 0) {
        return json{
            {"error", "No records found in database"},
            {"count", 0}
        };
    }

    if (count > totalCount) {
        logDebug("Reducing requested count to available record count");
        count = totalCount;
    }

    if (count > max) {
        logDebug("Reducing requested count to configured maximum of " + std::to_string(max));
        count = max;
    }

    json response;
    if (count <= med) {
        response = getRandomRecordsSlow(conn, rng, count, totalCount);
    } else {
        response = getRandomRecordsFast(conn, count);
    }
    if (response.contains("error")) {
        return response;
    }

    response["requested_count"] = count;
    response["total_available"] = totalCount;

    return response;
}

int Database::getTotalCount(pqxx::connection& conn) {
    pqxx::work txn(conn);
    auto result = txn.exec("SELECT COUNT(*) FROM random_data");
    int count = result[0][0].as<int>();
    txn.commit();
    return count;
}

std::string Database::getQueryRandomIds(const int count) {
    std::string query =
        "SELECT id, name, text, description, created_at "
        "FROM random_data WHERE id IN (";

    for (size_t i = 0; i < count; ++i) {
        if (i > 0) query += ",";
        query += "$" + std::to_string(i + 1);
    }
    query += ") ORDER BY id";
    return query;
}

void Database::logDebug(const std::string& message) const {
    if (debug) {
        std::osyncstream(std::cout) << "[debug][database] " << message << '\n';
    }
}