// Rotating a refresh token leaves the spent one redeemable for a grace window,
// answered with the replacement that rotation already named. Without it, a
// client with two requests in flight — or one whose refresh response was lost —
// holds a token the server has retired and is signed out for good.

import { describe, expect, it } from 'vitest';
import { createRefreshJwt, readJwtId } from '../packages/core/src/auth.js';
import { PersonalDataServer } from '../packages/core/src/pds.js';

const SECRET = 'test-secret';
const DID = 'did:plc:x';
const HOUR = 60 * 60 * 1000;

/**
 * A session store over a plain Map. Its spendSession tests and writes with no
 * await in between, which is what the SQLite-backed stores get from running a
 * statement to completion.
 * @param {{gateReads?: number}} [opts] - Hold the first N getSession calls
 *   until all N have been made, so two rotations read the same row before
 *   either writes.
 */
function makeStore({ gateReads = 0 } = {}) {
  /** @type {Map<string, any>} */
  const rows = new Map();
  /** @type {Array<() => void>} */
  const waiting = [];
  const gate = async () => {
    if (waiting.length >= gateReads) return;
    await new Promise((resolve) => {
      waiting.push(/** @type {() => void} */ (resolve));
      if (waiting.length === gateReads)
        for (const release of waiting) release();
    });
  };
  return {
    rows,
    storage: {
      putSession: async (/** @type {any} */ s) =>
        void rows.set(s.jti, { ...s }),
      getSession: async (/** @type {string} */ jti) => {
        await gate();
        return rows.get(jti) || null;
      },
      listSessions: async (/** @type {string} */ did) =>
        [...rows.values()].filter((s) => s.did === did),
      deleteSession: async (/** @type {string} */ jti) => rows.delete(jti),
      spendSession: async (
        /** @type {string} */ jti,
        /** @type {number} */ expiresAt,
        /** @type {string} */ nextJti,
      ) => {
        const row = rows.get(jti);
        if (!row || (row.nextJti && row.nextJti !== nextJti)) return false;
        rows.set(jti, { ...row, expiresAt, nextJti });
        return true;
      },
      listAppPasswords: async () => [],
      deleteAppPassword: async () => true,
    },
  };
}

/** @param {any} storage */
function makeServer(storage) {
  return new PersonalDataServer({
    actorStorage: /** @type {any} */ ({ getHandle: async () => 'x.test' }),
    blobs: /** @type {any} */ ({}),
    jwtSecret: SECRET,
    hostname: 'pds.example.com',
    sharedStorage: /** @type {any} */ (storage),
  });
}

const refreshRequest = (
  /** @type {string} */ jwt,
  /** @type {string} */ ua = 'client/1.0',
) =>
  new Request(
    'https://pds.example.com/xrpc/com.atproto.server.refreshSession',
    {
      method: 'POST',
      headers: { authorization: `Bearer ${jwt}`, 'user-agent': ua },
    },
  );

/**
 * Open a session directly against the store, as a login would.
 * @param {any} pds
 * @param {string} [label]
 * @param {string} [ua]
 */
async function openSession(pds, label = 'Account password', ua = 'client/1.0') {
  const jwt = await createRefreshJwt(DID, SECRET, 90 * 24 * 60 * 60);
  await pds._sessions.recordSession(jwt, {
    did: DID,
    label,
    scope: 'com.atproto.access',
    request: refreshRequest(jwt, ua),
    dedupe: true,
  });
  return jwt;
}

describe('refresh token rotation', () => {
  it('answers a replayed refresh with the same replacement', async () => {
    const { storage } = makeStore();
    const pds = makeServer(storage);
    const original = await openSession(pds);

    const first = await (await pds.fetch(refreshRequest(original))).json();
    // The same token again: a second in-flight request, or a lost response.
    const replay = await (await pds.fetch(refreshRequest(original))).json();

    expect(readJwtId(first.refreshJwt)).toBe(readJwtId(replay.refreshJwt));

    // Both callers hold a token that still works, and it is one session.
    const next = await pds.fetch(refreshRequest(replay.refreshJwt));
    expect(next.status).toBe(200);
    expect(await pds._sessions.listLiveSessions(DID)).toHaveLength(1);
  });

  it('settles two simultaneous rotations on one successor', async () => {
    // Both requests read the session before either writes — the window the
    // conditional write closes.
    const { storage } = makeStore({ gateReads: 2 });
    const pds = makeServer(storage);
    // A login reads no session, so the gate is still armed for the refreshes.
    const original = await openSession(pds);

    const [a, b] = await Promise.all([
      pds.fetch(refreshRequest(original)),
      pds.fetch(refreshRequest(original)),
    ]);

    expect([a.status, b.status]).toEqual([200, 200]);
    const [tokenA, tokenB] = [await a.json(), await b.json()];
    // One successor, not two: the loser adopts the id the winner recorded.
    expect(readJwtId(tokenA.refreshJwt)).toBe(readJwtId(tokenB.refreshJwt));
    expect(await pds._sessions.listLiveSessions(DID)).toHaveLength(1);

    // And that one token still drives the session forward.
    const next = await pds.fetch(refreshRequest(tokenB.refreshJwt));
    expect(next.status).toBe(200);
  });

  it('replaying a token whose successor already rotated does not fork it', async () => {
    const { storage } = makeStore();
    const pds = makeServer(storage);
    const original = await openSession(pds);

    const second = await (await pds.fetch(refreshRequest(original))).json();
    const third = await (
      await pds.fetch(refreshRequest(second.refreshJwt))
    ).json();

    // The first token comes back around late, after the session has moved on
    // twice. It reissues the token it named, which by now is itself spent.
    const late = await pds.fetch(refreshRequest(original));
    expect(late.status).toBe(200);
    expect(readJwtId((await late.json()).refreshJwt)).toBe(
      readJwtId(second.refreshJwt),
    );

    // Still one session, and the newest token is still the live one.
    expect(await pds._sessions.listLiveSessions(DID)).toHaveLength(1);
    expect((await pds._sessions.listLiveSessions(DID))[0].jti).toBe(
      readJwtId(third.refreshJwt),
    );
  });

  it('does not list the spent row as a second login', async () => {
    const { storage, rows } = makeStore();
    const pds = makeServer(storage);
    const original = await openSession(pds);

    await pds.fetch(refreshRequest(original));

    // Kept in storage for the grace window, but not a login of its own.
    expect(rows.size).toBe(2);
    expect(await pds._sessions.listLiveSessions(DID)).toHaveLength(1);
  });

  it('stops honoring a spent token once its grace window closes', async () => {
    const { storage, rows } = makeStore();
    const pds = makeServer(storage);
    const original = await openSession(pds);
    await pds.fetch(refreshRequest(original));

    // Age the spent row past its grace.
    const spent = rows.get(/** @type {string} */ (readJwtId(original)));
    spent.expiresAt = Math.floor((Date.now() - HOUR) / 1000);

    const res = await pds.fetch(refreshRequest(original));
    expect(res.status).toBe(400);
    expect((await res.json()).error).toBe('ExpiredToken');
  });

  it('never extends a token past its own expiry', async () => {
    const { storage, rows } = makeStore();
    const pds = makeServer(storage);
    const original = await openSession(pds);
    const jti = /** @type {string} */ (readJwtId(original));
    // A token with less life left than the grace window would grant.
    const nearlyDone = Math.floor((Date.now() + 5 * 60 * 1000) / 1000);
    rows.get(jti).expiresAt = nearlyDone;

    await pds.fetch(refreshRequest(original));

    expect(rows.get(jti).expiresAt).toBe(nearlyDone);
  });

  it('a rotation does not retire a separate login it comes to resemble', async () => {
    const { storage } = makeStore();
    const pds = makeServer(storage);
    // Two devices, same credential, told apart only by User-Agent.
    const laptop = await openSession(pds, 'Account password', 'app/1.0');
    await openSession(pds, 'Account password', 'app/2.0');
    expect(await pds._sessions.listLiveSessions(DID)).toHaveLength(2);

    // The first device updates, so its refresh now reports the same
    // User-Agent as the second. Only a fresh login may collapse a session.
    await pds.fetch(refreshRequest(laptop, 'app/2.0'));

    expect(await pds._sessions.listLiveSessions(DID)).toHaveLength(2);
  });

  it('revoking a session closes the chain that could reopen it', async () => {
    const { storage } = makeStore();
    const pds = makeServer(storage);
    const original = await openSession(pds);
    const rotated = await (await pds.fetch(refreshRequest(original))).json();

    await pds._sessions.revokeSessionChain(
      DID,
      /** @type {string} */ (readJwtId(rotated.refreshJwt)),
    );

    // Replaying the spent token must not resurrect the revoked session.
    const res = await pds.fetch(refreshRequest(original));
    expect(res.status).toBe(401);
    expect(await pds._sessions.listLiveSessions(DID)).toHaveLength(0);
  });
});

describe('app password revocation', () => {
  it('ends the logins that app password opened, and only those', async () => {
    const { storage } = makeStore();
    const pds = makeServer(storage);
    await openSession(pds, 'App password: bot');
    await openSession(pds, 'Account password');

    const ended = await pds._sessions.revokeAppPasswordSessions(DID, 'bot');

    expect(ended).toBe(1);
    expect(
      (await pds._sessions.listLiveSessions(DID)).map((s) => s.label),
    ).toEqual(['Account password']);
  });

  it('takes the spent rows with it', async () => {
    const { storage, rows } = makeStore();
    const pds = makeServer(storage);
    const original = await openSession(pds, 'App password: bot');
    await pds.fetch(refreshRequest(original));
    expect(rows.size).toBe(2);

    expect(await pds._sessions.revokeAppPasswordSessions(DID, 'bot')).toBe(1);
    expect(rows.size).toBe(0);
  });
});