/** @file
 * @brief Sharded database backend
 */
/* Copyright (C) 2017,2019,2020,2022,2024,2026 Olly Betts
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, see
 * <https://www.gnu.org/licenses/>.
 */

#include <config.h>

#include "multi_database.h"

#include "backends/backends.h"
#include "backends/multi.h"
#include "clamp_cast.h"
#include "expand/ortermlist.h"
#include "expand/termlistmerger.h"
#include "multi_alltermslist.h"
#include "multi_postlist.h"
#include "multi_termlist.h"
#include "multi_valuelist.h"
#include "negate_unsigned.h"

#include <memory>
#include <string_view>

using namespace std;

MultiDatabase::size_type
MultiDatabase::size() const
{
    return clamp_cast<MultiDatabase::size_type>(shards.size());
}

bool
MultiDatabase::reopen()
{
    bool result = false;
    for (auto&& shard : shards) {
        if (shard->reopen()) {
            result = true;
        }
    }
    return result;
}

void
MultiDatabase::close()
{
    for (auto&& shard : shards) {
        shard->close();
    }
}

PostList*
MultiDatabase::open_post_list(string_view term) const
{
    PostList** postlists = new PostList*[shards.size()];
    size_t count = 0;
    try {
        for (auto&& shard : shards) {
            postlists[count] = shard->open_post_list(term);
            ++count;
        }
        return new MultiPostList(count, postlists);
    } catch (...) {
        while (count)
            delete postlists[--count];
        delete [] postlists;
        throw;
    }
}

LeafPostList*
MultiDatabase::open_leaf_post_list(string_view, bool) const
{
    // This should never get called.
    Assert(false);
    return NULL;
}

TermList*
MultiDatabase::open_term_list(Xapian::docid did) const
{
    return new MultiTermList(this, MultiDatabase::open_term_list_direct(did));
}

TermList*
MultiDatabase::open_term_list_direct(Xapian::docid did) const
{
    Xapian::doccount n_shards = shards.size();
    auto shard_index = shard_number(did, n_shards);
    auto shard = shards[shard_index];
    Xapian::docid shard_did = shard_docid(did, n_shards);
    TermList* res = shard->open_term_list(shard_did);
    res->shard_index = shard_index;
    return res;
}

TermList*
MultiDatabase::open_allterms(string_view prefix) const
{
    size_t count = 0;
    TermList** termlists = new TermList*[shards.size()];
    try {
        for (auto&& shard : shards) {
            termlists[count] = shard->open_allterms(prefix);
            ++count;
        }
        return new MultiAllTermsList(count, termlists);
    } catch (...) {
        while (count)
            delete termlists[--count];
        delete [] termlists;
        throw;
    }
}

bool
MultiDatabase::has_positions() const
{
    for (auto&& shard : shards) {
        if (shard->has_positions()) {
            return true;
        }
    }
    return false;
}

PositionList*
MultiDatabase::open_position_list(Xapian::docid did, string_view term) const
{
    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    auto shard_did = shard_docid(did, n_shards);
    return shard->open_position_list(shard_did, term);
}

Xapian::doccount
MultiDatabase::get_doccount() const
{
    Xapian::doccount result = 0;
    for (auto&& shard : shards) {
        auto old_result = result;
        result += shard->get_doccount();
        if (result < old_result)
            throw Xapian::DatabaseError("doccount overflowed!");
    }
    return result;
}

Xapian::docid
MultiDatabase::get_lastdocid() const
{
    Xapian::docid result = 0;
    Xapian::doccount n_shards = shards.size();
    for (Xapian::doccount shard = 0; shard != n_shards; ++shard) {
        Xapian::docid shard_lastdocid = shards[shard]->get_lastdocid();
        if (shard_lastdocid == 0) {
            // This shard is empty, so doesn't influence lastdocid for the
            // combined database.
            continue;
        }
        result = max(result, unshard(shard_lastdocid, shard, n_shards));
    }
    return result;
}

Xapian::totallength
MultiDatabase::get_total_length() const
{
    Xapian::totallength result = 0;
    for (auto&& shard : shards) {
        auto old_result = result;
        result += shard->get_total_length();
        if (result < old_result)
            throw Xapian::DatabaseError("Total document length overflowed!");
    }
    return result;
}

void
MultiDatabase::get_freqs(string_view term,
                         Xapian::doccount* tf_ptr,
                         Xapian::termcount* cf_ptr) const
{
    Assert(!term.empty());

    Xapian::doccount shard_tf;
    Xapian::doccount* shard_tf_ptr = tf_ptr ? &shard_tf : NULL;
    Xapian::doccount total_tf = 0;

    Xapian::termcount shard_cf;
    Xapian::termcount* shard_cf_ptr = cf_ptr ? &shard_cf : NULL;
    Xapian::termcount total_cf = 0;

    for (auto&& shard : shards) {
        shard->get_freqs(term, shard_tf_ptr, shard_cf_ptr);
        if (shard_tf_ptr) {
            auto old_tf = total_tf;
            total_tf += *shard_tf_ptr;
            if (total_tf < old_tf)
                throw Xapian::DatabaseError("termfreq overflowed!");
        }
        if (shard_cf_ptr) {
            auto old_cf = total_cf;
            total_cf += *shard_cf_ptr;
            if (total_cf < old_cf)
                throw Xapian::DatabaseError("Collection freq overflowed!");
        }
    }
    if (tf_ptr) {
        *tf_ptr = total_tf;
    }
    if (cf_ptr) {
        *cf_ptr = total_cf;
    }
}

Xapian::doccount
MultiDatabase::get_value_freq(Xapian::valueno slot) const
{
    Xapian::termcount result = 0;
    for (auto&& shard : shards) {
        auto old_result = result;
        result += shard->get_value_freq(slot);
        if (result < old_result)
            throw Xapian::DatabaseError("Value freq overflowed!");
    }
    return result;
}

string
MultiDatabase::get_value_lower_bound(Xapian::valueno slot) const
{
    string result;
    for (auto&& shard : shards) {
        string shard_result = shard->get_value_lower_bound(slot);
        if (shard_result.empty())
            continue;
        if (result.empty() || shard_result < result)
            result = std::move(shard_result);
    }
    return result;
}

string
MultiDatabase::get_value_upper_bound(Xapian::valueno slot) const
{
    string result;
    for (auto&& shard : shards) {
        string shard_result = shard->get_value_upper_bound(slot);
        if (shard_result > result)
            result = std::move(shard_result);
    }
    return result;
}

Xapian::termcount
MultiDatabase::get_doclength_lower_bound() const
{
    // We want the smallest answer from amongst the shards, except that 0 means
    // that all documents have length 0 (including the special case of there
    // being no documents), so any non-zero answer should "beat" 0.  To achieve
    // this we find the *maximum* after negating each of the values (which
    // since Xapian::termcount is an unsigned type leaves 0 alone but flips the
    // order of all other values), then negate the answer again at the end.
    static_assert(std::is_unsigned_v<Xapian::termcount>,
                  "Unsigned type required");
    Xapian::termcount result = 0;
    for (auto&& shard : shards) {
        Xapian::termcount shard_result = shard->get_doclength_lower_bound();
        result = max(result, negate_unsigned(shard_result));
    }
    return negate_unsigned(result);
}

Xapian::termcount
MultiDatabase::get_doclength_upper_bound() const
{
    Xapian::termcount result = 0;
    for (auto&& shard : shards) {
        result = max(result, shard->get_doclength_upper_bound());
    }
    return result;
}

Xapian::termcount
MultiDatabase::get_wdf_upper_bound(string_view term) const
{
    Assert(!term.empty());

    Xapian::termcount result = 0;
    for (auto&& shard : shards) {
        result = max(result, shard->get_wdf_upper_bound(term));
    }
    return result;
}

Xapian::termcount
MultiDatabase::get_unique_terms_lower_bound() const
{
    // We want the smallest answer from amongst the shards, except that 0 means
    // that all documents have no unique terms (including the special case of
    // there being no documents), so any non-zero answer should "beat" 0.  To
    // achieve this we find the *maximum* after negating each of the values
    // (which since Xapian::termcount is an unsigned type leaves 0 alone but
    // flips the order of all other values), then negate the answer again at
    // the end.
    static_assert(std::is_unsigned_v<Xapian::termcount>,
                  "Unsigned type required");
    Xapian::termcount result = 0;
    for (auto&& shard : shards) {
        Xapian::termcount shard_result = shard->get_unique_terms_lower_bound();
        result = max(result, negate_unsigned(shard_result));
    }
    return negate_unsigned(result);
}

Xapian::termcount
MultiDatabase::get_unique_terms_upper_bound() const
{
    Xapian::termcount result = 0;
    for (auto&& shard : shards) {
        result = max(result, shard->get_unique_terms_upper_bound());
    }
    return result;
}

ValueList*
MultiDatabase::open_value_list(Xapian::valueno slot) const
{
    SubValueList** valuelists = new SubValueList*[shards.size()];
    unsigned count = 0;
    try {
        for (auto&& shard : shards) {
            ValueList* vl = shard->open_value_list(slot);
            valuelists[count] = new SubValueList(vl, count);
            ++count;
        }
        return new MultiValueList(count, valuelists, slot);
    } catch (...) {
        while (count)
            delete valuelists[--count];
        delete [] valuelists;
        throw;
    }
}

Xapian::termcount
MultiDatabase::get_doclength(Xapian::docid did) const
{
    Assert(did != 0);

    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    auto shard_did = shard_docid(did, n_shards);
    return shard->get_doclength(shard_did);
}

Xapian::termcount
MultiDatabase::get_unique_terms(Xapian::docid did) const
{
    Assert(did != 0);

    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    auto shard_did = shard_docid(did, n_shards);
    return shard->get_unique_terms(shard_did);
}

Xapian::termcount
MultiDatabase::get_wdfdocmax(Xapian::docid did) const
{
    Assert(did != 0);

    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    auto shard_did = shard_docid(did, n_shards);
    return shard->get_wdfdocmax(shard_did);
}

Xapian::Document::Internal*
MultiDatabase::open_document(Xapian::docid did, bool lazy) const
{
    Assert(did != 0);

    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    auto shard_did = shard_docid(did, n_shards);
    return shard->open_document(shard_did, lazy);
}

bool
MultiDatabase::term_exists(string_view term) const
{
    for (auto&& shard : shards) {
        if (shard->term_exists(term))
            return true;
    }
    return false;
}

void
MultiDatabase::keep_alive()
{
    for (auto&& shard : shards) {
        shard->keep_alive();
    }
}

TermList*
MultiDatabase::open_spelling_termlist(string_view word) const
{
    vector<TermList*> termlists;
    termlists.reserve(shards.size());

    try {
        for (auto&& shard : shards) {
            TermList* termlist = shard->open_spelling_termlist(word);
            if (!termlist)
                continue;
            termlists.push_back(termlist);
        }

        return make_termlist_merger(termlists);
    } catch (...) {
        for (auto&& termlist : termlists)
            delete termlist;
        throw;
    }
}

TermList*
MultiDatabase::open_spelling_wordlist() const
{
    vector<TermList*> termlists;
    termlists.reserve(shards.size());

    try {
        for (auto&& shard : shards) {
            TermList* termlist = shard->open_spelling_wordlist();
            if (!termlist)
                continue;
            termlists.push_back(termlist);
        }

        return make_termlist_merger<FreqAdderOrTermList>(termlists);
    } catch (...) {
        for (auto&& termlist : termlists)
            delete termlist;
        throw;
    }
}

Xapian::doccount
MultiDatabase::get_spelling_frequency(string_view word) const
{
    Xapian::doccount result = 0;
    for (auto&& shard : shards) {
        auto old_result = result;
        result += shard->get_spelling_frequency(word);
        if (result < old_result)
            throw Xapian::DatabaseError("Spelling frequency overflowed!");
    }
    return result;
}

TermList*
MultiDatabase::open_synonym_termlist(string_view term) const
{
    vector<TermList*> termlists;
    termlists.reserve(shards.size());

    try {
        for (auto&& shard : shards) {
            TermList* termlist = shard->open_synonym_termlist(term);
            if (!termlist)
                continue;
            termlists.push_back(termlist);
        }

        return make_termlist_merger(termlists);
    } catch (...) {
        for (auto&& termlist : termlists)
            delete termlist;
        throw;
    }
}

TermList*
MultiDatabase::open_synonym_keylist(string_view prefix) const
{
    vector<TermList*> termlists;
    termlists.reserve(shards.size());

    try {
        for (auto&& shard : shards) {
            TermList* termlist = shard->open_synonym_keylist(prefix);
            if (!termlist)
                continue;
            termlists.push_back(termlist);
        }

        return make_termlist_merger(termlists);
    } catch (...) {
        for (auto&& termlist : termlists)
            delete termlist;
        throw;
    }
}

string
MultiDatabase::get_metadata(string_view key) const
{
    return shards[0]->get_metadata(key);
}

TermList*
MultiDatabase::open_metadata_keylist(string_view prefix) const
{
    return shards[0]->open_metadata_keylist(prefix);
}

string
MultiDatabase::get_uuid() const
{
    string uuid;
    for (auto&& shard : shards) {
        const string& sub_uuid = shard->get_uuid();
        // If any of the sub-databases have no uuid, we can't make a uuid for
        // the combined database.
        if (sub_uuid.empty())
            return sub_uuid;
        if (!uuid.empty())
            uuid += ':';
        uuid += sub_uuid;
    }
    return uuid;
}

bool
MultiDatabase::locked() const
{
    for (auto&& shard : shards) {
        if (shard->locked()) {
            return true;
        }
    }
    return false;
}

void
MultiDatabase::write_changesets_to_fd(int,
                                      std::string_view,
                                      bool,
                                      Xapian::ReplicationInfo*)
{
    throw Xapian::InvalidOperationError("write_changesets_to_fd() with "
                                        "more than one subdatabase");
}

Xapian::rev
MultiDatabase::get_revision() const
{
    throw Xapian::InvalidOperationError("Database::get_revision() with "
                                        "more than one subdatabase");
}

void
MultiDatabase::invalidate_doc_object(Xapian::Document::Internal*) const
{
    // This method should only be called on a single shard.
    Assert(false);
}

int
MultiDatabase::get_backend_info(string*) const
{
    // This method should only be called on a single shard.
    Assert(false);
    return BACKEND_UNKNOWN;
}

void
MultiDatabase::commit()
{
    for (auto&& shard : shards) {
        shard->commit();
    }
}

void
MultiDatabase::cancel()
{
    for (auto&& shard : shards) {
        shard->cancel();
    }
}

void
MultiDatabase::begin_transaction(bool flushed)
{
    for (auto&& shard : shards) {
        shard->begin_transaction(flushed);
    }
}

void
MultiDatabase::end_transaction(bool do_commit)
{
    for (auto&& shard : shards) {
        shard->end_transaction(do_commit);
    }
}

Xapian::docid
MultiDatabase::add_document(const Xapian::Document& doc)
{
    // With a single shard, add_document() uses docid (get_lastdocid() + 1)
    // which seems a sensible invariant to preserve with multiple shards.
    Xapian::docid did = UNSIGNED_OVERFLOW_OK(get_lastdocid() + 1);
    if (rare(did == 0)) {
        throw Xapian::DatabaseError("Run out of docids - you'll have to use "
                                    "copydatabase to eliminate any gaps "
                                    "before you can add more documents");
    }

    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    shard->replace_document(shard_docid(did, n_shards), doc);
    return did;
}

void
MultiDatabase::delete_document(Xapian::docid did)
{
    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    shard->delete_document(shard_docid(did, n_shards));
}

void
MultiDatabase::delete_document(string_view term)
{
    for (auto&& shard : shards) {
        shard->delete_document(term);
    }
}

void
MultiDatabase::replace_document(Xapian::docid did, const Xapian::Document& doc)
{
    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    shard->replace_document(shard_docid(did, n_shards), doc);
}

Xapian::docid
MultiDatabase::replace_document(string_view term, const Xapian::Document& doc)
{
    auto n_shards = shards.size();
    unique_ptr<PostList> pl(open_post_list(term));
    if (!pl || (pl->next(), pl->at_end())) {
        // unique_term not in the database, so this is just an add_document().
        // Calculate which shard the next never used docid maps to.
        Xapian::docid did = UNSIGNED_OVERFLOW_OK(get_lastdocid() + 1);
        if (rare(did == 0)) {
            throw Xapian::DatabaseError("Run out of docids - you'll have to "
                                        "use copydatabase to eliminate any "
                                        "gaps before you can add more "
                                        "documents");
        }
        auto shard = shards[shard_number(did, n_shards)];
        return shard->add_document(doc);
    }

    Xapian::docid result = pl->get_docid();
    auto replacing_shard = shards[shard_number(result, n_shards)];
    replacing_shard->replace_document(shard_docid(result, n_shards), doc);

    // Delete any other occurrences of the unique term.
    while (pl->next(), !pl->at_end()) {
        Xapian::docid did = pl->get_docid();
        auto shard = shards[shard_number(did, n_shards)];
        shard->delete_document(shard_docid(did, n_shards));
    }

    return result;
}

void
MultiDatabase::request_document(Xapian::docid did) const
{
    Assert(did != 0);

    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    auto shard_did = shard_docid(did, n_shards);
    shard->request_document(shard_did);
}

void
MultiDatabase::add_spelling(string_view word,
                            Xapian::termcount freqinc) const
{
    shards[0]->add_spelling(word, freqinc);
}

Xapian::termcount
MultiDatabase::remove_spelling(string_view word,
                               Xapian::termcount freqdec) const
{
    for (auto&& shard : shards) {
        freqdec = shard->remove_spelling(word, freqdec);
        if (freqdec == 0)
            break;
    }
    return freqdec;
}

void
MultiDatabase::add_synonym(string_view term,
                           string_view synonym) const
{
    shards[0]->add_synonym(term, synonym);
}

void
MultiDatabase::remove_synonym(string_view term,
                              string_view synonym) const
{
    for (auto&& shard : shards) {
        shard->remove_synonym(term, synonym);
    }
}

void
MultiDatabase::clear_synonyms(string_view term) const
{
    for (auto&& shard : shards) {
        shard->clear_synonyms(term);
    }
}

void
MultiDatabase::set_metadata(string_view key, string_view value)
{
    shards[0]->set_metadata(key, value);
}

string
MultiDatabase::reconstruct_text(Xapian::docid did,
                                size_t length,
                                string_view prefix,
                                Xapian::termpos start_pos,
                                Xapian::termpos end_pos) const
{
    Assert(did != 0);

    auto n_shards = shards.size();
    auto shard = shards[shard_number(did, n_shards)];
    auto shard_did = shard_docid(did, n_shards);
    return shard->reconstruct_text(shard_did, length, prefix,
                                   start_pos, end_pos);
}

string
MultiDatabase::get_description() const
{
    string desc;
    for (auto&& shard : shards) {
        if (!desc.empty()) {
            desc += ", ";
        }
        desc += shard->get_description();
    }
    desc += ')';
    return desc;
}
