// Copyright 2020 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chromeos/ash/components/local_search_service/inverted_index.h"

#include <numeric>
#include <string>
#include <tuple>
#include <vector>

#include "base/functional/bind.h"
#include "base/functional/callback.h"
#include "base/strings/utf_string_conversions.h"
#include "base/task/task_traits.h"
#include "base/task/thread_pool.h"
#include "chromeos/ash/components/local_search_service/search_utils.h"

namespace ash::local_search_service {

namespace {

// (document-score, posting-of-all-matching-terms).
using ScoreWithPosting = std::pair<double, Posting>;

// Calculates TF-IDF scores for a term
std::vector<TfidfResult> CalculateTfidf(const std::u16string& term,
                                        const DocLength& doc_length,
                                        const Dictionary& dictionary) {
  std::vector<TfidfResult> results;
  // We don't apply weights to idf because the effect is likely small.
  const float idf =
      1.0 + log((1.0 + doc_length.size()) / (1.0 + dictionary.at(term).size()));

  for (const auto& item : dictionary.at(term)) {
    // If a term has a very low content weight in a doc, its effective number of
    // occurrences in the doc should be lower. Strictly speaking, the effective
    // length of the doc should be smaller too. However, for performance
    // reasons, we only apply the weight to the term occurrences but not doc
    // length.
    // TODO(jiameng): this is an expensive operation, we will need to monitor
    // its performance and optimize it.
    const double effective_term_occ = std::accumulate(
        item.second.begin(), item.second.end(), 0.0,
        [](double sum, const WeightedPosition& weighted_position) {
          return sum + weighted_position.weight;
        });
    const float tf = effective_term_occ / doc_length.at(item.first);
    results.push_back({item.first, item.second, tf * idf});
  }
  return results;
}

// Builds TF-IDF cache given the data. Since this function is expensive, it
// should run on a non-blocking thread that is different than the main thread.
TfidfCache BuildTfidf(uint32_t num_docs_from_last_update,
                      const DocLength& doc_length,
                      const Dictionary& dictionary,
                      const TermSet& terms_to_be_updated,
                      const TfidfCache& tfidf_cache) {
  // TODO(crbug.com/40152719): consider moving the helper functions inside the
  // class so that we can use SequenceChecker.
  TfidfCache new_cache(tfidf_cache);
  // If number of documents doesn't change from the last time index was built,
  // we only need to update terms in |terms_to_be_updated|. Otherwise we need
  // to rebuild the index.
  if (num_docs_from_last_update == doc_length.size()) {
    for (const auto& term : terms_to_be_updated) {
      if (dictionary.find(term) != dictionary.end()) {
        new_cache[term] = CalculateTfidf(term, doc_length, dictionary);
      } else {
        new_cache.erase(term);
      }
    }
  } else {
    new_cache.clear();
    for (const auto& item : dictionary) {
      new_cache[item.first] =
          CalculateTfidf(item.first, doc_length, dictionary);
    }
  }
  return new_cache;
}

// Removes a document from document state variables given it's ID. Don't do
// anything if the ID doesn't exist. Return true if the document is removed.
bool RemoveDocumentIfExist(const std::string& document_id,
                           DocLength* doc_length,
                           Dictionary* dictionary,
                           TermSet* terms_to_be_updated) {
  CHECK(doc_length);
  CHECK(dictionary);
  CHECK(terms_to_be_updated);
  bool document_removed = false;
  if (doc_length->find(document_id) == doc_length->end())
    return document_removed;
  doc_length->erase(document_id);
  for (auto it = dictionary->begin(); it != dictionary->end();) {
    if (it->second.find(document_id) != it->second.end()) {
      terms_to_be_updated->insert(it->first);
      it->second.erase(document_id);
      document_removed = true;
    }

    // Removes term from the dictionary if its posting list is empty.
    if (it->second.empty()) {
      it = dictionary->erase(it);
    } else {
      it++;
    }
  }
  return document_removed;
}

// Given list of documents to update and document state variables, returns new
// document state variables and number of deleted documents.
std::pair<DocumentStateVariables, uint32_t> UpdateDocumentStateVariables(
    DocumentToUpdate&& documents_to_update,
    const DocLength& doc_length,
    Dictionary&& dictionary,
    TermSet&& terms_to_be_updated) {
  DocLength new_doc_length(doc_length);
  uint32_t num_deleted = 0u;
  for (const auto& document : documents_to_update) {
    const std::string document_id(document.first);
    bool is_deleted = RemoveDocumentIfExist(document_id, &new_doc_length,
                                            &dictionary, &terms_to_be_updated);

    // Update the document if necessary.
    if (!document.second.empty()) {
      // If document content is not empty, it is being updated but not
      // deleted.
      is_deleted = false;
      for (const auto& token : document.second) {
        dictionary[token.content][document_id] = token.positions;
        new_doc_length[document_id] += token.positions.size();
        terms_to_be_updated.insert(token.content);
      }
    }
    num_deleted += (is_deleted) ? 1 : 0;
  }

  return std::make_pair(
      std::make_tuple(std::move(new_doc_length), std::move(dictionary),
                      std::move(terms_to_be_updated)),
      num_deleted);
}

// Given the index variables, clear all the data.
std::pair<DocumentStateVariables, TfidfCache> ClearData(
    DocumentToUpdate&& documents_to_update,
    const DocLength& doc_length,
    Dictionary&& dictionary,
    TermSet&& terms_to_be_updated,
    TfidfCache&& tfidf_cache) {
  DocLength new_doc_length;
  documents_to_update.clear();
  dictionary.clear();
  terms_to_be_updated.clear();
  tfidf_cache.clear();
  return std::make_pair(
      std::make_tuple(std::move(new_doc_length), std::move(dictionary),
                      std::move(terms_to_be_updated)),
      std::move(tfidf_cache));
}

}  // namespace

InvertedIndex::InvertedIndex() {
  task_runner_ = base::ThreadPool::CreateSequencedTaskRunner(
      {base::TaskPriority::BEST_EFFORT, base::MayBlock(),
       base::TaskShutdownBehavior::CONTINUE_ON_SHUTDOWN});
}
InvertedIndex::~InvertedIndex() = default;

PostingList InvertedIndex::FindTerm(const std::u16string& term) const {
  auto it = dictionary_.find(term);
  if (it != dictionary_.end()) {
    return it->second;
  }

  return {};
}

std::vector<Result> InvertedIndex::FindMatchingDocumentsApproximately(
    const std::unordered_set<std::u16string>& terms,
    double prefix_threshold,
    double block_threshold) const {
  // For each document, its score is the sum of the scores of its terms that
  // match one of more query term. Each term's score is the product of its
  // TF-IDF score and its match relevance score.
  // The map is keyed by the document id.
  std::unordered_map<std::string, ScoreWithPosting> matching_docs;
  for (const auto& kv : tfidf_cache_) {
    const std::u16string& index_term = kv.first;
    const std::vector<TfidfResult>& tfidf_results = kv.second;
    for (const auto& term : terms) {
      const float relevance = RelevanceCoefficient(
          term, index_term, prefix_threshold, block_threshold);
      if (relevance > 0) {
        // If the |index_term| is relevant, all of the enclosing documents will
        // have their ranking scores updated.
        for (const auto& docid_tfidf : tfidf_results) {
          const std::string& docid = std::get<0>(docid_tfidf);
          const Posting& posting = std::get<1>(docid_tfidf);
          const float tfidf = std::get<2>(docid_tfidf);
          auto it = matching_docs.find(docid);
          if (it == matching_docs.end()) {
            it = matching_docs.emplace(docid, ScoreWithPosting(0.0, {})).first;
          }

          auto& score_posting = it->second;
          // TODO(jiameng): add position penalty.
          score_posting.first += tfidf * relevance;
          // Also update matching positions.
          auto& existing_posting = score_posting.second;
          existing_posting.insert(existing_posting.end(), posting.begin(),
                                  posting.end());
        }
        // Break out from inner loop, i.e. no need to check other query terms.
        break;
      }
    }
  }

  std::vector<Result> sorted_matching_docs;
  for (const auto& kv : matching_docs) {
    // We don't need to include weights in the search results.
    std::vector<Position> positions;
    for (const auto& weighted_position : kv.second.second) {
      positions.emplace_back(weighted_position.position);
    }
    sorted_matching_docs.emplace_back(
        Result(kv.first, kv.second.first, positions));
  }
  std::sort(sorted_matching_docs.begin(), sorted_matching_docs.end(),
            CompareResults);
  return sorted_matching_docs;
}

void InvertedIndex::AddDocuments(const DocumentToUpdate& documents,
                                 base::OnceCallback<void()> callback) {
  if (documents.empty())
    return;

  task_runner_->PostTaskAndReplyWithResult(
      FROM_HERE,
      base::BindOnce(&UpdateDocumentStateVariables, documents,
                     std::move(doc_length_), std::move(dictionary_),
                     std::move(terms_to_be_updated_)),
      base::BindOnce(&InvertedIndex::OnAddDocumentsComplete,
                     weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
}

void InvertedIndex::RemoveDocuments(
    const std::vector<std::string>& document_ids,
    base::OnceCallback<void(uint32_t)> callback) {
  DocumentToUpdate documents;
  for (const auto& id : document_ids) {
    documents.push_back({id, std::vector<Token>()});
  }

  task_runner_->PostTaskAndReplyWithResult(
      FROM_HERE,
      base::BindOnce(&UpdateDocumentStateVariables, documents,
                     std::move(doc_length_), std::move(dictionary_),
                     std::move(terms_to_be_updated_)),
      base::BindOnce(&InvertedIndex::OnUpdateDocumentsComplete,
                     weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
}

void InvertedIndex::UpdateDocuments(
    const DocumentToUpdate& documents,
    base::OnceCallback<void(uint32_t)> callback) {
  task_runner_->PostTaskAndReplyWithResult(
      FROM_HERE,
      base::BindOnce(&UpdateDocumentStateVariables, documents,
                     std::move(doc_length_), std::move(dictionary_),
                     std::move(terms_to_be_updated_)),
      base::BindOnce(&InvertedIndex::OnUpdateDocumentsComplete,
                     weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
}

std::vector<TfidfResult> InvertedIndex::GetTfidf(
    const std::u16string& term) const {
  auto it = tfidf_cache_.find(term);
  if (it != tfidf_cache_.end()) {
    return it->second;
  }

  return {};
}

void InvertedIndex::BuildInvertedIndex(base::OnceCallback<void()> callback) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  task_runner_->PostTaskAndReplyWithResult(
      FROM_HERE,
      base::BindOnce(&BuildTfidf, num_docs_from_last_update_, doc_length_,
                     dictionary_, std::move(terms_to_be_updated_),
                     tfidf_cache_),
      base::BindOnce(&InvertedIndex::OnBuildTfidfComplete,
                     weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
}

void InvertedIndex::ClearInvertedIndex(base::OnceCallback<void()> callback) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  task_runner_->PostTaskAndReplyWithResult(
      FROM_HERE,
      base::BindOnce(&ClearData, std::move(documents_to_update_), doc_length_,
                     std::move(dictionary_), std::move(terms_to_be_updated_),
                     std::move(tfidf_cache_)),
      base::BindOnce(&InvertedIndex::OnDataCleared,
                     weak_ptr_factory_.GetWeakPtr(), std::move(callback)));
}

void InvertedIndex::OnBuildTfidfComplete(base::OnceCallback<void()> callback,
                                         TfidfCache&& new_cache) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  num_docs_from_last_update_ = doc_length_.size();
  tfidf_cache_ = std::move(new_cache);

  std::move(callback).Run();
}

void InvertedIndex::OnUpdateDocumentsComplete(
    base::OnceCallback<void(uint32_t)> callback,
    std::pair<DocumentStateVariables, uint32_t>&&
        document_state_variables_and_num_deleted) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  doc_length_ =
      std::move(std::get<0>(document_state_variables_and_num_deleted.first));
  dictionary_ =
      std::move(std::get<1>(document_state_variables_and_num_deleted.first));
  terms_to_be_updated_ =
      std::move(std::get<2>(document_state_variables_and_num_deleted.first));

  BuildInvertedIndex(base::BindOnce(
      [](base::OnceCallback<void(uint32_t)> callback, uint32_t num_deleted) {
        std::move(callback).Run(num_deleted);
      },
      std::move(callback), document_state_variables_and_num_deleted.second));
}

void InvertedIndex::OnAddDocumentsComplete(
    base::OnceCallback<void()> callback,
    std::pair<DocumentStateVariables, uint32_t>&&
        document_state_variables_and_num_deleted) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  DCHECK_EQ(document_state_variables_and_num_deleted.second, 0u);
  doc_length_ =
      std::move(std::get<0>(document_state_variables_and_num_deleted.first));
  dictionary_ =
      std::move(std::get<1>(document_state_variables_and_num_deleted.first));
  terms_to_be_updated_ =
      std::move(std::get<2>(document_state_variables_and_num_deleted.first));

  BuildInvertedIndex(std::move(callback));
}

void InvertedIndex::OnDataCleared(
    base::OnceCallback<void()> callback,
    std::pair<DocumentStateVariables, TfidfCache>&& inverted_index_data) {
  DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
  doc_length_ = std::move(std::get<0>(inverted_index_data.first));
  dictionary_ = std::move(std::get<1>(inverted_index_data.first));
  terms_to_be_updated_ = std::move(std::get<2>(inverted_index_data.first));
  tfidf_cache_ = std::move(inverted_index_data.second);
  num_docs_from_last_update_ = 0;

  std::move(callback).Run();
}

}  // namespace ash::local_search_service