#include "ResultSourceMapping.h"

#include "Pgsql_Result.h"
#include "catalog/PgDatabaseCatalog.h"
#include "catalog/PgConstraintContainer.h"

#include <algorithm>
#include <stdexcept>

namespace {

/// Whether a non-key column of the table is selected more than once.
/// A table that appears twice in a join (through aliases) shows up like that, and libpq can not tell the instances apart.
bool hasAmbiguousColumn(const std::vector<ResultColumnSource> &columns, Oid tableOid, const std::vector<int> &pkAttnums)
{
    for (const auto &col : columns) {
        if (col.tableOid != tableOid || std::find(pkAttnums.begin(), pkAttnums.end(), col.attnum) != pkAttnums.end())
            continue;
        auto count = std::count_if(columns.begin(), columns.end(), [&](const ResultColumnSource &other) {
            return other.tableOid == tableOid && other.attnum == col.attnum;
        });
        if (count > 1)
            return true;
    }
    return false;
}

}

ResultSourceMapping::ResultSourceMapping(std::vector<ResultColumnSource> columns, const std::vector<TablePrimaryKey> &primaryKeys)
    : m_columns(std::move(columns))
{
    for (const auto &col : m_columns) {
        if (col.tableOid == InvalidOid || findTable(col.tableOid))
            continue;

        TableInfo info;
        info.tableOid = col.tableOid;

        auto pk = std::find_if(primaryKeys.begin(), primaryKeys.end(), [&](const TablePrimaryKey &key) {
            return key.tableOid == col.tableOid;
        });
        if (pk != primaryKeys.end() && !pk->attnums.empty()) {
            bool complete = true;
            for (int attnum : pk->attnums) {
                auto found = std::find_if(m_columns.begin(), m_columns.end(), [&](const ResultColumnSource &source) {
                    return source.tableOid == col.tableOid && source.attnum == attnum;
                });
                if (found == m_columns.end()) {
                    complete = false;
                    break;
                }
                info.pkColumns.push_back(static_cast<int>(std::distance(m_columns.begin(), found)));
            }
            info.editable = complete && !hasAmbiguousColumn(m_columns, col.tableOid, pk->attnums);
        }
        if (!info.editable)
            info.pkColumns.clear();

        m_tables.push_back(info);
    }
}

int ResultSourceMapping::columnCount() const
{
    return static_cast<int>(m_columns.size());
}

Oid ResultSourceMapping::tableOid(int col) const
{
    if (col < 0 || col >= columnCount())
        throw std::range_error("column index out of range");
    return m_columns[col].tableOid;
}

int ResultSourceMapping::attnum(int col) const
{
    if (col < 0 || col >= columnCount())
        throw std::range_error("column index out of range");
    return m_columns[col].attnum;
}

bool ResultSourceMapping::isEditable(int col) const
{
    if (col < 0 || col >= columnCount())
        throw std::range_error("column index out of range");

    const auto &source = m_columns[col];
    if (source.attnum <= 0)
        return false;

    const TableInfo *info = findTable(source.tableOid);
    return info && info->editable;
}

std::vector<int> ResultSourceMapping::primaryKeyColumns(Oid tableOid) const
{
    const TableInfo *info = findTable(tableOid);
    return info ? info->pkColumns : std::vector<int>{};
}

const ResultSourceMapping::TableInfo *ResultSourceMapping::findTable(Oid tableOid) const
{
    auto found = std::find_if(m_tables.begin(), m_tables.end(), [tableOid](const TableInfo &info) {
        return info.tableOid == tableOid;
    });
    return found == m_tables.end() ? nullptr : &*found;
}

ResultSourceMapping createResultSourceMapping(const Pgsql::Result &result, const PgDatabaseCatalog &catalog)
{
    std::vector<ResultColumnSource> columns;
    std::vector<TablePrimaryKey> primaryKeys;

    for (int col = 0; col < result.cols(); ++col) {
        ResultColumnSource source;
        source.tableOid = result.ftable(col);
        source.attnum = result.ftableCol(col);
        columns.push_back(source);

        if (source.tableOid == InvalidOid)
            continue;

        bool known = std::any_of(primaryKeys.begin(), primaryKeys.end(), [&](const TablePrimaryKey &key) {
            return key.tableOid == source.tableOid;
        });
        if (known)
            continue;

        TablePrimaryKey pk;
        pk.tableOid = source.tableOid;
        auto constraint = catalog.constraints()->getPrimaryForRelation(source.tableOid);
        if (constraint)
            pk.attnums.assign(constraint->key.begin(), constraint->key.end());
        primaryKeys.push_back(pk);
    }

    return ResultSourceMapping(std::move(columns), primaryKeys);
}