import { beforeEach, describe, expect, it, vi } from "vitest";
import { deriveKey } from "../crypto/argon2";

// PCH-01-008: a password change must refresh the MV3 session-key escrow so the
// next service-worker restore decrypts the new-key Slot A. The backing stores
// are closed over by the chrome mock on globalThis, so they survive the
// vi.resetModules() we use to simulate a service-worker restart mid-test.


type StorageShape = Record<string, unknown>;

function deferred(): { promise: Promise<void>; resolve: () => void } {
  let resolve!: () => void;
  const promise = new Promise<void>((done) => {
    resolve = done;
  });
  return { promise, resolve };
}

function pauseSessionKeyWrite(session: StorageShape): {
  started: Promise<void>;
  release: () => void;
} {
  const started = deferred();
  const release = deferred();
  const sessionSet = chrome.storage.session.set as unknown as ReturnType<typeof vi.fn>;
  sessionSet.mockImplementation(async (next: Record<string, unknown>) => {
    if ("peach_session_dk" in next) {
      started.resolve();
      await release.promise;
    }
    Object.assign(session, next);
  });
  return { started: started.promise, release: release.resolve };
}

describe("vault session escrow refresh (PCH-01-008)", () => {
  let local: StorageShape;
  let session: StorageShape;

  beforeEach(() => {
    vi.resetModules();
    local = {};
    session = {};
    const area = (store: StorageShape) => ({
      get: vi.fn(async (key?: string | string[]) => {
        if (!key) return store;
        if (typeof key === "string") return { [key]: store[key] };
        return key.reduce<Record<string, unknown>>((r, k) => { r[k] = store[k]; return r; }, {});
      }),
      set: vi.fn(async (next: Record<string, unknown>) => { Object.assign(store, next); }),
      remove: vi.fn(async (keys: string | string[]) => {
        for (const k of Array.isArray(keys) ? keys : [keys]) delete store[k];
      }),
    });
    (globalThis as { chrome?: typeof chrome }).chrome = {
      storage: { local: area(local), session: area(session) },
      alarms: { clear: vi.fn(), create: vi.fn(), onAlarm: { addListener: vi.fn() } },
      idle: { setDetectionInterval: vi.fn(), onStateChanged: { addListener: vi.fn() } },
    } as unknown as typeof chrome;
    vi.stubGlobal("fetch", vi.fn(async () => { throw new Error("network disabled"); }));
  });

  it("change_password_updates_session_escrow", async () => {
    const v1 = await import("./vault");
    await v1.createVault("old-pass-abc123");
    await v1.addEntry({ type: "login", title: "Acct", password: "pw" });
    expect(await v1.changePassword("old-pass-abc123", "new-pass-xyz789")).toBe(true);

    // Simulate a SW restart after the awaited escrow refresh: a fresh module
    // has derivedKey=null but session/local persist.
    vi.resetModules();
    const v2 = await import("./vault");
    expect(await v2.restoreFromSession()).toBe(true);
    expect(v2.getEntries().some((e) => e.title === "Acct")).toBe(true);
  });

  it("change_password_with_new_key_updates_session_escrow", async () => {
    const v1 = await import("./vault");
    await v1.createVault("orig-pass-abc123");
    await v1.addEntry({ type: "login", title: "Recovered", password: "pw" });

    const newSalt = crypto.getRandomValues(new Uint8Array(16));
    const newKey = await deriveKey("reset-pass-xyz789", newSalt);
    await v1.changePasswordWithNewKey(newSalt, newKey);

    vi.resetModules();
    const v2 = await import("./vault");
    expect(await v2.restoreFromSession()).toBe(true);
    expect(v2.getEntries().some((e) => e.title === "Recovered")).toBe(true);
  });

  it("restoreFromSession rehydrates rootKey so background sync survives a SW restart", async () => {
    // Background sync bug: entry.ts calls restoreFromSession() on every SW wake-up,
    // then triggerSyncIfReachable. The runner needs getRootKeyHex() to be non-null
    // — sync identity derives from rootKey, not derivedKey. The previous
    // restoreFromSession only set rootKey when the slot had recovery_salt +
    // recovery_wrapped_root_key; for any vault that hadn't been through
    // the recovery bootstrap, rootKey was silently null and the runner returned
    // VAULT_LOCKED. The fix is to always derive rootKey from pw_root_* (which
    // are present on every vault since createVault).
    const v1 = await import("./vault");
    await v1.createVault("strong-pass-abc123");
    // unlock() is the real-world path that saves the MV3 session key — the
    // user's popup opens, types the password, and the session key is parked
    // in chrome.storage.session for the next SW wake.
    await v1.unlock("strong-pass-abc123");
    const rootBeforeDead = v1.getRootKeyHex();
    expect(rootBeforeDead).not.toBeNull();

    // Simulate a SW kill: flush the fire-and-forget saveSessionKey promise,
    // then reset modules so a fresh vault.ts reads the persisted session key.
    await Promise.resolve();
    vi.resetModules();
    const v2 = await import("./vault");
    expect(await v2.restoreFromSession()).toBe(true);
    // The bug shape: rootKey was null after restore, so the runner short-circuited
    // to VAULT_LOCKED even though entries were present.
    expect(v2.getRootKeyHex()).not.toBeNull();
    expect(v2.getRootKeyHex()).toBe(rootBeforeDead);
  });

  it("lock() clears the MV3 session key so a manual lock survives a SW restart", async () => {
    // Bug shape: lock() used to keep the session key in chrome.storage.session
    // so background sync could re-wake a restored vault. MV3 idle-kills the SW
    // after ~30s, so a popup reopened a minute after locking woke a cold SW,
    // restoreFromSession() found the surviving key, and the vault silently
    // re-unlocked without the password. Every lock path (manual lock icon,
    // popup close with lockOnClose=true, auto-lock alarm, idle) flows through
    // lock(), and all of them are deliberate locks — none should be undone by
    // the next SW restart. clearSessionKey() mirrors the fire-and-forget
    // saveSessionKey() pattern used by unlock().
    const v1 = await import("./vault");
    await v1.createVault("strong-pass-abc123");
    await v1.unlock("strong-pass-abc123");
    const sessionKeyBeforeLock = await chrome.storage.session.get("peach_session_dk");
    expect(sessionKeyBeforeLock.peach_session_dk).toBeDefined();

    // The lock-on-close path: popup port disconnect, alarm tick, etc.
    v1.lock();
    expect(v1.isUnlocked()).toBe(false);
    // Flush the fire-and-forget clearSessionKey microtask, then confirm the
    // key was actually removed from chrome.storage.session.
    await Promise.resolve();
    const sessionKeyAfterLock = await chrome.storage.session.get("peach_session_dk");
    expect(sessionKeyAfterLock.peach_session_dk).toBeUndefined();

    // The next SW wake must NOT restore the vault — restoreFromSession() finds
    // no session key and the vault stays locked until the password is typed.
    vi.resetModules();
    const v2 = await import("./vault");
    expect(await v2.restoreFromSession()).toBe(false);
    expect(v2.isUnlocked()).toBe(false);
    expect(v2.getRootKeyHex()).toBeNull();
  });

  it("a lock invalidates a real unlock paused while writing session escrow", async () => {
    const vault = await import("./vault");
    await vault.createVault("strong-pass-abc123");
    await vault.lock();

    const write = pauseSessionKeyWrite(session);
    const unlocking = vault.unlock("strong-pass-abc123");
    await write.started;

    let lockSettled = false;
    const locking = vault.lock().then(() => {
      lockSettled = true;
    });
    await new Promise((resolve) => setTimeout(resolve, 0));
    const lockSettledBeforeWrite = lockSettled;
    write.release();

    expect(lockSettledBeforeWrite).toBe(false);
    await expect(unlocking).resolves.toBe(false);
    await locking;
    expect(vault.isUnlocked()).toBe(false);
    expect(session.peach_session_dk).toBeUndefined();

    vi.resetModules();
    const restarted = await import("./vault");
    expect(await restarted.restoreFromSession()).toBe(false);
  });

  it("a lock invalidates a duress unlock paused while writing session escrow", async () => {
    const vault = await import("./vault");
    await vault.createVault("real-pass-abc123");
    await vault.setupDuress("duress-pass-xyz789");
    await vault.lock();

    const write = pauseSessionKeyWrite(session);
    const unlocking = vault.unlock("duress-pass-xyz789");
    await write.started;

    const locking = vault.lock();
    await Promise.resolve();
    write.release();

    await expect(unlocking).resolves.toBe(false);
    await locking;
    expect(vault.isUnlocked()).toBe(false);
    expect(session.peach_session_dk).toBeUndefined();
  });

  it("a lock invalidates a key unlock paused while writing session escrow", async () => {
    const vault = await import("./vault");
    await vault.createVault("strong-pass-abc123");
    const keyHex = vault.getDerivedKeyHex();
    await vault.lock();

    const write = pauseSessionKeyWrite(session);
    const unlocking = vault.unlockWithKey(keyHex);
    await write.started;

    const locking = vault.lock();
    await Promise.resolve();
    write.release();

    await expect(unlocking).resolves.toEqual({ ok: false, error: "invalid_key" });
    await locking;
    expect(vault.isUnlocked()).toBe(false);
    expect(session.peach_session_dk).toBeUndefined();
  });

  it("a lock invalidates a password change paused while writing session escrow", async () => {
    const vault = await import("./vault");
    await vault.createVault("old-pass-abc123");

    const write = pauseSessionKeyWrite(session);
    const changing = vault.changePassword("old-pass-abc123", "new-pass-xyz789");
    await write.started;

    const locking = vault.lock();
    await Promise.resolve();
    write.release();

    await expect(changing).resolves.toBe(false);
    await locking;
    expect(vault.isUnlocked()).toBe(false);
    expect(session.peach_session_dk).toBeUndefined();

    vi.resetModules();
    const restarted = await import("./vault");
    expect(await restarted.restoreFromSession()).toBe(false);
  });

  it("a lock clears a recovery re-key paused while writing session escrow", async () => {
    const vault = await import("./vault");
    await vault.createVault("old-pass-abc123");
    const newSalt = crypto.getRandomValues(new Uint8Array(16));
    const newKey = await deriveKey("new-pass-xyz789", newSalt);

    const write = pauseSessionKeyWrite(session);
    const changing = vault.changePasswordWithNewKey(newSalt, newKey);
    await write.started;

    let lockSettled = false;
    const locking = vault.lock().then(() => {
      lockSettled = true;
    });
    await new Promise((resolve) => setTimeout(resolve, 0));
    const lockSettledBeforeWrite = lockSettled;
    write.release();

    expect(lockSettledBeforeWrite).toBe(false);
    await expect(changing).rejects.toThrow("vault_locked");
    await locking;
    expect(vault.isUnlocked()).toBe(false);
    expect(session.peach_session_dk).toBeUndefined();

    vi.resetModules();
    const restarted = await import("./vault");
    expect(await restarted.restoreFromSession()).toBe(false);
  });

  it("a lock invalidates a restore paused while reading the session key", async () => {
    const v1 = await import("./vault");
    await v1.createVault("strong-pass-abc123");
    await v1.unlock("strong-pass-abc123");
    await vi.waitFor(() => expect(session.peach_session_dk).toBeDefined());

    const staleSessionKey = session.peach_session_dk;
    const readStarted = deferred();
    const releaseRead = deferred();
    const sessionGet = chrome.storage.session.get as unknown as ReturnType<typeof vi.fn>;
    sessionGet.mockImplementation(async (key?: string | string[]) => {
      if (key === "peach_session_dk") {
        readStarted.resolve();
        await releaseRead.promise;
        return { peach_session_dk: staleSessionKey };
      }
      if (!key) return session;
      const keys = Array.isArray(key) ? key : [key];
      return keys.reduce<StorageShape>((result, entry) => {
        result[entry] = session[entry];
        return result;
      }, {});
    });

    vi.resetModules();
    const v2 = await import("./vault");
    const restoring = v2.restoreFromSession();
    await readStarted.promise;

    await v2.lock();
    releaseRead.resolve();

    await expect(restoring).resolves.toBe(false);
    expect(v2.isUnlocked()).toBe(false);
    expect(v2.getRootKeyHex()).toBeNull();
  });

  it("keeps decrypted restore state private until the final epoch-checked commit", async () => {
    const v1 = await import("./vault");
    await v1.createVault("strong-pass-abc123");
    await v1.unlock("strong-pass-abc123");
    await vi.waitFor(() => expect(session.peach_session_dk).toBeDefined());

    const settingsReadStarted = deferred();
    const releaseSettingsRead = deferred();
    const localGet = chrome.storage.local.get as unknown as ReturnType<typeof vi.fn>;
    localGet.mockImplementation(async (key?: string | string[]) => {
      if (key === "peach_settings") {
        settingsReadStarted.resolve();
        await releaseSettingsRead.promise;
      }
      if (!key) return local;
      const keys = Array.isArray(key) ? key : [key];
      return keys.reduce<StorageShape>((result, entry) => {
        result[entry] = local[entry];
        return result;
      }, {});
    });

    vi.resetModules();
    const v2 = await import("./vault");
    const restoring = v2.restoreFromSession();
    await settingsReadStarted.promise;

    try {
      expect(v2.isUnlocked()).toBe(false);
      const locking = Promise.resolve(v2.lock());
      expect(v2.isUnlocked()).toBe(false);
      releaseSettingsRead.resolve();
      await locking;

      await expect(restoring).resolves.toBe(false);
      expect(v2.isUnlocked()).toBe(false);
      expect(v2.getRootKeyHex()).toBeNull();
    } finally {
      releaseSettingsRead.resolve();
      await Promise.allSettled([restoring]);
    }
  });

  it("does not resolve lock until session escrow is absent", async () => {
    const vault = await import("./vault");
    const pending = await import("./recovery-pending");
    await vault.createVault("strong-pass-abc123");
    await vault.unlock("strong-pass-abc123");
    await pending.setRecoveryConfirmPending();
    await vi.waitFor(() => {
      expect(session.peach_session_dk).toBeDefined();
      expect(session.peach_recovery_pending).toBeDefined();
    });

    const removalStarted = deferred();
    const releaseRemoval = deferred();
    const sessionRemove = chrome.storage.session.remove as unknown as ReturnType<typeof vi.fn>;
    sessionRemove.mockImplementation(async (keys: string | string[]) => {
      removalStarted.resolve();
      await releaseRemoval.promise;
      for (const key of Array.isArray(keys) ? keys : [keys]) delete session[key];
    });

    let settled = false;
    const locking = Promise.resolve(vault.lock()).then(() => {
      settled = true;
    });
    expect(vault.isUnlocked()).toBe(false);
    await removalStarted.promise;

    try {
      await Promise.resolve();
      expect(settled).toBe(false);
    } finally {
      releaseRemoval.resolve();
      await locking;
    }

    expect(session.peach_session_dk).toBeUndefined();
    expect(session.peach_recovery_pending).toBeDefined();
  });

  it("disabled escrow clears the stored key and blocks save and restore", async () => {
    const v1 = await import("./vault");
    await v1.createVault("strong-pass-abc123");
    await v1.unlock("strong-pass-abc123");
    await vi.waitFor(() => expect(session.peach_session_dk).toBeDefined());
    const parked = session.peach_session_dk;

    await v1.applyStorageHardening(false);
    expect(session.peach_session_dk).toBeUndefined();

    // A later unlock must not re-park the key while escrow is off.
    await v1.unlock("strong-pass-abc123");
    await Promise.resolve();
    expect(session.peach_session_dk).toBeUndefined();

    // A restarted worker with the same stored key restores only when the
    // hardening result allows escrow; a disabled escrow fails closed.
    session.peach_session_dk = parked;
    vi.resetModules();
    const v2 = await import("./vault");
    expect(await v2.restoreFromSession()).toBe(true);

    vi.resetModules();
    const v3 = await import("./vault");
    await v3.applyStorageHardening(false);
    expect(session.peach_session_dk).toBeUndefined();
    expect(await v3.restoreFromSession()).toBe(false);
    expect(v3.isUnlocked()).toBe(false);
  });

  it("deduplicates simultaneous restores and runs the unlock hook once", async () => {
    const v1 = await import("./vault");
    await v1.createVault("strong-pass-abc123");
    await v1.unlock("strong-pass-abc123");
    await vi.waitFor(() => expect(session.peach_session_dk).toBeDefined());

    vi.resetModules();
    const v2 = await import("./vault");
    const afterUnlock = vi.fn(async () => undefined);
    v2.setLockedSyncHooks({
      refresh: vi.fn(async () => undefined),
      afterUnlock,
    });
    const sessionGet = chrome.storage.session.get as unknown as ReturnType<typeof vi.fn>;
    sessionGet.mockClear();

    await expect(Promise.all([
      v2.restoreFromSession(),
      v2.restoreFromSession(),
    ])).resolves.toEqual([true, true]);

    expect(sessionGet).toHaveBeenCalledTimes(1);
    expect(afterUnlock).toHaveBeenCalledTimes(1);
  });
});