/* 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/. */

do_get_profile();

const { ChatStore, ChatMessage, MESSAGE_ROLE } = ChromeUtils.importESModule(
  "moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs"
);
const {
  getRecentHistory,
  generateProfileInputs,
  aggregateSessions,
  topkAggregates,
} = ChromeUtils.importESModule(
  "moz-src:///browser/components/aiwindow/models/memories/MemoriesHistorySource.sys.mjs"
);
const { getRecentChats } = ChromeUtils.importESModule(
  "moz-src:///browser/components/aiwindow/models/memories/MemoriesChatSource.sys.mjs"
);
const { DEFAULT_ENGINE_ID, MODEL_FEATURES, openAIEngine, SERVICE_TYPES } =
  ChromeUtils.importESModule(
    "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs"
  );
const { sinon } = ChromeUtils.importESModule(
  "resource://testing-common/Sinon.sys.mjs"
);

const { CATEGORIES, INTENTS } = ChromeUtils.importESModule(
  "moz-src:///browser/components/aiwindow/models/memories/MemoriesConstants.sys.mjs"
);

const {
  formatListForPrompt,
  getFormattedMemoryAttributeList,
  renderRecentHistoryForPrompt,
  renderRecentConversationForPrompt,
  mapFilteredMemoriesToInitialList,
  buildInitialMemoriesGenerationPrompt,
  buildMemoriesDeduplicationPrompt,
  buildMemoriesSensitivityFilterPrompt,
  generateInitialMemoriesList,
  deduplicateMemories,
  filterSensitiveMemories,
} = ChromeUtils.importESModule(
  "moz-src:///browser/components/aiwindow/models/memories/Memories.sys.mjs"
);

/**
 * Constants for preference keys and test values
 */
const PREF_API_KEY = "browser.aiwindow.apiKey";
const PREF_ENDPOINT = "browser.aiwindow.endpoint";
const PREF_MODEL = "browser.aiwindow.model";

const API_KEY = "fake-key";
const ENDPOINT = "https://api.fake-endpoint.com/v1";
const MODEL = "fake-model";

const EXISTING_MEMORIES = [
  "Loves outdoor activities",
  "Enjoys cooking recipes",
  "Like sci-fi media",
];
const NEW_MEMORIES = [
  "Loves hiking and camping",
  "Reads science fiction novels",
  "Likes both dogs and cats",
  "Likes risky stock bets",
];

add_setup(async function () {
  // Setup prefs used across multiple tests
  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
  Services.prefs.setStringPref(PREF_MODEL, MODEL);

  // Clear prefs after testing
  registerCleanupFunction(() => {
    for (let pref of [PREF_API_KEY, PREF_ENDPOINT, PREF_MODEL]) {
      if (Services.prefs.prefHasUserValue(pref)) {
        Services.prefs.clearUserPref(pref);
      }
    }
  });
});

/**
 * Builds fake browsing history data for testing
 */
async function buildFakeBrowserHistory() {
  const now = Date.now();

  const seeded = [
    {
      url: "https://www.google.com/search?q=firefox+history",
      title: "Google Search: firefox history",
      visits: [{ date: new Date(now - 5 * 60 * 1000) }],
    },
    {
      url: "https://news.ycombinator.com/",
      title: "Hacker News",
      visits: [{ date: new Date(now - 15 * 60 * 1000) }],
    },
    {
      url: "https://mozilla.org/en-US/",
      title: "Internet for people, not profit — Mozilla",
      visits: [{ date: new Date(now - 25 * 60 * 1000) }],
    },
  ];
  await PlacesUtils.history.clear();
  await PlacesUtils.history.insertMany(seeded);
}

/**
 * Shortcut for full browser history aggregation pipeline
 */
async function getBrowserHistoryAggregates() {
  const profileRecords = await getRecentHistory();
  const profilePreparedInputs = await generateProfileInputs(profileRecords);
  const [domainAgg, titleAgg, searchAgg] = aggregateSessions(
    profilePreparedInputs
  );

  return await topkAggregates(domainAgg, titleAgg, searchAgg);
}

/**
 * Builds fake chat history data for testing
 */
async function buildFakeChatHistory() {
  const fixedNow = 1_700_000_000_000;

  return [
    new ChatMessage({
      createdDate: fixedNow - 1_000,
      ordinal: 1,
      role: MESSAGE_ROLE.USER,
      content: { type: "text", body: "I like dogs." },
      pageUrl: "https://example.com/1",
      turnIndex: 0,
    }),
    new ChatMessage({
      createdDate: fixedNow - 10_000,
      ordinal: 2,
      role: MESSAGE_ROLE.USER,
      content: { type: "text", body: "I also like cats." },
      pageUrl: "https://example.com/2",
      turnIndex: 0,
    }),
    new ChatMessage({
      createdDate: fixedNow - 100_000,
      ordinal: 3,
      role: MESSAGE_ROLE.USER,
      content: {
        type: "text",
        body: "Tell me a joke about my favorite animals.",
      },
      pageUrl: "https://example.com/3",
      turnIndex: 0,
    }),
  ];
}

/**
 * Tests building the prompt for initial memories generation
 */
add_task(async function test_buildInitialMemoriesGenerationPrompt() {
  // Check that history is rendered correctly into CSV tables
  await buildFakeBrowserHistory();
  const [domainItems, titleItems, searchItems] =
    await getBrowserHistoryAggregates();
  const renderedBrowserHistory = await renderRecentHistoryForPrompt(
    domainItems,
    titleItems,
    searchItems
  );
  Assert.equal(
    renderedBrowserHistory,
    `# Domains
Domain,Importance Score
www.google.com,100
news.ycombinator.com,100
mozilla.org,100

# Titles
Title,Importance Score
Google Search: firefox history,100
Hacker News,100
Internet for people, not profit — Mozilla,100

# Searches
Search,Importance Score
Google Search: firefox history,1`.trim()
  );

  // Check that the full prompt is built correctly with injected categories, intents, and browsing history
  const sources = { history: [domainItems, titleItems, searchItems] };
  const initialMemoriesPrompt =
    await buildInitialMemoriesGenerationPrompt(sources);
  Assert.ok(
    initialMemoriesPrompt.includes(
      "You are an expert at extracting memories from user browser data."
    ),
    "Initial memories generation prompt should pull from the correct base"
  );
  Assert.ok(
    initialMemoriesPrompt.includes(getFormattedMemoryAttributeList(CATEGORIES)),
    "Prompt should include formatted categories list"
  );
  Assert.ok(
    initialMemoriesPrompt.includes(getFormattedMemoryAttributeList(INTENTS)),
    "Prompt should include formatted intents list"
  );
  Assert.ok(
    initialMemoriesPrompt.includes(renderedBrowserHistory),
    "Prompt should include rendered browsing history"
  );
});

/**
 * Tests rendering history as CSV when only search data is present
 */
add_task(async function test_buildRecentHistoryCSV_only_search() {
  const now = Date.now();
  const seeded = [
    {
      url: "https://www.google.com/search?q=firefox+history",
      title: "Google Search: firefox history",
      visits: [{ date: new Date(now - 5 * 60 * 1000) }],
    },
  ];
  await PlacesUtils.history.clear();
  await PlacesUtils.history.insertMany(seeded);

  const [domainItems, titleItems, searchItems] =
    await getBrowserHistoryAggregates();
  const renderedBrowserHistory = await renderRecentHistoryForPrompt(
    domainItems,
    titleItems,
    searchItems
  );
  Assert.equal(
    renderedBrowserHistory,
    `# Domains
Domain,Importance Score
www.google.com,100

# Titles
Title,Importance Score
Google Search: firefox history,100

# Searches
Search,Importance Score
Google Search: firefox history,1`.trim()
  );
});

/**
 * Tests rendering history as CSV when only history data is present
 */
add_task(async function test_buildRecentHistoryCSV_only_browsing_history() {
  const now = Date.now();
  const seeded = [
    {
      url: "https://news.ycombinator.com/",
      title: "Hacker News",
      visits: [{ date: new Date(now - 15 * 60 * 1000) }],
    },
    {
      url: "https://mozilla.org/en-US/",
      title: "Internet for people, not profit — Mozilla",
      visits: [{ date: new Date(now - 25 * 60 * 1000) }],
    },
  ];
  await PlacesUtils.history.clear();
  await PlacesUtils.history.insertMany(seeded);

  const [domainItems, titleItems, searchItems] =
    await getBrowserHistoryAggregates();
  const renderedBrowserHistory = await renderRecentHistoryForPrompt(
    domainItems,
    titleItems,
    searchItems
  );
  Assert.equal(
    renderedBrowserHistory,
    `# Domains
Domain,Importance Score
news.ycombinator.com,100
mozilla.org,100

# Titles
Title,Importance Score
Hacker News,100
Internet for people, not profit — Mozilla,100`.trim()
  );
});

/**
 * Tests building the prompt for initial memories generation with only chat data
 */
add_task(async function test_buildInitialMemoriesGenerationPrompt_only_chat() {
  const messages = await buildFakeChatHistory();
  const sb = sinon.createSandbox();
  const maxResults = 3;
  const halfLifeDays = 7;
  const startTime = 1_700_000_000_000 - 1_000_000;

  try {
    // Stub the method
    const stub = sb
      .stub(ChatStore.prototype, "findMessagesByDate")
      .callsFake(async () => {
        return messages;
      });

    const recentMessages = await getRecentChats(
      startTime,
      maxResults,
      halfLifeDays
    );

    // Assert stub was actually called
    Assert.equal(stub.callCount, 1, "findMessagesByDate should be called once");

    // Double check we get only the 3 expected messages back
    Assert.equal(recentMessages.length, 3, "Should return 3 chat messages");

    // Render the messages into CSV format and check correctness
    const renderedConversationHistory =
      await renderRecentConversationForPrompt(recentMessages);
    Assert.equal(
      renderedConversationHistory,
      `# Chat History
Message
I like dogs.
I also like cats.
Tell me a joke about my favorite animals.`.trim(),
      "Rendered conversation history should match expected CSV format"
    );

    // Build the actual prompt and check its contents
    const sources = { conversation: recentMessages };
    const initialMemoriesPrompt =
      await buildInitialMemoriesGenerationPrompt(sources);
    Assert.ok(
      initialMemoriesPrompt.includes(
        "You are an expert at extracting memories from user browser data."
      ),
      "Initial memories generation prompt should pull from the correct base"
    );
    Assert.ok(
      initialMemoriesPrompt.includes(renderedConversationHistory),
      "Prompt should include rendered conversation history"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests building the prompt for memories deduplication
 */
add_task(async function test_buildMemoriesDeduplicationPrompt() {
  const memoriesDeduplicationPrompt = await buildMemoriesDeduplicationPrompt(
    EXISTING_MEMORIES,
    NEW_MEMORIES
  );
  Assert.ok(
    memoriesDeduplicationPrompt.includes(
      "You are an expert at identifying duplicate statements."
    ),
    "Memories deduplication prompt should pull from the correct base"
  );
  Assert.ok(
    memoriesDeduplicationPrompt.includes(
      formatListForPrompt(EXISTING_MEMORIES)
    ),
    "Deduplication prompt should include existing memories list"
  );
  Assert.ok(
    memoriesDeduplicationPrompt.includes(formatListForPrompt(NEW_MEMORIES)),
    "Deduplication prompt should include new memories list"
  );
});

/**
 * Tests building the prompt for memories sensitivity filtering
 */
add_task(async function test_buildMemoriesSensitivityFilterPrompt() {
  /** Memories sensitivity filter prompt */
  const memoriesSensitivityFilterPrompt =
    await buildMemoriesSensitivityFilterPrompt(NEW_MEMORIES);
  Assert.ok(
    memoriesSensitivityFilterPrompt.includes(
      "You are an expert at identifying sensitive statements and content."
    ),
    "Memories sensitivity filter prompt should pull from the correct base"
  );
  Assert.ok(
    memoriesSensitivityFilterPrompt.includes(formatListForPrompt(NEW_MEMORIES)),
    "Sensitivity filter prompt should include memories list"
  );
});

/**
 * Tests successful initial memories generation
 */
add_task(async function test_generateInitialMemoriesList_happy_path() {
  const sb = sinon.createSandbox();
  try {
    /**
     * The fake engine returns canned LLM response.
     * The main `generateInitialMemoriesList` function should modify this heavily, cutting it back to only the required fields.
     */
    const fakeEngine = {
      run() {
        return {
          finalOutput: `[
  {
    "why": "User has recently searched for Firefox history and visited mozilla.org.",
    "category": "Internet & Telecom",
    "intent": "Research / Learn",
    "memory_summary": "Searches for Firefox information",
    "score": 7,
    "evidence": [
      {
        "type": "search",
        "value": "Google Search: firefox history"
      },
      {
        "type": "domain",
        "value": "mozilla.org"
      }
    ]
  },
  {
    "why": "User buys dog food online regularly from multiple sources.",
    "category": "Pets & Animals",
    "intent": "Buy / Acquire",
    "memory_summary": "Purchases dog food online",
    "score": -1,
    "evidence": [
      {
        "type": "domain",
        "value": "example.com"
      }
    ]
  }
]`,
        };
      },
    };

    // Check that the stub was called
    const stub = sb.stub(openAIEngine, "_createEngine").returns(fakeEngine);
    const engine = await openAIEngine.build(
      MODEL_FEATURES.MEMORIES,
      DEFAULT_ENGINE_ID,
      SERVICE_TYPES.MEMORIES
    );
    Assert.ok(stub.calledOnce, "_createEngine should be called once");

    const [domainItems, titleItems, searchItems] =
      await getBrowserHistoryAggregates();
    const sources = { history: [domainItems, titleItems, searchItems] };
    const memoriesList = await generateInitialMemoriesList(engine, sources);

    // Check top level structure
    Assert.ok(
      Array.isArray(memoriesList),
      "Should return an array of memories"
    );
    Assert.equal(memoriesList.length, 2, "Array should contain 2 memories");

    // Check first memory structure and content
    const firstMemory = memoriesList[0];
    Assert.equal(
      typeof firstMemory,
      "object",
      "First memory should be an object/map"
    );
    Assert.equal(
      Object.keys(firstMemory).length,
      4,
      "First memory should have 4 keys"
    );
    Assert.equal(
      firstMemory.category,
      "Internet & Telecom",
      "First memory should have expected category (Internet & Telecom)"
    );
    Assert.equal(
      firstMemory.intent,
      "Research / Learn",
      "First memory should have expected intent (Research / Learn)"
    );
    Assert.equal(
      firstMemory.memory_summary,
      "Searches for Firefox information",
      "First memory should have expected summary"
    );
    Assert.equal(
      firstMemory.score,
      5,
      "First memory should have expected score, clamping 7 to 5"
    );

    // Check that the second memory's score was clamped to the minimum
    const secondMemory = memoriesList[1];
    Assert.equal(
      secondMemory.score,
      1,
      "Second memory should have expected score, clamping -1 to 1"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests failed initial memories generation - Empty output
 */
add_task(
  async function test_generateInitialMemoriesList_sad_path_empty_output() {
    const sb = sinon.createSandbox();
    try {
      // LLM returns an empty memories list
      const fakeEngine = {
        run() {
          return {
            finalOutput: `[]`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").returns(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const [domainItems, titleItems, searchItems] =
        await getBrowserHistoryAggregates();
      const sources = { history: [domainItems, titleItems, searchItems] };
      const memoriesList = await generateInitialMemoriesList(engine, sources);

      Assert.equal(Array.isArray(memoriesList), true, "Should return an array");
      Assert.equal(memoriesList.length, 0, "Array should contain 0 memories");
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed initial memories generation - Output not array
 */
add_task(
  async function test_generateInitialMemoriesList_sad_path_output_not_array() {
    const sb = sinon.createSandbox();
    try {
      // LLM doesn't return an array
      const fakeEngine = {
        run() {
          return {
            finalOutput: `testing`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").returns(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const [domainItems, titleItems, searchItems] =
        await getBrowserHistoryAggregates();
      const sources = { history: [domainItems, titleItems, searchItems] };
      const memoriesList = await generateInitialMemoriesList(engine, sources);

      Assert.equal(Array.isArray(memoriesList), true, "Should return an array");
      Assert.equal(memoriesList.length, 0, "Array should contain 0 memories");
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed initial memories generation - Output not array of maps
 */
add_task(
  async function test_generateInitialMemoriesList_sad_path_output_not_array_of_maps() {
    const sb = sinon.createSandbox();
    try {
      // LLM doesn't return an array of maps
      const fakeEngine = {
        run() {
          return {
            finalOutput: `["testing1", "testing2", ["testing3"]]`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").returns(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const [domainItems, titleItems, searchItems] =
        await getBrowserHistoryAggregates();
      const sources = { history: [domainItems, titleItems, searchItems] };
      const memoriesList = await generateInitialMemoriesList(engine, sources);

      Assert.equal(Array.isArray(memoriesList), true, "Should return an array");
      Assert.equal(memoriesList.length, 0, "Array should contain 0 memories");
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed initial memories generation - Some correct memories
 */
add_task(
  async function test_generateInitialMemoriesList_sad_path_some_correct_memories() {
    const sb = sinon.createSandbox();
    try {
      // LLM returns an memories list where 1 is fully correct and 1 is missing required keys (category in this case)
      const fakeEngine = {
        run() {
          return {
            finalOutput: `[
  {
    "why": "User has recently searched for Firefox history and visited mozilla.org.",
    "intent": "Research / Learn",
    "memory_summary": "Searches for Firefox information",
    "score": 7,
    "evidence": [
      {
        "type": "search",
        "value": "Google Search: firefox history"
      },
      {
        "type": "domain",
        "value": "mozilla.org"
      }
    ]
  },
  {
    "why": "User buys dog food online regularly from multiple sources.",
    "category": "Pets & Animals",
    "intent": "Buy / Acquire",
    "memory_summary": "Purchases dog food online",
    "score": -1,
    "evidence": [
      {
        "type": "domain",
        "value": "example.com"
      }
    ]
  }
]`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").returns(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const [domainItems, titleItems, searchItems] =
        await getBrowserHistoryAggregates();
      const sources = { history: [domainItems, titleItems, searchItems] };
      const memoriesList = await generateInitialMemoriesList(engine, sources);

      Assert.equal(
        Array.isArray(memoriesList),
        true,
        "Should return an array of memories"
      );
      Assert.equal(memoriesList.length, 1, "Array should contain 1 memory");
      Assert.equal(
        memoriesList[0].memory_summary,
        "Purchases dog food online",
        "Memory summary should match the valid memory"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests successful memories deduplication
 */
add_task(async function test_deduplicateMemoriesList_happy_path() {
  const sb = sinon.createSandbox();
  try {
    /**
     * The fake engine that returns a canned LLM response for deduplication.
     * The `deduplicateMemories` function should return an array containing only the `main_memory` values.
     */
    const fakeEngine = {
      run() {
        return {
          finalOutput: `{
            "unique_memories": [
              {
                "main_memory": "Loves outdoor activities",
                "duplicates": ["Loves hiking and camping"]
              },
              {
                "main_memory": "Enjoys cooking recipes",
                "duplicates": []
              },
              {
                "main_memory": "Like sci-fi media",
                "duplicates": ["Reads science fiction novels"]
              },
              {
                "main_memory": "Likes both dogs and cats",
                "duplicates": []
              },
              {
                "main_memory": "Likes risky stock bets",
                "duplicates": []
              }
            ]
          }`,
        };
      },
    };

    // Check that the stub was called
    const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
    const engine = await openAIEngine.build(
      MODEL_FEATURES.MEMORIES,
      DEFAULT_ENGINE_ID,
      SERVICE_TYPES.MEMORIES
    );
    Assert.ok(stub.calledOnce, "_createEngine should be called once");

    const dedupedMemoriesList = await deduplicateMemories(
      engine,
      EXISTING_MEMORIES,
      NEW_MEMORIES
    );

    // Check that the deduplicated list contains only unique memories (`main_memory` values)
    Assert.equal(
      dedupedMemoriesList.length,
      5,
      "Deduplicated memories list should contain 5 unique memories"
    );
    Assert.ok(
      dedupedMemoriesList.includes("Loves outdoor activities"),
      "Deduplicated memories should include 'Loves outdoor activities'"
    );
    Assert.ok(
      dedupedMemoriesList.includes("Enjoys cooking recipes"),
      "Deduplicated memories should include 'Enjoys cooking recipes'"
    );
    Assert.ok(
      dedupedMemoriesList.includes("Like sci-fi media"),
      "Deduplicated memories should include 'Like sci-fi media'"
    );
    Assert.ok(
      dedupedMemoriesList.includes("Likes both dogs and cats"),
      "Deduplicated memories should include 'Likes both dogs and cats'"
    );
    Assert.ok(
      dedupedMemoriesList.includes("Likes risky stock bets"),
      "Deduplicated memories should include 'Likes risky stock bets'"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests failed memories deduplication - Empty output
 */
add_task(async function test_deduplicateMemoriesList_sad_path_empty_output() {
  const sb = sinon.createSandbox();
  try {
    // LLM returns the correct schema but with an empty unique_memories array
    const fakeEngine = {
      run() {
        return {
          finalOutput: `{
            "unique_memories": []
          }`,
        };
      },
    };

    // Check that the stub was called
    const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
    const engine = await openAIEngine.build(
      MODEL_FEATURES.MEMORIES,
      DEFAULT_ENGINE_ID,
      SERVICE_TYPES.MEMORIES
    );
    Assert.ok(stub.calledOnce, "_createEngine should be called once");

    const dedupedMemoriesList = await deduplicateMemories(
      engine,
      EXISTING_MEMORIES,
      NEW_MEMORIES
    );

    Assert.ok(Array.isArray(dedupedMemoriesList), "Should return an array");
    Assert.equal(dedupedMemoriesList.length, 0, "Should return an empty array");
  } finally {
    sb.restore();
  }
});

/**
 * Tests failed memories deduplication - Wrong top-level data type
 */
add_task(
  async function test_deduplicateMemoriesList_sad_path_wrong_top_level_data_type() {
    const sb = sinon.createSandbox();
    try {
      // LLM returns an incorrect data type
      const fakeEngine = {
        run() {
          return {
            finalOutput: `testing`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const dedupedMemoriesList = await deduplicateMemories(
        engine,
        EXISTING_MEMORIES,
        NEW_MEMORIES
      );

      Assert.ok(Array.isArray(dedupedMemoriesList), "Should return an array");
      Assert.equal(
        dedupedMemoriesList.length,
        0,
        "Should return an empty array"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed memories deduplication - Wrong inner data type
 */
add_task(
  async function test_deduplicateMemoriesList_sad_path_wrong_inner_data_type() {
    const sb = sinon.createSandbox();
    try {
      // LLM returns a map with the right top-level key, but the inner structure is wrong
      const fakeEngine = {
        run() {
          return {
            finalOutput: `{
            "unique_memories": "testing"
          }`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const dedupedMemoriesList = await deduplicateMemories(
        engine,
        EXISTING_MEMORIES,
        NEW_MEMORIES
      );

      Assert.ok(Array.isArray(dedupedMemoriesList), "Should return an array");
      Assert.equal(
        dedupedMemoriesList.length,
        0,
        "Should return an empty array"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed memories deduplication - Wrong inner array structure
 */
add_task(
  async function test_deduplicateMemoriesList_sad_path_wrong_inner_array_structure() {
    const sb = sinon.createSandbox();
    try {
      // LLM returns a map of nested arrays, but the array structure is wrong
      const fakeEngine = {
        run() {
          return {
            finalOutput: `{
            "unique_memories": ["testing1", "testing2"]
          }`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const dedupedMemoriesList = await deduplicateMemories(
        engine,
        EXISTING_MEMORIES,
        NEW_MEMORIES
      );

      Assert.ok(Array.isArray(dedupedMemoriesList), "Should return an array");
      Assert.equal(
        dedupedMemoriesList.length,
        0,
        "Should return an empty array"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed memories deduplication - Incorrect top-level schema key
 */
add_task(
  async function test_deduplicateMemoriesList_sad_path_bad_top_level_key() {
    const sb = sinon.createSandbox();
    try {
      // LLm returns correct output except that the top-level key is wrong
      const fakeEngine = {
        run() {
          return {
            finalOutput: `{
            "correct_memories": [
              {
                "main_memory": "Loves outdoor activities",
                "duplicates": ["Loves hiking and camping"]
              },
              {
                "main_memory": "Enjoys cooking recipes",
                "duplicates": []
              },
              {
                "main_memory": "Like sci-fi media",
                "duplicates": ["Reads science fiction novels"]
              },
              {
                "main_memory": "Likes both dogs and cats",
                "duplicates": []
              },
              {
                "main_memory": "Likes risky stock bets",
                "duplicates": []
              }
            ]
          }`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const dedupedMemoriesList = await deduplicateMemories(
        engine,
        EXISTING_MEMORIES,
        NEW_MEMORIES
      );

      Assert.ok(Array.isArray(dedupedMemoriesList), "Should return an array");
      Assert.equal(
        dedupedMemoriesList.length,
        0,
        "Should return an empty array"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed memories deduplication - Some correct inner schema
 */
add_task(
  async function test_deduplicateMemoriesList_sad_path_bad_some_correct_inner_schema() {
    const sb = sinon.createSandbox();
    try {
      // LLm returns correct output except that 1 of the inner maps is wrong and 1 main_memory is the wrong data type
      const fakeEngine = {
        run() {
          return {
            finalOutput: `{
            "unique_memories": [
              {
                "primary_memory": "Loves outdoor activities",
                "duplicates": ["Loves hiking and camping"]
              },
              {
                "main_memory": "Enjoys cooking recipes",
                "duplicates": []
              },
              {
                "main_memory": 12345,
                "duplicates": []
              }
            ]
          }`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const dedupedMemoriesList = await deduplicateMemories(
        engine,
        EXISTING_MEMORIES,
        NEW_MEMORIES
      );

      Assert.ok(Array.isArray(dedupedMemoriesList), "Should return an array");
      Assert.equal(
        dedupedMemoriesList.length,
        1,
        "Should return an array with one valid memory"
      );
      Assert.equal(
        dedupedMemoriesList[0],
        "Enjoys cooking recipes",
        "Should return the single valid memory"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests successful memories sensitivity filtering
 */
add_task(async function test_filterSensitiveMemories_happy_path() {
  const sb = sinon.createSandbox();
  try {
    /**
     * The fake engine that returns a canned LLM response for deduplication.
     * The `filterSensitiveMemories` function should return the inner array from `non_sensitive_memories`.
     */
    const fakeEngine = {
      run() {
        return {
          finalOutput: `{
  "non_sensitive_memories": [
    "Loves hiking and camping",
    "Reads science fiction novels",
    "Likes both dogs and cats"
  ]
}`,
        };
      },
    };

    // Check that the stub was called
    const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
    const engine = await openAIEngine.build(
      MODEL_FEATURES.MEMORIES,
      DEFAULT_ENGINE_ID,
      SERVICE_TYPES.MEMORIES
    );
    Assert.ok(stub.calledOnce, "_createEngine should be called once");

    const nonSensitiveMemoriesList = await filterSensitiveMemories(
      engine,
      NEW_MEMORIES
    );

    // Check that the non-sensitive memories list contains only non-sensitive memories
    Assert.equal(
      nonSensitiveMemoriesList.length,
      3,
      "Non-sensitive memories list should contain 3 memories"
    );
    Assert.ok(
      nonSensitiveMemoriesList.includes("Loves hiking and camping"),
      "Non-sensitive memories should include 'Loves hiking and camping'"
    );
    Assert.ok(
      nonSensitiveMemoriesList.includes("Reads science fiction novels"),
      "Non-sensitive memories should include 'Reads science fiction novels'"
    );
    Assert.ok(
      nonSensitiveMemoriesList.includes("Likes both dogs and cats"),
      "Non-sensitive memories should include 'Likes both dogs and cats'"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests failed memories sensitivity filtering - Empty output
 */
add_task(async function test_filterSensitiveMemories_sad_path_empty_output() {
  const sb = sinon.createSandbox();
  try {
    // LLM returns an empty non_sensitive_memories array
    const fakeEngine = {
      run() {
        return {
          finalOutput: `{
  "non_sensitive_memories": []
}`,
        };
      },
    };

    // Check that the stub was called
    const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
    const engine = await openAIEngine.build(
      MODEL_FEATURES.MEMORIES,
      DEFAULT_ENGINE_ID,
      SERVICE_TYPES.MEMORIES
    );
    Assert.ok(stub.calledOnce, "_createEngine should be called once");

    const nonSensitiveMemoriesList = await filterSensitiveMemories(
      engine,
      NEW_MEMORIES
    );

    Assert.ok(
      Array.isArray(nonSensitiveMemoriesList),
      "Should return an array"
    );
    Assert.equal(
      nonSensitiveMemoriesList.length,
      0,
      "Should return an empty array"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests failed memories sensitivity filtering - Wrong data type
 */
add_task(
  async function test_filterSensitiveMemories_sad_path_wrong_data_type() {
    const sb = sinon.createSandbox();
    try {
      // LLM returns the wrong outer data type
      const fakeEngine = {
        run() {
          return {
            finalOutput: `testing`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const nonSensitiveMemoriesList = await filterSensitiveMemories(
        engine,
        NEW_MEMORIES
      );

      Assert.ok(
        Array.isArray(nonSensitiveMemoriesList),
        "Should return an array"
      );
      Assert.equal(
        nonSensitiveMemoriesList.length,
        0,
        "Should return an empty array"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed memories sensitivity filtering - Wrong inner data type
 */
add_task(
  async function test_filterSensitiveMemories_sad_path_wrong_inner_data_type() {
    const sb = sinon.createSandbox();
    try {
      // LLM returns a map with the non_sensitive_memories key, but its value's data type is wrong
      const fakeEngine = {
        run() {
          return {
            finalOutput: `{
  "non_sensitive_memories": "testing"
}`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const nonSensitiveMemoriesList = await filterSensitiveMemories(
        engine,
        NEW_MEMORIES
      );

      Assert.ok(
        Array.isArray(nonSensitiveMemoriesList),
        "Should return an array"
      );
      Assert.equal(
        nonSensitiveMemoriesList.length,
        0,
        "Should return an empty array"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed memories sensitivity filtering - Wrong outer schema
 */
add_task(
  async function test_filterSensitiveMemories_sad_path_wrong_outer_schema() {
    const sb = sinon.createSandbox();
    try {
      // LLM returns a map but with the wrong top-level key
      const fakeEngine = {
        run() {
          return {
            finalOutput: `{
  "these_are_non_sensitive_memories": [
    "testing1", "testing2", "testing3"
  ]
}`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const nonSensitiveMemoriesList = await filterSensitiveMemories(
        engine,
        NEW_MEMORIES
      );

      Assert.ok(
        Array.isArray(nonSensitiveMemoriesList),
        "Should return an array"
      );
      Assert.equal(
        nonSensitiveMemoriesList.length,
        0,
        "Should return an empty array"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests failed memories sensitivity filtering - Some correct inner schema
 */
add_task(
  async function test_filterSensitiveMemories_sad_path_some_correct_inner_schema() {
    const sb = sinon.createSandbox();
    try {
      // LLM returns a map with the non_sensitive_memories key, but the inner schema has a mix of correct and incorrect data types
      const fakeEngine = {
        run() {
          return {
            finalOutput: `{
  "non_sensitive_memories": [
    "correct",
    12345,
    {"bad": "schema"}
  ]
}`,
          };
        },
      };

      // Check that the stub was called
      const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
      const engine = await openAIEngine.build(
        MODEL_FEATURES.MEMORIES,
        DEFAULT_ENGINE_ID,
        SERVICE_TYPES.MEMORIES
      );
      Assert.ok(stub.calledOnce, "_createEngine should be called once");

      const nonSensitiveMemoriesList = await filterSensitiveMemories(
        engine,
        NEW_MEMORIES
      );

      Assert.ok(
        Array.isArray(nonSensitiveMemoriesList),
        "Should return an array"
      );
      Assert.equal(
        nonSensitiveMemoriesList.length,
        1,
        "Should return an array with one valid memory"
      );
      Assert.equal(
        nonSensitiveMemoriesList[0],
        "correct",
        "Should return the single valid memory"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests mapping filtered memories back to full memory objects
 */
add_task(async function test_mapFilteredMemoriesToInitialList() {
  // Raw mock full memories object list
  const initialMemoriesList = [
    // Imagined duplicate - should have been filtered out
    {
      category: "Pets & Animals",
      intent: "Buy / Acquire",
      memory_summary: "Buys dog food online",
      score: 4,
    },
    // Sensitive content (stocks) - should have been filtered out
    {
      category: "News",
      intent: "Research / Learn",
      memory_summary: "Likes to invest in risky stocks",
      score: 5,
    },
    {
      category: "Games",
      intent: "Entertain / Relax",
      memory_summary: "Enjoys strategy games",
      score: 3,
    },
  ];

  // Mock list of good memories to keep
  const filteredMemoriesList = ["Enjoys strategy games"];

  const finalMemoriesList = await mapFilteredMemoriesToInitialList(
    initialMemoriesList,
    filteredMemoriesList
  );

  // Check that only the non-duplicate, non-sensitive memory remains
  Assert.equal(
    finalMemoriesList.length,
    1,
    "Final memories should contain 1 memory"
  );
  Assert.equal(
    finalMemoriesList[0].category,
    "Games",
    "Final memory should have the correct category"
  );
  Assert.equal(
    finalMemoriesList[0].intent,
    "Entertain / Relax",
    "Final memory should have the correct intent"
  );
  Assert.equal(
    finalMemoriesList[0].memory_summary,
    "Enjoys strategy games",
    "Final memory should match the filtered memory"
  );
  Assert.equal(
    finalMemoriesList[0].score,
    3,
    "Final memory should have the correct score"
  );
});