import { PersonalDataServer } from '@pdsjs/core';
import { afterEach, describe, expect, test, vi } from 'vitest';
import {
  base64UrlDecode,
  generateKeyPair,
  importPrivateKey,
} from '../src/crypto.js';
import { createIdentityHandlers } from '../src/handlers/xrpc-identity.js';
import { generateKemKey, kemDidKey } from '../src/kem.js';
import {
  buildPlcOperation,
  computePlcOperationCid,
  publicKeyToDidKey,
} from '../src/plc.js';
import { cborEncodeDagCbor } from '../src/repo.js';
import { createVerifier } from '../src/verify.js';

const DID = 'did:plc:kemkeytestaccount2345672';
const PLC = 'https://plc.example.com';

afterEach(() => {
  vi.unstubAllGlobals();
  vi.restoreAllMocks();
});

describe('KEM keys', () => {
  test('a fresh key names its own P-256 did:key', async () => {
    const { privateKey, didKey } = await generateKemKey();

    expect(privateKey).toHaveLength(32);
    expect(didKey).toMatch(/^did:key:zDn/);
    expect(await kemDidKey(privateKey)).toBe(didKey);
  });

  test('refuses a key of the wrong length', async () => {
    await expect(kemDidKey(new Uint8Array(31))).rejects.toThrow('32 bytes');
  });

  test('a genesis operation publishes the KEM key beside the signing key', () => {
    const op = /** @type {{verificationMethods: Record<string, string>}} */ (
      buildPlcOperation({
        didKey: 'did:key:zDnaeSIGN',
        kemDidKey: 'did:key:zDnaeKEM',
        handle: 'alice.example.com',
        pdsEndpoint: 'https://pds.example.com',
      })
    );

    expect(op.verificationMethods).toEqual({
      atproto: 'did:key:zDnaeSIGN',
      atproto_kem: 'did:key:zDnaeKEM',
    });
  });
});

/**
 * @param {{
 *   did?: string|null, kemKey?: Uint8Array|null, readOnly?: boolean,
 *   deactivated?: boolean, hostname?: string,
 * }} [options]
 */
async function createIdentity({
  did = DID,
  kemKey = null,
  readOnly = false,
  deactivated = false,
  hostname = 'pds.example.com',
} = {}) {
  const { privateKey, publicKey } = await generateKeyPair();
  const signingKey = await importPrivateKey(privateKey);
  const serverKey = publicKeyToDidKey(publicKey);
  const stored = { kemKey };
  /** @type {Array<{type: string, body: Object}>} */
  const events = [];
  const identity = createIdentityHandlers({
    actorStorage: /** @type {any} */ ({
      async getHandle() {
        return 'alice.example.com';
      },
      async getPrivateKey() {
        return privateKey;
      },
      async getKemKey() {
        return stored.kemKey;
      },
      async setKemKey(/** @type {Uint8Array} */ key) {
        stored.kemKey = key;
      },
    }),
    sharedStorage: /** @type {any} */ ({}),
    hostname,
    plcUrl: PLC,
    jwtSecret: 'test-secret',
    readOnly,
    emailer: null,
    k256Signer: null,
    getDid: async () => did,
    getSigningKey: async () => signingKey,
    isDeactivated: async () => deactivated,
    handleProxyHeader: async () => null,
    emitEvent: async (type, body) => {
      events.push({ type, body });
    },
    readJsonBounded: async (request) => request.json(),
    readOnlyError: () => Response.json({}, { status: 401 }),
  });
  return { identity, stored, events, serverKey };
}

/**
 * The last operation in a PLC log this server can update.
 * @param {string} serverKey
 * @param {Record<string, unknown>} [overrides]
 */
function lastOp(serverKey, overrides = {}) {
  return {
    type: 'plc_operation',
    rotationKeys: ['did:key:zQ3shRECOVERY', serverKey],
    verificationMethods: { atproto: serverKey },
    alsoKnownAs: ['at://alice.example.com'],
    services: {
      atproto_pds: {
        type: 'AtprotoPersonalDataServer',
        endpoint: 'https://pds.example.com',
      },
    },
    prev: null,
    sig: 'c2ln',
    ...overrides,
  };
}

/**
 * A PLC directory answering the log read with `last` and a submission with
 * `submitStatus`.
 * @param {Object|null} last - null answers the log read with a 500
 * @param {number} [submitStatus]
 */
function stubPlc(last, submitStatus = 200) {
  /** @type {Array<Record<string, any>>} */
  const submitted = [];
  const fetchMock = vi.fn(
    async (/** @type {string} */ url, /** @type {RequestInit} */ init) => {
      if (url === `${PLC}/${encodeURIComponent(DID)}/log/last`) {
        return last
          ? Response.json(last)
          : new Response('down', { status: 500 });
      }
      if (
        url === `${PLC}/${encodeURIComponent(DID)}` &&
        init?.method === 'POST'
      ) {
        submitted.push(JSON.parse(/** @type {string} */ (init.body)));
        return new Response('', { status: submitStatus });
      }
      return new Response('not found', { status: 404 });
    },
  );
  vi.stubGlobal('fetch', fetchMock);
  return { submitted, fetchMock };
}

describe('recommended DID credentials', () => {
  test('carry the KEM key once the account has one', async () => {
    stubPlc(null);
    const kem = await generateKemKey();
    const { identity } = await createIdentity({ kemKey: kem.privateKey });
    const creds = await identity.getRecommendedDidCredentials();

    expect(creds?.verificationMethods.atproto_kem).toBe(kem.didKey);
  });
});

describe('submitPlcOperation', () => {
  /**
   * @param {ReturnType<typeof createIdentityHandlers>} identity
   * @param {Object} operation
   */
  async function submit(identity, operation) {
    const path = '/xrpc/com.atproto.identity.submitPlcOperation';
    const url = new URL(`https://pds.example.com${path}`);
    const handler =
      /** @type {(request: Request, url: URL, auth: unknown) => Promise<Response>} */ (
        identity.routes[path].handler
      );
    return handler(
      new Request(url, { method: 'POST', body: JSON.stringify({ operation }) }),
      url,
      { did: DID, scope: 'com.atproto.access' },
    );
  }

  test('refuses a caller with no session, server or operation', async () => {
    stubPlc(null);
    const path = '/xrpc/com.atproto.identity.submitPlcOperation';
    const url = new URL(`https://pds.example.com${path}`);
    /**
     * @param {ReturnType<typeof createIdentityHandlers>} identity
     * @param {Object} body
     * @param {unknown} auth
     */
    const post = (identity, body, auth) =>
      /** @type {(request: Request, url: URL, auth: unknown) => Promise<Response>} */ (
        identity.routes[path].handler
      )(
        new Request(url, { method: 'POST', body: JSON.stringify(body) }),
        url,
        auth,
      );
    const auth = { did: DID, scope: 'com.atproto.access' };

    const { identity } = await createIdentity();
    expect((await post(identity, { operation: {} }, null)).status).toBe(401);
    expect((await post(identity, {}, auth)).status).toBe(400);
    const empty = await createIdentity({ did: null });
    expect((await post(empty.identity, { operation: {} }, auth)).status).toBe(
      400,
    );
  });

  test('reports what the directory said when it refuses', async () => {
    stubPlc(null, 409);
    const { identity } = await createIdentity();
    const response = await submit(identity, { verificationMethods: {} });

    expect(response.status).toBe(400);
    expect((await response.json()).message).toContain('409');
  });

  test('refuses an operation that drops the KEM key', async () => {
    const { submitted } = stubPlc(null);
    const kem = await generateKemKey();
    const { identity } = await createIdentity({ kemKey: kem.privateKey });
    const response = await submit(identity, {
      verificationMethods: { atproto: 'did:key:zDnaeSIGN' },
    });

    expect(response.status).toBe(400);
    expect((await response.json()).message).toBe('Incorrect KEM key');
    expect(submitted).toEqual([]);
  });

  test('passes an operation that keeps the KEM key', async () => {
    const { submitted } = stubPlc(null);
    const kem = await generateKemKey();
    const { identity } = await createIdentity({ kemKey: kem.privateKey });
    const response = await submit(identity, {
      verificationMethods: {
        atproto: 'did:key:zDnaeSIGN',
        atproto_kem: kem.didKey,
      },
    });

    expect(response.status).toBe(200);
    expect(submitted).toHaveLength(1);
  });
});

describe('KEM key backfill', () => {
  test('publishes a key, then stores it and announces the identity', async () => {
    const created = await createIdentity();
    const last = lastOp(created.serverKey);
    const { submitted } = stubPlc(last);

    expect(await created.identity.backfillKemKey()).toBe('published');
    expect(submitted).toHaveLength(1);
    const [op] = submitted;
    const kem = await kemDidKey(
      /** @type {Uint8Array} */ (created.stored.kemKey),
    );
    expect(op.verificationMethods).toEqual({
      atproto: created.serverKey,
      atproto_kem: kem,
    });
    expect(op.rotationKeys).toEqual(last.rotationKeys);
    expect(op.prev).toBe(await computePlcOperationCid(last));

    const { sig, ...unsigned } = op;
    const valid = await createVerifier().verify(
      created.serverKey,
      cborEncodeDagCbor(unsigned),
      base64UrlDecode(sig),
    );
    expect(valid).toBe(true);
    expect(created.events.map((e) => e.type)).toEqual(['identity']);
  });

  test('stores a key it cannot publish when another key rotates the DID', async () => {
    const created = await createIdentity();
    const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
    const { submitted } = stubPlc(
      lastOp(created.serverKey, { rotationKeys: ['did:key:zQ3shOTHER'] }),
    );

    expect(await created.identity.backfillKemKey()).toBe('stored');
    expect(submitted).toEqual([]);
    expect(created.stored.kemKey).toHaveLength(32);
    expect(created.events).toEqual([]);
    expect(warn).toHaveBeenCalled();
  });

  test('stores a key it cannot publish when the DID names another server', async () => {
    const created = await createIdentity();
    vi.spyOn(console, 'warn').mockImplementation(() => {});
    const { submitted } = stubPlc(
      lastOp(created.serverKey, {
        services: {
          atproto_pds: {
            type: 'AtprotoPersonalDataServer',
            endpoint: 'https://elsewhere.example.com',
          },
        },
      }),
    );

    expect(await created.identity.backfillKemKey()).toBe('stored');
    expect(submitted).toEqual([]);
  });

  test('stores a key for a did:web, whose document is kept elsewhere', async () => {
    const created = await createIdentity({ did: 'did:web:alice.example.com' });
    vi.spyOn(console, 'warn').mockImplementation(() => {});
    const { fetchMock } = stubPlc(null);

    expect(await created.identity.backfillKemKey()).toBe('stored');
    expect(fetchMock).not.toHaveBeenCalled();
  });

  test('keeps nothing when the directory refuses, so the next run retries', async () => {
    const created = await createIdentity();
    stubPlc(lastOp(created.serverKey), 503);

    await expect(created.identity.backfillKemKey()).rejects.toThrow('503');
    expect(created.stored.kemKey).toBeNull();
  });

  test('keeps nothing when the log cannot be read', async () => {
    const created = await createIdentity();
    stubPlc(null);

    await expect(created.identity.backfillKemKey()).rejects.toThrow(
      'Could not read the PLC log',
    );
    expect(created.stored.kemKey).toBeNull();
  });

  test('leaves alone an account that has a key, none, or has moved away', async () => {
    const { fetchMock } = stubPlc(null);
    const kem = await generateKemKey();
    for (const options of [
      { kemKey: kem.privateKey },
      { did: null },
      { deactivated: true },
      { readOnly: true },
    ]) {
      const { identity } = await createIdentity(options);
      expect(await identity.backfillKemKey()).toBe('skipped');
    }
    expect(fetchMock).not.toHaveBeenCalled();
  });
});

describe('the server', () => {
  /** @param {Record<string, any>} actorStorage */
  function server(actorStorage) {
    /** @type {import('../src/ports.js').ExtensionHost|undefined} */
    let host;
    const pds = new PersonalDataServer({
      actorStorage: /** @type {any} */ (actorStorage),
      sharedStorage: /** @type {any} */ ({}),
      blobs: /** @type {any} */ ({}),
      jwtSecret: 'test-secret',
      hostname: 'pds.example.com',
      plcUrl: PLC,
      password: 'pw',
      extensions: [
        (h) => {
          host = h;
          return { name: 'probe' };
        },
      ],
    });
    return { pds, host: () => /** @type {NonNullable<typeof host>} */ (host) };
  }

  test('hands extensions the KEM key, read once', async () => {
    const kem = await generateKemKey();
    const getKemKey = vi.fn(async () => kem.privateKey);
    const { pds, host } = server({ getKemKey });

    expect(await host().getKemKey()).toEqual(kem);
    expect(await pds.getKemKey()).toEqual(kem);
    expect(getKemKey).toHaveBeenCalledTimes(1);
  });

  test('answers null while the account has no KEM key', async () => {
    const { pds } = server({ getKemKey: async () => null });
    expect(await pds.getKemKey()).toBeNull();
  });

  test('logs a failed backfill and reports it', async () => {
    const error = vi.spyOn(console, 'error').mockImplementation(() => {});
    stubPlc(null);
    const { privateKey } = await generateKeyPair();
    const { pds } = server({
      getDid: async () => DID,
      getKemKey: async () => null,
      getPrivateKey: async () => privateKey,
      getAccountStatus: async () => 'active',
    });

    expect(await pds.maybeBackfillKemKey()).toBe('failed');
    expect(error).toHaveBeenCalled();
  });

  test('reads the new key after a backfill stores one', async () => {
    vi.spyOn(console, 'warn').mockImplementation(() => {});
    const { privateKey } = await generateKeyPair();
    /** @type {Uint8Array|null} */
    let kemKey = null;
    const { pds } = server({
      getDid: async () => 'did:web:alice.example.com',
      getKemKey: async () => kemKey,
      setKemKey: async (/** @type {Uint8Array} */ key) => {
        kemKey = key;
      },
      getPrivateKey: async () => privateKey,
      getAccountStatus: async () => 'active',
    });

    expect(await pds.getKemKey()).toBeNull();
    expect(await pds.maybeBackfillKemKey()).toBe('stored');
    expect((await pds.getKemKey())?.privateKey).toBe(kemKey);
  });
});