/* 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 {
  NewTabStarterGenerator,
  trimConversation,
  addMemoriesToPrompt,
  cleanInferenceOutput,
  generateConversationStartersSidebar,
  generateFollowupPrompts,
  MemoriesGetterForSuggestionPrompts,
} = ChromeUtils.importESModule(
  "moz-src:///browser/components/aiwindow/models/ConversationSuggestions.sys.mjs"
);

const { openAIEngine } = ChromeUtils.importESModule(
  "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs"
);
const { MemoriesManager } = ChromeUtils.importESModule(
  "moz-src:///browser/components/aiwindow/models/memories/MemoriesManager.sys.mjs"
);
const { MESSAGE_ROLE } = ChromeUtils.importESModule(
  "moz-src:///browser/components/aiwindow/ui/modules/ChatConstants.sys.mjs"
);
const { sinon } = ChromeUtils.importESModule(
  "resource://testing-common/Sinon.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 PREF_HISTORY_ENABLED = "places.history.enabled";
const PREF_PRIVATE_BROWSING = "browser.privatebrowsing.autostart";

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

async function loadRemoteSettingsSnapshot() {
  const file = do_get_file("ai-window-prompts-remote-settings-snapshot.json");
  const data = await IOUtils.readUTF8(file.path);
  return JSON.parse(data);
}

let REAL_REMOTE_SETTINGS_SNAPSHOT;

add_setup(async function () {
  REAL_REMOTE_SETTINGS_SNAPSHOT = await loadRemoteSettingsSnapshot();
});

/**
 * Cleans up preferences after testing
 */
registerCleanupFunction(() => {
  for (let pref of [
    PREF_API_KEY,
    PREF_ENDPOINT,
    PREF_MODEL,
    PREF_HISTORY_ENABLED,
    PREF_PRIVATE_BROWSING,
  ]) {
    if (Services.prefs.prefHasUserValue(pref)) {
      Services.prefs.clearUserPref(pref);
    }
  }
});

/**
 * Tests for trimConversation function
 */
add_task(async function test_trimConversation() {
  const cases = [
    {
      input: [
        // empty case
      ],
      expected: [],
    },
    {
      input: [
        // more than 15 messages
        ...Array.from({ length: 20 }, (_, i) => ({
          role: i % 2 === 0 ? MESSAGE_ROLE.USER : MESSAGE_ROLE.ASSISTANT,
          content: `Message ${i + 1}`,
        })),
      ],
      expected: [
        ...Array.from({ length: 15 }, (_, i) => ({
          role: (i + 5) % 2 === 0 ? "user" : "assistant",
          content: `Message ${i + 6}`,
        })),
      ],
    },
    {
      input: [
        // should remove tool call/responses
        { role: MESSAGE_ROLE.USER, content: "What's the weather like?" },
        {
          role: MESSAGE_ROLE.ASSISTANT,
          content: "",
          tool_call: { name: "get_weather", arguments: "{}" },
        },
        {
          role: MESSAGE_ROLE.TOOL,
          content: "Here are the latest news articles.",
        },
      ],
      expected: [{ role: "user", content: "What's the weather like?" }],
    },
    {
      input: [
        // should remove system message
        { role: MESSAGE_ROLE.SYSTEM, content: "System message" },
        { role: MESSAGE_ROLE.USER, content: "Hello" },
        { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
      ],
      expected: [
        { role: "user", content: "Hello" },
        { role: "assistant", content: "Hi there!" },
      ],
    },
    {
      input: [
        // should remove messages with empty content
        { role: MESSAGE_ROLE.USER, content: "\n" },
        { role: MESSAGE_ROLE.ASSISTANT, content: "   " },
      ],
      expected: [],
    },
    {
      input: [
        // no valid messages
        { role: MESSAGE_ROLE.SYSTEM, content: "System message" },
        {
          role: MESSAGE_ROLE.ASSISTANT,
          content: "",
          tool_call: { name: "get_info", arguments: "{}" },
        },
      ],
      expected: [],
    },
    {
      input: [
        // should slice after filtering invalid messages
        ...Array(10)
          .fill(0)
          .flatMap((_, i) => [
            {
              role: MESSAGE_ROLE.USER,
              content: `User message ${i + 1}`,
            },
            { role: MESSAGE_ROLE.SYSTEM, content: "System message" },
          ]),
      ],
      expected: [
        ...Array.from({ length: 10 }, (_, i) => ({
          role: "user",
          content: `User message ${i + 1}`,
        })),
      ],
    },
  ];

  for (const { input, expected } of cases) {
    const result = trimConversation(input);
    Assert.deepEqual(
      result,
      expected,
      "trimConversation should return only user/assistant messages with content"
    );
  }
});

/**
 * Test for addMemoriesToPrompt function when there are memories
 */
add_task(async function test_addMemoriesToPrompt_have_memories() {
  const sb = sinon.createSandbox();
  try {
    const basePrompt = "Base prompt content.";
    const fakeMemories = ["Memory summary 1", "Memory summary 2"];
    const memoriesStub = sb
      .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
      .resolves(fakeMemories);
    const conversationMemoriesPrompt = "Memories block:\n{memories}";
    const promptWithMemories = await addMemoriesToPrompt(
      basePrompt,
      conversationMemoriesPrompt
    );

    Assert.ok(
      memoriesStub.calledOnce,
      "getMemorySummariesForPrompt should be called"
    );
    Assert.ok(
      promptWithMemories.includes("- Memory summary 1") &&
        promptWithMemories.includes("- Memory summary 2"),
      "Prompt should include memories"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Test for addMemoriesToPrompt function when there are no memories
 */
add_task(async function test_addMemoriesToPrompt_dont_have_memories() {
  const sb = sinon.createSandbox();
  try {
    const basePrompt = "Base prompt content.";
    const fakeMemories = [];
    const memoriesStub = sb
      .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
      .resolves(fakeMemories);
    const conversationMemoriesPrompt = "Memories block:\n{memories}";
    const promptWithMemories = await addMemoriesToPrompt(
      basePrompt,
      conversationMemoriesPrompt
    );

    Assert.ok(
      memoriesStub.calledOnce,
      "getMemorySummariesForPrompt should be called"
    );
    Assert.equal(
      promptWithMemories,
      basePrompt,
      "Prompt should be unchanged when no memories"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for cleanInferenceOutput function
 */
add_task(function test_cleanInferenceOutput() {
  const sb = sinon.createSandbox();
  try {
    const cases = [
      {
        input: {
          finalOutput: `- Suggestion 1\n\n* Suggestion 2\n1. Suggestion 3.\n2)Suggestion 4\n[5] Suggestion 5\n6.\nLabel: Suggestion 6\nSuggestion 7.\nSuggestion 8?`,
        },
        expected: [
          "Suggestion 1",
          "Suggestion 2",
          "Suggestion 3",
          "Suggestion 4",
          "Suggestion 5",
          "Suggestion 6",
          "Suggestion 7",
          "Suggestion 8?",
        ],
      },
      {
        input: { finalOutput: `Suggestion X\nSuggestion Y\nSuggestion Z` },
        expected: ["Suggestion X", "Suggestion Y", "Suggestion Z"],
      },
      {
        input: { finalOutput: "" },
        expected: [],
      },
    ];

    for (const { input, expected } of cases) {
      const result = cleanInferenceOutput(input);
      Assert.deepEqual(
        result,
        expected,
        "cleanInferenceOutput should return expected output"
      );
    }
  } finally {
    sb.restore();
  }
});

/**
 * Tests for createNewTabPromptGenerator generating prompts based on tab count
 */
add_task(
  async function test_createNewTabPromptGenerator_with_history_enabled() {
    const sb = sinon.createSandbox();
    const writingPrompts = [
      "Write a first draft",
      "Improve writing",
      "Proofread a message",
    ];

    const planningPrompts = [
      "Simplify a topic",
      "Brainstorm ideas",
      "Help make a plan",
    ];
    try {
      const cases = [
        {
          input: { tabCount: -1 },
          expectedBrowsing: [],
        },
        {
          input: { tabCount: 0 },
          expectedBrowsing: ["Find tabs in history"],
        },
        {
          input: { tabCount: 1 },
          expectedBrowsing: ["Find tabs in history", "Summarize tabs"],
        },
        {
          input: { tabCount: 2 },
          expectedBrowsing: [
            "Find tabs in history",
            "Summarize tabs",
            "Compare tabs",
          ],
        },
        {
          input: { tabCount: 3 },
          expectedBrowsing: [
            "Find tabs in history",
            "Summarize tabs",
            "Compare tabs",
          ],
        },
      ];
      const promptGenerator = NewTabStarterGenerator;
      for (const { input, expectedBrowsing } of cases) {
        const results = await promptGenerator.getPrompts(input.tabCount);
        if (input.tabCount <= -1) {
          Assert.equal(results.length, 2, "Should return 2 suggestions");
        } else {
          Assert.equal(results.length, 3, "Should return 3 suggestions");
        }
        for (const result of results) {
          Assert.equal(
            result.type,
            "chat",
            "Each result should have type 'chat'"
          );
        }
        Assert.ok(
          writingPrompts.includes(results[0].text),
          "Results should include a valid writing prompt"
        );
        Assert.ok(
          planningPrompts.includes(results[1].text),
          "Results should include a valid planning prompt"
        );
        if (results[2]) {
          Assert.ok(
            expectedBrowsing.includes(results[2].text),
            "Results should include a valid browsing prompt"
          );
        }
      }
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests for createNewTabPromptGenerator generating prompts based on tab count, with history disabled
 */
add_task(
  async function test_createNewTabPromptGenerator_with_history_disabled() {
    const sb = sinon.createSandbox();
    const writingPrompts = [
      "Write a first draft",
      "Improve writing",
      "Proofread a message",
    ];

    const planningPrompts = [
      "Simplify a topic",
      "Brainstorm ideas",
      "Help make a plan",
    ];
    try {
      const cases = [
        {
          input: { tabCount: -1 },
          expectedBrowsing: [],
        },
        {
          input: { tabCount: 0 },
          expectedBrowsing: [],
        },
        {
          input: { tabCount: 1 },
          expectedBrowsing: ["Summarize tabs"],
        },
        {
          input: { tabCount: 2 },
          expectedBrowsing: ["Summarize tabs", "Compare tabs"],
        },
        {
          input: { tabCount: 3 },
          expectedBrowsing: ["Summarize tabs", "Compare tabs"],
        },
      ];
      for (const pref of [
        [{ key: PREF_HISTORY_ENABLED, value: false }],
        [{ key: PREF_PRIVATE_BROWSING, value: true }],
        [
          { key: PREF_HISTORY_ENABLED, value: false },
          { key: PREF_PRIVATE_BROWSING, value: true },
        ],
      ]) {
        for (const p of pref) {
          Services.prefs.setBoolPref(p.key, p.value);
        }
        const promptGenerator = NewTabStarterGenerator;
        for (const { input, expectedBrowsing } of cases) {
          const results = await promptGenerator.getPrompts(input.tabCount);
          if (input.tabCount <= 0) {
            Assert.equal(results.length, 2, "Should return 2 suggestions");
          } else {
            Assert.equal(results.length, 3, "Should return 3 suggestions");
          }
          for (const result of results) {
            Assert.equal(
              result.type,
              "chat",
              "Each result should have type 'chat'"
            );
          }
          Assert.ok(
            writingPrompts.includes(results[0].text),
            "Results should include a valid writing prompt"
          );
          Assert.ok(
            planningPrompts.includes(results[1].text),
            "Results should include a valid planning prompt"
          );
          if (results[2]) {
            Assert.ok(
              expectedBrowsing.includes(results[2].text),
              "Results should include a valid browsing prompt"
            );
          }
        }
        for (const p of pref) {
          Services.prefs.clearUserPref(p.key);
        }
      }
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests for generateConversationStartersSidebar successfully generating suggestions
 */
add_task(async function test_generateConversationStartersSidebar_happy_path() {
  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
  Services.prefs.setStringPref(PREF_MODEL, MODEL);

  const sb = sinon.createSandbox();
  try {
    // Mock the openAIEngine and memories response
    const fakeEngine = {
      run: sb.stub().resolves({
        finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
      }),
    };
    sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

    const fakeMemories = ["Memory summary 1", "Memory summary 2"];
    const memoriesStub = sb
      .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
      .resolves(fakeMemories);

    const n = 3;
    const contextTabs = [
      { title: "Current Tab", url: "https://current.example.com" },
      { title: "Tab 2", url: "https://tab2.example.com" },
    ];

    const result = await generateConversationStartersSidebar(
      contextTabs,
      n,
      true
    );
    Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
    Assert.ok(
      memoriesStub.calledOnce,
      "getMemorySummariesForPrompt should be called once"
    );

    // Verify the prompt content
    const callArgs = fakeEngine.run.firstCall.args[0];
    Assert.equal(
      callArgs.messages.length,
      2,
      "run should be called with 2 messages"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '{"title":"Current Tab","url":"https://current.example.com"}'
      ),
      "Prompt should include current tab info"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '{"title":"Tab 2","url":"https://tab2.example.com"}'
      ),
      "Prompt should include other tab info"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        "\n- Memory summary 1\n- Memory summary 2"
      ),
      "Prompt should include memory summaries"
    );

    Assert.deepEqual(
      result,
      [
        { text: "Suggestion 1", type: "chat" },
        { text: "Suggestion 2", type: "chat" },
        { text: "Suggestion 3", type: "chat" },
      ],
      "Suggestions should match expected values"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for generateConversationStartersSidebar without including memories
 */
add_task(
  async function test_generateConversationStartersSidebar_without_memories() {
    Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
    Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
    Services.prefs.setStringPref(PREF_MODEL, MODEL);

    const sb = sinon.createSandbox();
    try {
      // Mock the openAIEngine and memories response
      const fakeEngine = {
        run: sb.stub().resolves({
          finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
        }),
      };
      sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

      const fakeMemories = ["Memory summary 1", "Memory summary 2"];
      const memoriesStub = sb
        .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
        .resolves(fakeMemories);

      const n = 3;
      const contextTabs = [
        { title: "Current Tab", url: "https://current.example.com" },
        { title: "Tab 2", url: "https://tab2.example.com" },
      ];

      const result = await generateConversationStartersSidebar(
        contextTabs,
        n,
        false
      );
      Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
      Assert.ok(
        !memoriesStub.calledOnce,
        "getMemorySummariesForPrompt shouldn't be called"
      );

      // Verify the prompt content
      const callArgs = fakeEngine.run.firstCall.args[0];
      Assert.equal(
        callArgs.messages.length,
        2,
        "run should be called with 2 messages"
      );
      Assert.ok(
        callArgs.messages[1].content.includes(
          '{"title":"Current Tab","url":"https://current.example.com"}'
        ),
        "Prompt should include current tab info"
      );
      Assert.ok(
        callArgs.messages[1].content.includes(
          '{"title":"Tab 2","url":"https://tab2.example.com"}'
        ),
        "Prompt should include other tab info"
      );
      Assert.ok(
        !callArgs.messages[1].content.includes(
          "\n- Memory summary 1\n- Memory summary 2"
        ),
        "Prompt should not include memory summaries"
      );

      Assert.deepEqual(
        result,
        [
          { text: "Suggestion 1", type: "chat" },
          { text: "Suggestion 2", type: "chat" },
          { text: "Suggestion 3", type: "chat" },
        ],
        "Suggestions should match expected values"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests for generateConversationStartersSidebar when no memories are returned
 */
add_task(
  async function test_generateConversationStartersSidebar_no_memories_returned() {
    Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
    Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
    Services.prefs.setStringPref(PREF_MODEL, MODEL);

    const sb = sinon.createSandbox();
    try {
      // Mock the openAIEngine and memories response
      const fakeEngine = {
        run: sb.stub().resolves({
          finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
        }),
      };
      sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

      const fakeMemories = [];
      const memoriesStub = sb
        .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
        .resolves(fakeMemories);

      const n = 3;
      const contextTabs = [
        { title: "Current Tab", url: "https://current.example.com" },
        { title: "Tab 2", url: "https://tab2.example.com" },
      ];

      const result = await generateConversationStartersSidebar(
        contextTabs,
        n,
        true
      );
      Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
      Assert.ok(
        memoriesStub.calledOnce,
        "getMemorySummariesForPrompt should be called once"
      );

      // Verify the prompt content
      const callArgs = fakeEngine.run.firstCall.args[0];
      Assert.equal(
        callArgs.messages.length,
        2,
        "run should be called with 2 messages"
      );
      Assert.ok(
        callArgs.messages[1].content.includes(
          '{"title":"Current Tab","url":"https://current.example.com"}'
        ),
        "Prompt should include current tab info"
      );
      Assert.ok(
        callArgs.messages[1].content.includes(
          '{"title":"Tab 2","url":"https://tab2.example.com"}'
        ),
        "Prompt should include other tab info"
      );
      Assert.ok(
        !callArgs.messages[1].content.includes("\nUser Memories:\n"),
        "Prompt shouldn't include user memories block"
      );

      Assert.deepEqual(
        result,
        [
          { text: "Suggestion 1", type: "chat" },
          { text: "Suggestion 2", type: "chat" },
          { text: "Suggestion 3", type: "chat" },
        ],
        "Suggestions should match expected values"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests for generateConversationStartersSidebar when no tabs are provided
 */
add_task(async function test_generateConversationStartersSidebar_no_tabs() {
  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
  Services.prefs.setStringPref(PREF_MODEL, MODEL);

  const sb = sinon.createSandbox();
  try {
    // Mock the openAIEngine and memories response
    const fakeEngine = {
      run: sb.stub().resolves({
        finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
      }),
    };
    sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

    const fakeMemories = ["Memory summary 1", "Memory summary 2"];
    const memoriesStub = sb
      .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
      .resolves(fakeMemories);

    const n = 3;
    const contextTabs = [];

    const result = await generateConversationStartersSidebar(
      contextTabs,
      n,
      true
    );
    Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
    Assert.ok(
      memoriesStub.calledOnce,
      "getMemorySummariesForPrompt should be called once"
    );

    // Verify the prompt content
    const callArgs = fakeEngine.run.firstCall.args[0];
    Assert.equal(
      callArgs.messages.length,
      2,
      "run should be called with 2 messages"
    );
    Assert.ok(
      callArgs.messages[1].content.includes("\nNo current tab\n"),
      "Prompt should indicate no current tab"
    );
    Assert.ok(
      callArgs.messages[1].content.includes("\nNo tabs available\n"),
      "Prompt should indicate no tabs available"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        "\n- Memory summary 1\n- Memory summary 2"
      ),
      "Prompt should include memory summaries"
    );

    Assert.deepEqual(
      result,
      [
        { text: "Suggestion 1", type: "chat" },
        { text: "Suggestion 2", type: "chat" },
        { text: "Suggestion 3", type: "chat" },
      ],
      "Suggestions should match expected values"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for generateConversationStartersSidebar with one tab provided
 */
add_task(async function test_generateConversationStartersSidebar_one_tab() {
  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
  Services.prefs.setStringPref(PREF_MODEL, MODEL);

  const sb = sinon.createSandbox();
  try {
    // Mock the openAIEngine and memories response
    const fakeEngine = {
      run: sb.stub().resolves({
        finalOutput: `1. Suggestion 1\n\n- Suggestion 2\nLabel: Suggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
      }),
    };
    sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

    const fakeMemories = ["Memory summary 1", "Memory summary 2"];
    const memoriesStub = sb
      .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
      .resolves(fakeMemories);

    const n = 3;
    const contextTabs = [
      { title: "Only Tab", url: "https://only.example.com" },
    ];

    const result = await generateConversationStartersSidebar(
      contextTabs,
      n,
      true
    );
    Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
    Assert.ok(
      memoriesStub.calledOnce,
      "getMemorySummariesForPrompt should be called once"
    );

    // Verify the prompt content
    const callArgs = fakeEngine.run.firstCall.args[0];
    Assert.equal(
      callArgs.messages.length,
      2,
      "run should be called with 2 messages"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '\n{"title":"Only Tab","url":"https://only.example.com"}'
      ),
      "Prompt should include current tab info"
    );
    Assert.ok(
      callArgs.messages[1].content.includes("\nOnly current tab is open\n"),
      "Prompt should indicate only current tab is open"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        "\n- Memory summary 1\n- Memory summary 2"
      ),
      "Prompt should include memory summaries"
    );

    Assert.deepEqual(
      result,
      [
        { text: "Suggestion 1", type: "chat" },
        { text: "Suggestion 2", type: "chat" },
        { text: "Suggestion 3", type: "chat" },
      ],
      "Suggestions should match expected values"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests that generateConversationStartersSidebar handles engine errors gracefully
 */
add_task(
  async function test_generateConversationStartersSidebar_engine_error() {
    Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
    Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
    Services.prefs.setStringPref(PREF_MODEL, MODEL);

    const sb = sinon.createSandbox();
    try {
      // Mock the openAIEngine and memories response
      const fakeEngine = {
        run: sb.stub().rejects(new Error("Engine failure")),
      };
      sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

      const fakeMemories = ["Memory summary 1", "Memory summary 2"];
      sb.stub(
        MemoriesGetterForSuggestionPrompts,
        "getMemorySummariesForPrompt"
      ).resolves(fakeMemories);

      const n = 3;
      const contextTabs = [
        { title: "Only Tab", url: "https://only.example.com" },
      ];

      const result = await generateConversationStartersSidebar(
        contextTabs,
        n,
        true
      );
      Assert.deepEqual(result, [], "Should return empty array on engine error");
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests that assistant limitations are included in conversation starter prompts
 */
add_task(
  async function test_generateConversationStartersSidebar_includes_assistant_limitations() {
    Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
    Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
    Services.prefs.setStringPref(PREF_MODEL, MODEL);

    const sb = sinon.createSandbox();
    try {
      const fakeEngine = {
        run: sb.stub().resolves({
          finalOutput: `Suggestion 1\nSuggestion 2\nSuggestion 3`,
        }),
      };
      sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

      sb.stub(openAIEngine, "getRemoteClient").returns({
        get: sb.stub().resolves(REAL_REMOTE_SETTINGS_SNAPSHOT),
      });

      sb.stub(
        MemoriesGetterForSuggestionPrompts,
        "getMemorySummariesForPrompt"
      ).resolves([]);

      const n = 3;
      const contextTabs = [
        { title: "Test Tab", url: "https://test.example.com" },
      ];

      await generateConversationStartersSidebar(contextTabs, n, false);

      Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");

      const callArgs = fakeEngine.run.firstCall.args[0];
      Assert.ok(
        callArgs.messages[1].content.includes(
          "You can do this and cannot do that."
        ),
        "Prompt should include assistant limitations from remote settings"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests for generateFollowupPrompts successfully generating suggestions
 */
add_task(async function test_generateFollowupPrompts_happy_path() {
  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
  Services.prefs.setStringPref(PREF_MODEL, MODEL);

  const sb = sinon.createSandbox();
  try {
    // Mock the openAIEngine and memories response
    const fakeEngine = {
      run: sb.stub().resolves({
        finalOutput: `1. Suggestion 1\n\n- Suggestion 2.\nSuggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
      }),
    };
    sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

    const fakeMemories = ["Memory summary 1", "Memory summary 2"];
    const memoriesStub = sb
      .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
      .resolves(fakeMemories);

    const n = 2;
    const conversationHistory = [
      { role: MESSAGE_ROLE.USER, content: "Hello" },
      { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
    ];
    const currentTab = {
      title: "Current Tab",
      url: "https://current.example.com",
    };

    // Using memories
    const result = await generateFollowupPrompts(
      conversationHistory,
      currentTab,
      n,
      true
    );
    Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
    Assert.ok(
      memoriesStub.calledOnce,
      "getMemorySummariesForPrompt should be called once"
    );

    const callArgs = fakeEngine.run.firstCall.args[0];
    Assert.equal(
      callArgs.messages.length,
      2,
      "run should be called with 2 messages"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '{"title":"Current Tab","url":"https://current.example.com"}'
      ),
      "Prompt should include current tab info"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"}]'
      ),
      "Prompt should include conversation history"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        "\n- Memory summary 1\n- Memory summary 2"
      ),
      "Prompt should include memory summaries"
    );

    Assert.deepEqual(
      result,
      [
        { text: "Suggestion 1", type: "chat" },
        { text: "Suggestion 2", type: "chat" },
      ],
      "Suggestions should match"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for generateFollowupPrompts without including memories
 */
add_task(async function test_generateFollowupPrompts_no_memories() {
  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
  Services.prefs.setStringPref(PREF_MODEL, MODEL);

  const sb = sinon.createSandbox();
  try {
    // Mock the openAIEngine and memories response
    const fakeEngine = {
      run: sb.stub().resolves({
        finalOutput: `1. Suggestion 1\n\n- Suggestion 2.\nSuggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
      }),
    };
    sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

    const fakeMemories = ["Memory summary 1", "Memory summary 2"];
    const memoriesStub = sb
      .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
      .resolves(fakeMemories);

    const n = 2;
    const conversationHistory = [
      { role: MESSAGE_ROLE.USER, content: "Hello" },
      { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
    ];
    const currentTab = {
      title: "Current Tab",
      url: "https://current.example.com",
    };

    const result = await generateFollowupPrompts(
      conversationHistory,
      currentTab,
      n,
      false
    );
    Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
    Assert.ok(
      !memoriesStub.calledOnce,
      "getMemorySummariesForPrompt shouldn't be called"
    );

    const callArgs = fakeEngine.run.firstCall.args[0];
    Assert.equal(
      callArgs.messages.length,
      2,
      "run should be called with 2 messages"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '{"title":"Current Tab","url":"https://current.example.com"}'
      ),
      "Prompt should include current tab info"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"}]'
      ),
      "Prompt should include conversation history"
    );
    Assert.ok(
      !callArgs.messages[1].content.includes(
        "\n- Memory summary 1\n- Memory summary 2"
      ),
      "Prompt shouldn't include memory summaries"
    );

    Assert.deepEqual(
      result,
      [
        { text: "Suggestion 1", type: "chat" },
        { text: "Suggestion 2", type: "chat" },
      ],
      "Suggestions should match"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for generateFollowupPrompts when no memories are returned
 */
add_task(async function test_generateFollowupPrompts_no_memories_returned() {
  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
  Services.prefs.setStringPref(PREF_MODEL, MODEL);

  const sb = sinon.createSandbox();
  try {
    // Mock the openAIEngine and memories response
    const fakeEngine = {
      run: sb.stub().resolves({
        finalOutput: `1. Suggestion 1\n\n- Suggestion 2.\nSuggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
      }),
    };
    sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

    const fakeMemories = [];
    const memoriesStub = sb
      .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
      .resolves(fakeMemories);

    const n = 2;
    const conversationHistory = [
      { role: MESSAGE_ROLE.USER, content: "Hello" },
      { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
    ];
    const currentTab = {
      title: "Current Tab",
      url: "https://current.example.com",
    };

    // Using memories
    const result = await generateFollowupPrompts(
      conversationHistory,
      currentTab,
      n,
      true
    );
    Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
    Assert.ok(
      memoriesStub.calledOnce,
      "getMemorySummariesForPrompt should be called once"
    );

    const callArgs = fakeEngine.run.firstCall.args[0];
    Assert.equal(
      callArgs.messages.length,
      2,
      "run should be called with 2 messages"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '{"title":"Current Tab","url":"https://current.example.com"}'
      ),
      "Prompt should include current tab info"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"}]'
      ),
      "Prompt should include conversation history"
    );
    Assert.ok(
      !callArgs.messages[1].content.includes("\nUser Memories:\n"),
      "Prompt shouldn't include user memories block"
    );

    Assert.deepEqual(
      result,
      [
        { text: "Suggestion 1", type: "chat" },
        { text: "Suggestion 2", type: "chat" },
      ],
      "Suggestions should match"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for generateFollowupPrompts without a current tab
 */
add_task(async function test_generateFollowupPrompts_no_current_tab() {
  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
  Services.prefs.setStringPref(PREF_MODEL, MODEL);

  const sb = sinon.createSandbox();
  try {
    // Mock the openAIEngine and memories response
    const fakeEngine = {
      run: sb.stub().resolves({
        finalOutput: `1. Suggestion 1\n\n- Suggestion 2.\nSuggestion 3.\nSuggestion 4\nSuggestion 5\nSuggestion 6`,
      }),
    };
    sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

    const fakeMemories = [];
    const memoriesStub = sb
      .stub(MemoriesGetterForSuggestionPrompts, "getMemorySummariesForPrompt")
      .resolves(fakeMemories);

    const n = 2;
    const conversationHistory = [
      { role: MESSAGE_ROLE.USER, content: "Hello" },
      { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
    ];
    const currentTab = {};

    const result = await generateFollowupPrompts(
      conversationHistory,
      currentTab,
      n,
      false
    );
    Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");
    Assert.ok(
      !memoriesStub.calledOnce,
      "getMemorySummariesForPrompt shouldn't be called"
    );

    const callArgs = fakeEngine.run.firstCall.args[0];
    Assert.equal(
      callArgs.messages.length,
      2,
      "run should be called with 2 messages"
    );
    Assert.ok(
      callArgs.messages[1].content.includes("\nNo tab\n"),
      "Prompt shouldn't include any tab info"
    );
    Assert.ok(
      callArgs.messages[1].content.includes(
        '[{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"}]'
      ),
      "Prompt should include conversation history"
    );
    Assert.ok(
      !callArgs.messages[1].content.includes("\nUser Memories:\n"),
      "Prompt shouldn't include user memories block"
    );

    Assert.deepEqual(
      result,
      [
        { text: "Suggestion 1", type: "chat" },
        { text: "Suggestion 2", type: "chat" },
      ],
      "Suggestions should match"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests that generateFollowupPrompts handles engine errors gracefully
 */
add_task(async function test_generateFollowupPrompts_engine_error() {
  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
  Services.prefs.setStringPref(PREF_MODEL, MODEL);

  const sb = sinon.createSandbox();
  try {
    // Mock the openAIEngine and memories response
    const fakeEngine = {
      run: sb.stub().rejects(new Error("Engine failure")),
    };
    sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

    const fakeMemories = [];
    sb.stub(
      MemoriesGetterForSuggestionPrompts,
      "getMemorySummariesForPrompt"
    ).resolves(fakeMemories);

    const n = 2;
    const conversationHistory = [
      { role: MESSAGE_ROLE.USER, content: "Hello" },
      { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
    ];
    const currentTab = {};

    const result = await generateFollowupPrompts(
      conversationHistory,
      currentTab,
      n,
      false
    );
    Assert.deepEqual(result, [], "Should return empty array on engine error");
  } finally {
    sb.restore();
  }
});

/**
 * Tests that assistant limitations are included in followup prompts
 */
add_task(
  async function test_generateFollowupPrompts_includes_assistant_limitations() {
    Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
    Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
    Services.prefs.setStringPref(PREF_MODEL, MODEL);

    const sb = sinon.createSandbox();
    try {
      const fakeEngine = {
        run: sb.stub().resolves({
          finalOutput: `Suggestion 1\nSuggestion 2`,
        }),
      };
      sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);

      sb.stub(openAIEngine, "getRemoteClient").returns({
        get: sb.stub().resolves(REAL_REMOTE_SETTINGS_SNAPSHOT),
      });

      sb.stub(
        MemoriesGetterForSuggestionPrompts,
        "getMemorySummariesForPrompt"
      ).resolves([]);

      const n = 2;
      const conversationHistory = [
        { role: MESSAGE_ROLE.USER, content: "Hello" },
        { role: MESSAGE_ROLE.ASSISTANT, content: "Hi there!" },
      ];
      const currentTab = { title: "Test", url: "https://test.example.com" };

      await generateFollowupPrompts(conversationHistory, currentTab, n, false);

      Assert.ok(fakeEngine.run.calledOnce, "Engine run should be called once");

      const callArgs = fakeEngine.run.firstCall.args[0];
      Assert.ok(
        callArgs.messages[1].content.includes(
          "You can do this and cannot do that."
        ),
        "Prompt should include assistant limitations from remote settings"
      );
    } finally {
      sb.restore();
    }
  }
);

/**
 * Tests for getMemorySummariesForPrompt happy path
 */
add_task(async function test_getMemorySummariesForPrompt_happy_path() {
  const sb = sinon.createSandbox();
  try {
    // Mock the MemoryStore to return fixed memories
    const fakeMemories = [
      {
        memory_summary: "Memory summary 1",
      },
      {
        memory_summary: "Memory summary 2",
      },
      {
        memory_summary: "Memory summary 3",
      },
    ];

    sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

    const maxMemories = 2;
    const summaries =
      await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
        maxMemories
      );

    Assert.deepEqual(
      summaries,
      ["Memory summary 1", "Memory summary 2"],
      "Memory summaries should match expected values"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for getMemorySummariesForPrompt when no memories are returned
 */
add_task(async function test_getMemorySummariesForPrompt_no_memories() {
  const sb = sinon.createSandbox();
  try {
    // Mock the MemoryStore to return fixed memories
    const fakeMemories = [];

    sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

    const maxMemories = 2;
    const summaries =
      await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
        maxMemories
      );

    Assert.equal(
      summaries.length,
      0,
      `getMemorySummariesForPrompt(${maxMemories}) should return 0 summaries`
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for getMemorySummariesForPrompt with fewer memories than maxMemories
 */
add_task(async function test_getMemorySummariesForPrompt_too_few_memories() {
  const sb = sinon.createSandbox();
  try {
    // Mock the MemoryStore to return fixed memories
    const fakeMemories = [
      {
        memory_summary: "Memory summary 1",
      },
    ];

    sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

    const maxMemories = 2;
    const summaries =
      await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
        maxMemories
      );

    Assert.deepEqual(
      summaries,
      ["Memory summary 1"],
      "Memory summaries should match expected values"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for getMemorySummariesForPrompt handling duplicate summaries
 */
add_task(async function test_getMemorySummariesForPrompt_duplicates() {
  const sb = sinon.createSandbox();
  try {
    // Mock the MemoryStore to return fixed memories
    const fakeMemories = [
      {
        memory_summary: "Duplicate summary",
      },
      {
        memory_summary: "duplicate summary",
      },
      {
        memory_summary: "Unique summary",
      },
    ];

    sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

    const maxMemories = 2;
    const summaries =
      await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
        maxMemories
      );

    Assert.deepEqual(
      summaries,
      ["Duplicate summary", "Unique summary"],
      "Memory summaries should match expected values"
    );
  } finally {
    sb.restore();
  }
});

/**
 * Tests for getMemorySummariesForPrompt handling empty and whitespace-only summaries
 */
add_task(
  async function test_getMemorySummariesForPrompt_empty_and_whitespace() {
    const sb = sinon.createSandbox();
    try {
      // Mock the MemoryStore to return fixed memories
      const fakeMemories = [
        {
          memory_summary: "   \n",
        },
        {
          memory_summary: "",
        },
        {
          memory_summary: "Valid summary",
        },
      ];

      sb.stub(MemoriesManager, "getAllMemories").resolves(fakeMemories);

      const maxMemories = 2;
      const summaries =
        await MemoriesGetterForSuggestionPrompts.getMemorySummariesForPrompt(
          maxMemories
        );

      Assert.deepEqual(
        summaries,
        ["Valid summary"],
        "Memory summaries should match expected values"
      );
    } finally {
      sb.restore();
    }
  }
);