/* 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 https://mozilla.org/MPL/2.0/. */

/**
 * This module defines functions to generate, deduplicate, and filter memories.
 *
 * The primary method in this module is `generateMemories`, which orchestrates the entire pipeline:
 * 1. Generates initial memories from a specified user data user
 * 2. Deduplicates the newly generated memories against all existing memories
 * 3. Filters out memories with sensitive content (i.e. financial, medical, etc.)
 * 4. Returns the final list of memories objects
 *
 * `generateMemories` requires 3 arguments:
 * 1. `engine`: an instance of `openAIEngine` to call the LLM API
 * 2. `sources`: an object mapping user data source types to aggregated records (i.e., {history: [domainItems, titleItems, searchItems]})
 * 3. `existingMemoriesList`: an array of existing memory summary strings to deduplicate against
 *
 * Example Usage:
 * const engine = await openAIEngine.build(MODEL_FEATURES.MEMORIES, DEFAULT_ENGINE_ID, SERVICE_TYPES.MEMORIES);
 * const sources = {history: [domainItems, titleItems, searchItems]};
 * const existingMemoriesList = [...]; // Array of existing memory summary strings; this should be fetched from memory storage
 * const newMemories = await generateMemories(engine, sources, existingMemoriesList);
 *
 */

import { renderPrompt, openAIEngine } from "../Utils.sys.mjs";

import {
  HISTORY,
  CONVERSATION,
  ALL_SOURCES,
  CATEGORIES,
  CATEGORIES_LIST,
  INTENTS,
  INTENTS_LIST,
} from "./MemoriesConstants.sys.mjs";

import {
  initialMemoriesGenerationSystemPrompt,
  initialMemoriesGenerationPrompt,
  memoriesDeduplicationSystemPrompt,
  memoriesDeduplicationPrompt,
  memoriesSensitivityFilterSystemPrompt,
  memoriesSensitivityFilterPrompt,
} from "moz-src:///browser/components/aiwindow/models/prompts/MemoriesPrompts.sys.mjs";

import {
  INITIAL_MEMORIES_SCHEMA,
  MEMORIES_DEDUPLICATION_SCHEMA,
  MEMORIES_NON_SENSITIVE_SCHEMA,
} from "moz-src:///browser/components/aiwindow/models/memories/MemoriesSchemas.sys.mjs";

/**
 * Generates, deduplicates, and filters memories end-to-end
 *
 * This is the main pipeline function.
 *
 * @param {OpenAIEngine} engine                 openAIEngine instance to call LLM API
 * @param {object} sources                      User data source type to aggregrated records (i.e., {history: [domainItems, titleItems, searchItems]})
 * @param {Array<string>} existingMemoriesList  List of existing memory summary strings to deduplicate against
 * @returns {Promise<Array<Map<{
 *  category: string,
 *  intent: string,
 *  memory_summary: string,
 *  score: number,
 * }>>>}                                        Promise resolving the final list of generated, deduplicated, and filtered memory objects
 */
export async function generateMemories(engine, sources, existingMemoriesList) {
  // Step 1: Generate initial memories
  const initialMemories = await generateInitialMemoriesList(engine, sources);
  // If we don't generate any new memories, just return an empty list immediately instead of doing the rest of the steps
  if (!initialMemories || initialMemories.length === 0) {
    return [];
  }

  // Step 2: Deduplicate against existing memories
  const initialMemoriesSummaries = initialMemories.map(
    memory => memory.memory_summary
  );
  const dedupedMemoriesSummaries = await deduplicateMemories(
    engine,
    existingMemoriesList,
    initialMemoriesSummaries
  );
  // If we don't have any deduped memories, no new memories were generated or we ran into an unexpected JSON parse error, so return an empty list
  if (!dedupedMemoriesSummaries || dedupedMemoriesSummaries.length === 0) {
    return [];
  }

  // Step 3: Filter out sensitive memories
  const nonSensitiveMemoriesSummaries = await filterSensitiveMemories(
    engine,
    dedupedMemoriesSummaries
  );

  // Step 4: Map back to full memory objects and return
  return await mapFilteredMemoriesToInitialList(
    initialMemories,
    nonSensitiveMemoriesSummaries
  );
}

/**
 * Formats a list of strings into a prompt-friendly bullet list
 *
 * @param {List<string>} list
 * @returns {string}
 */
export function formatListForPrompt(list) {
  return list.map(item => `- "${item}"`).join("\n");
}

/**
 * Utility function to cleanly get bullet-formatted category and memory lists
 *
 * @param {string} attributeName  "categories" or "intents"
 * @returns {string}              Formatted list string
 */
export function getFormattedMemoryAttributeList(attributeName) {
  if (attributeName === CATEGORIES) {
    return formatListForPrompt(CATEGORIES_LIST);
  } else if (attributeName === INTENTS) {
    return formatListForPrompt(INTENTS_LIST);
  }
  throw new Error(`Unsupported memory attribute name: ${attributeName}`);
}

/**
 * Extracts a JSON as a map from an LLM response (handles markdown-formatted code blocks)
 *
 * @param {any} response  LLM response
 * @param {any} fallback  Fallback value if parsing fails to protect downstream code
 * @returns {Map}         Parsed JSON object
 */
export function parseAndExtractJSON(response, fallback) {
  const rawContent = response?.finalOutput ?? "";
  const markdownMatch = rawContent.match(/```(?:json)?\s*([\s\S]*?)\s*```/i);
  const payload = markdownMatch ? markdownMatch[1] : rawContent;
  try {
    return JSON.parse(payload);
  } catch (e) {
    // If we can't parse a JSON from the LLM response, return a tailored fallback value to prevent downstream code failures
    if (e instanceof SyntaxError) {
      console.warn(
        `Could not parse JSON from LLM response; using fallback (${fallback}): ${e.message}`
      );
      return fallback;
    }
    throw new Error(
      `Unexpected error parsing JSON from LLM response: ${e.message}`
    );
  }
}

/**
 * Renders recent history records into CSV tables for prompt input
 *
 * @param {Array<Array<string>>} domainItems   List of aggregated domain items
 * @param {Array<Array<string>>} titleItems    List of aggregated title items
 * @param {Array<object>} searchItems          List of aggregated search items
 * @returns {Promise<string>}   Promise resolving recent browser history rendered as CSV tables
 */
export async function renderRecentHistoryForPrompt(
  domainItems,
  titleItems,
  searchItems
) {
  let finalCSV = "";

  if (domainItems.length) {
    let domainRecordsTable = ["Domain,Importance Score"];
    for (const domainItem of domainItems) {
      domainRecordsTable.push(domainItem.join(","));
    }
    finalCSV += "# Domains\n" + domainRecordsTable.join("\n") + "\n\n";
  }

  if (titleItems.length) {
    let titleRecordsTable = ["Title,Importance Score"];
    for (const titleItem of titleItems) {
      titleRecordsTable.push(titleItem.join(","));
    }
    finalCSV += "# Titles\n" + titleRecordsTable.join("\n") + "\n\n";
  }

  if (searchItems.length) {
    let searchRecordsTable = ["Search,Importance Score"];
    for (const searchItem of searchItems) {
      searchRecordsTable.push(`${searchItem.q},${searchItem.r}`);
    }
    finalCSV += "# Searches\n" + searchRecordsTable.join("\n");
  }

  return finalCSV.trim();
}

export async function renderRecentConversationForPrompt(conversationMessages) {
  let finalCSV = "";
  if (conversationMessages.length) {
    let conversationRecordsTable = ["Message"];
    for (const message of conversationMessages) {
      conversationRecordsTable.push(`${message.content}`);
    }
    finalCSV += "# Chat History\n" + conversationRecordsTable.join("\n");
  }
  return finalCSV.trim();
}

/**
 * Builds the initial memories generation prompt, pulling profile information based on given source
 *
 * @param {object} sources      User data source type to aggregrated records (i.e., {history: [domainItems, titleItems, searchItems]})
 * @returns {Promise<string>}   Promise resolving the generated memories generation prompt with profile records injected
 */
export async function buildInitialMemoriesGenerationPrompt(sources) {
  if (ALL_SOURCES.intersection(new Set(Object.keys(sources))).size === 0) {
    throw new Error(
      `No valid sources provided to build memories generation prompt: ${Object.keys(sources).join(", ")}`
    );
  }

  let profileRecordsRenderedStr = "";

  // Allow for multiple sources in the future
  if (sources.hasOwnProperty(HISTORY)) {
    const [domainItems, titleItems, searchItems] = sources[HISTORY];
    profileRecordsRenderedStr += await renderRecentHistoryForPrompt(
      domainItems,
      titleItems,
      searchItems
    );
  }
  if (sources.hasOwnProperty(CONVERSATION)) {
    profileRecordsRenderedStr += await renderRecentConversationForPrompt(
      sources[CONVERSATION]
    );
  }

  return await renderPrompt(initialMemoriesGenerationPrompt, {
    categoriesList: getFormattedMemoryAttributeList(CATEGORIES),
    intentsList: getFormattedMemoryAttributeList(INTENTS),
    profileRecordsRenderedStr,
  });
}

/**
 * Builds the memories deduplication prompt
 *
 * @param {Array<string>} existingMemoriesList  List of existing memories
 * @param {Array<string>} newMemoriesList       List of newly generated memories
 * @returns {Promise<string>}                   Promise resolving the generated deduplication prompt with existing and new memories lists injected
 */
export async function buildMemoriesDeduplicationPrompt(
  existingMemoriesList,
  newMemoriesList
) {
  const existingMemoriesListStr = formatListForPrompt(existingMemoriesList);
  const newMemoriesListStr = formatListForPrompt(newMemoriesList);

  return await renderPrompt(memoriesDeduplicationPrompt, {
    existingMemoriesList: existingMemoriesListStr,
    newMemoriesList: newMemoriesListStr,
  });
}

/**
 * Builds the memories sensitivity filter prompt
 *
 * @param {Array<string>} memoriesList  List of memories to filter
 * @returns {Promise<string>}           Promise resolving the generated sensitivity filter prompt with memories list injected
 */
export async function buildMemoriesSensitivityFilterPrompt(memoriesList) {
  const memoriesListStr = formatListForPrompt(memoriesList);

  return await renderPrompt(memoriesSensitivityFilterPrompt, {
    memoriesList: memoriesListStr,
  });
}

/**
 * Sanitizes a single memory object from LLM output, checking required fields and normalizing score
 *
 * @param {*} memory               Raw memory object from LLM
 * @returns {Map<{
 *  category: string|null,
 *  intent: string|null,
 *  memory_summary: string|null,
 *  score: number,
 * }>|null}                         Sanitized memory or null if invalid
 */
function sanitizeMemory(memory) {
  // Shortcut to return nothing if memory is bad
  if (!memory || typeof memory !== "object") {
    return null;
  }

  // Check that the candidate memory object has all the required string fields
  for (const field of ["category", "intent", "memory_summary"]) {
    if (!(field in memory) && typeof memory[field] !== "string") {
      return null;
    }
  }

  // Clamp score to [1,5]; treat missing/invalid as 1
  let score = Number.isFinite(memory.score) ? Math.round(memory.score) : 1;
  if (score < 1) {
    score = 1;
  } else if (score > 5) {
    score = 5;
  }

  return {
    category: memory.category,
    intent: memory.intent,
    memory_summary: memory.memory_summary,
    score,
  };
}

/**
 * Normalizes and validates parsed LLM output into a list of memories to handle LLM output variability
 *
 * @param {*} parsed                JSON-parsed LLM output
 * @returns {Array<Map<{
 *  category: string,
 *  intent: string,
 *  memory_summary: string,
 *  score: number,
 * }>>}                             List of sanitized memories
 */
function normalizeMemoryList(parsed) {
  let list = parsed;
  if (!Array.isArray(list)) {
    // If list isn't an array, check that it's an object with a nested "items" array
    if (list && Array.isArray(list.items)) {
      list = list.items;
    } else if (list && typeof list === "object") {
      // If list isn't an array, check that it's a least a single object, so check that list has memory-like keys
      const looksLikeMemory =
        "category" in list || "intent" in list || "memory_summary" in list;
      if (looksLikeMemory) {
        list = [list];
      }
    }
  }
  if (!Array.isArray(list)) {
    return [];
  }

  return list.map(sanitizeMemory).filter(Boolean);
}

/**
 * Prompts an LLM to generate an initial, unfiltered list of candidate memories from user data
 *
 * @param {openAIEngine} engine     openAIEngine instance to call LLM API
 * @param {object} sources          User data source type to aggregrated records (i.e., {history: [domainItems, titleItems, searchItems]})
 * @returns {Promise<Array<Map<{
 *  category: string,
 *  intent: string,
 *  memory_summary: string,
 *  score: number,
 * }>>>}                            Promise resolving the list of generated memories
 */
export async function generateInitialMemoriesList(engine, sources) {
  const promptText = await buildInitialMemoriesGenerationPrompt(sources);
  const response = await engine.run({
    args: [
      {
        role: "system",
        content: initialMemoriesGenerationSystemPrompt,
      },
      { role: "user", content: promptText },
    ],
    responseFormat: { type: "json_schema", schema: INITIAL_MEMORIES_SCHEMA },
    fxAccountToken: await openAIEngine.getFxAccountToken(),
  });

  const parsed = parseAndExtractJSON(response, []);
  return normalizeMemoryList(parsed);
}

/**
 * Prompts an LLM to deduplicate new memories against existing ones
 *
 * @param {OpenAIEngine} engine                 openAIEngine instance to call LLM API
 * @param {Array<string>} existingMemoriesList  List of existing memory summary strings
 * @param {Array<string>} newMemoriesList       List of new memory summary strings to deduplicate
 * @returns {Promise<Array<string>>}            Promise resolving the final list of deduplicated memory summary strings
 */
export async function deduplicateMemories(
  engine,
  existingMemoriesList,
  newMemoriesList
) {
  const dedupPrompt = await buildMemoriesDeduplicationPrompt(
    existingMemoriesList,
    newMemoriesList
  );

  const response = await engine.run({
    args: [
      {
        role: "system",
        content: memoriesDeduplicationSystemPrompt,
      },
      { role: "user", content: dedupPrompt },
    ],
    responseFormat: {
      type: "json_schema",
      schema: MEMORIES_DEDUPLICATION_SCHEMA,
    },
    fxAccountToken: await openAIEngine.getFxAccountToken(),
  });

  const parsed = parseAndExtractJSON(response, { unique_memories: [] });

  // Able to extract a JSON, so the fallback wasn't used, but the LLM didn't follow the schema
  if (
    parsed.unique_memories === undefined ||
    !Array.isArray(parsed.unique_memories)
  ) {
    return [];
  }

  // Make sure we filter out any invalid main_memory entries before returning
  return parsed.unique_memories
    .filter(
      item =>
        item.main_memory !== undefined && typeof item.main_memory === "string"
    )
    .map(item => item.main_memory);
}

/**
 * Prompts an LLM to filter out sensitive memories from an memories list
 *
 * @param {OpenAIEngine} engine         openAIEngine instance to call LLM API
 * @param {Array<string>} memoriesList  List of memory summary strings to filter
 * @returns {Promise<Array<string>>}    Promise resolving the final list of non-sensitive memory summary strings
 */
export async function filterSensitiveMemories(engine, memoriesList) {
  const sensitivityFilterPrompt =
    await buildMemoriesSensitivityFilterPrompt(memoriesList);
  const response = await engine.run({
    args: [
      {
        role: "system",
        content: memoriesSensitivityFilterSystemPrompt,
      },
      { role: "user", content: sensitivityFilterPrompt },
    ],
    responseFormat: {
      type: "json_schema",
      schema: MEMORIES_NON_SENSITIVE_SCHEMA,
    },
    fxAccountToken: await openAIEngine.getFxAccountToken(),
  });

  const parsed = parseAndExtractJSON(response, { non_sensitive_memories: [] });

  // Able to extract a JSON, so the fallback wasn't used, but the LLM didn't follow the schema
  if (
    parsed.non_sensitive_memories === undefined ||
    !Array.isArray(parsed.non_sensitive_memories)
  ) {
    return [];
  }

  // Make sure we filter out any invalid entries before returning
  return parsed.non_sensitive_memories.filter(item => typeof item === "string");
}

/**
 *
 * @param {Map<string, any>} initialMemories    List of original, unfiltered memory objects
 * @param {Array<string>} filteredMemoriesList  List of deduplicated and sensitivity-filtered memory summary strings
 * @returns {Promise<Map<string, any>>}         Promise resolving the final list of memory objects
 */
export async function mapFilteredMemoriesToInitialList(
  initialMemories,
  filteredMemoriesList
) {
  return initialMemories.filter(memory =>
    filteredMemoriesList.includes(memory.memory_summary)
  );
}