import { describe, it, expect } from "vitest";
import { SqlMemory, NoopMemory } from "../src/ai/memory/memory.js";
import type { Db } from "../src/db/db.js";

/** In-memory fake Db implementing the three memory methods. */
function memoryFakeDb() {
  const store = new Map<string, string>();
  const db: Db = {
    botDbName: "bot",
    pool: {} as never,
    saveExchange: async () => ({ conversationId: 1, userId: 1 }),
    recentMessages: async () => [],
    memorySet: async (userId, key, value) => {
      store.set(`${userId}:${key}`, value);
    },
    memoryGet: async (userId, key) => store.get(`${userId}:${key}`),
    memoryGetAll: async (userId) => {
      const prefix = `${userId}:`;
      const out: Record<string, string> = {};
      for (const [k, v] of store) {
        if (k.startsWith(prefix)) out[k.slice(prefix.length)] = v;
      }
      return out;
    },
    close: async () => {},
  };
  return { db, store };
}

describe("SqlMemory", () => {
  it("remember + recall round-trip (per user)", async () => {
    const { db } = memoryFakeDb();
    const mem = new SqlMemory(db);

    await mem.remember(7, "lang", "zh");
    await mem.remember(7, "summary", "likes CDD01");
    await mem.remember(8, "lang", "en");

    expect(await mem.recall(7, "lang")).toBe("zh");
    expect(await mem.recall(7, "summary")).toBe("likes CDD01");
    // Other user's memory is isolated
    expect(await mem.recall(7, "lang")).toBe("zh");
    expect(await mem.recall(8, "lang")).toBe("en");
  });

  it("recall returns undefined for a missing key", async () => {
    const { db } = memoryFakeDb();
    const mem = new SqlMemory(db);
    expect(await mem.recall(1, "nope")).toBeUndefined();
  });

  it("remember upserts: second write with same (user,key) overwrites", async () => {
    const { db, store } = memoryFakeDb();
    const mem = new SqlMemory(db);

    await mem.remember(3, "lang", "zh");
    await mem.remember(3, "lang", "ms");

    expect(await mem.recall(3, "lang")).toBe("ms");
    // Only one entry ever existed
    let count = 0;
    for (const [k] of store) if (k.startsWith("3:")) count++;
    expect(count).toBe(1);
  });

  it("recallAll returns the full KV object without system keys leaking", async () => {
    const { db } = memoryFakeDb();
    const mem = new SqlMemory(db);
    await mem.remember(5, "lang", "ms");
    await mem.remember(5, "last_queries", '["a"]');

    const all = await mem.recallAll(5);
    expect(all).toEqual({ lang: "ms", last_queries: '["a"]' });
  });
});

describe("NoopMemory", () => {
  it("never stores and always recalls nothing", async () => {
    const mem = new NoopMemory();
    await mem.remember(1, "lang", "zh");
    expect(await mem.recall(1, "lang")).toBeUndefined();
    expect(await mem.recallAll(1)).toEqual({});
  });
});