// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

/**
 * Provides semantic search over Firefox preferences navigation entries to
 * answer user queries about where to find Smart Window settings, Memories, etc.
 *
 * Navigation data is loaded from PreferencesNavMap.sys.mjs, which is
 * auto-generated by browser/components/aiwindow/models/scripts/generate_prefs_nav_map.py.
 * Run that script and commit the result whenever preferences.js or FTL files change.
 */

const lazy = {};
ChromeUtils.defineESModuleGetters(lazy, {
  cosSim: "chrome://global/content/ml/NLPUtils.sys.mjs",
  EmbeddingsGenerator: "chrome://global/content/ml/EmbeddingsGenerator.sys.mjs",
  PREFERENCES_NAV_MAP:
    "moz-src:///browser/components/aiwindow/models/PreferencesNavMap.sys.mjs",
});

export const NAV_TOP_K = 3;
export const NAV_SIMILARITY_THRESHOLD = 0.3;

/**
 * Converts the flat nav map into an array of entries with embedding text.
 *
 * @returns {Array<{url, label, breadcrumb, description, embeddingText}>}
 */
function buildNavEntries() {
  return Object.entries(lazy.PREFERENCES_NAV_MAP).map(([url, info]) => ({
    url,
    label: info.label,
    breadcrumb: info.breadcrumb,
    description: info.description,
    embeddingText: [info.breadcrumb, info.description]
      .filter(Boolean)
      .join(". ")
      .toLowerCase(),
  }));
}

/**
 * Provides semantic search over Firefox preferences navigation
 * entries using embeddings to find settings pages relevant to a query.
 */
class NavigationInfoImpl {
  #embeddingsGenerator = null;
  #navEntries = null;
  #navEmbeddings = null;

  /**
   * Returns navigation entries semantically similar to the query,
   * ranked by cosine similarity.
   *
   * @param {string} query
   * @param {number} [topK=3]
   * @param {number} [similarityThreshold=NAV_SIMILARITY_THRESHOLD]
   * @returns {Promise<Array<{url, label, breadcrumb, description, similarity}>>}
   */
  async getRelevantNavigation(
    query,
    topK = NAV_TOP_K,
    similarityThreshold = NAV_SIMILARITY_THRESHOLD
  ) {
    if (!this.#navEntries) {
      this.#navEntries = buildNavEntries();
    }

    if (!this.#embeddingsGenerator) {
      this.#embeddingsGenerator = new lazy.EmbeddingsGenerator({
        backend: "onnx-native",
        embeddingSize: 384,
      });
    }

    if (!this.#navEmbeddings) {
      const result = await this.#embeddingsGenerator.embedMany(
        this.#navEntries.map(e => e.embeddingText)
      );
      this.#navEmbeddings = result.output ?? result;
    }

    const queryResult = await this.#embeddingsGenerator.embed(
      query.toLowerCase()
    );
    let queryEmbedding = queryResult.output ?? queryResult;
    if (Array.isArray(queryEmbedding) && queryEmbedding.length === 1) {
      queryEmbedding = queryEmbedding[0];
    }

    return this.#navEmbeddings
      .map((emb, idx) => ({
        ...this.#navEntries[idx],
        similarity: lazy.cosSim(queryEmbedding, emb),
      }))
      .filter(e => e.similarity >= similarityThreshold)
      .sort((a, b) => b.similarity - a.similarity)
      .slice(0, topK)
      .map(({ embeddingText: _omit, ...rest }) => rest);
  }
}

export const SmartWindowNavigationInfo = new NavigationInfoImpl();