import { beforeEach, describe, expect, it } from 'vitest';
import {
  createServiceJwt,
  verifyAccessJwt,
} from '../packages/core/src/auth.js';
import {
  base64UrlDecode,
  base64UrlEncode,
  generateKeyPair,
  importPrivateKey,
  sign,
} from '../packages/core/src/crypto.js';
import { PersonalDataServer } from '../packages/core/src/pds.js';
import { publicKeyToDidKey } from '../packages/core/src/plc.js';
import { createVerifier } from '../packages/core/src/verify.js';

const HOSTNAME = 'pds.example.com';
const PASSWORD = 'test-password';
const JWT_SECRET = 'test-secret';
const MIGRATING_DID = 'did:plc:migrateexampleaccount000';
const HANDLE = 'alice.pds.example.com';

/**
 * A stateful in-memory actor store, uninitialized until createAccount adopts a
 * DID. Only the fields createAccount, createSession, and describeServer read
 * are backed; the rest are inert stubs.
 */
function createEmptyStorage() {
  /** @type {Record<string, any>} */
  const state = {
    did: null,
    handle: null,
    privateKey: null,
    status: null,
    email: null,
  };
  return {
    _state: state,
    async getDid() {
      return state.did;
    },
    async setDid(/** @type {string} */ did) {
      state.did = did;
    },
    async getHandle() {
      return state.handle;
    },
    async setHandle(/** @type {string} */ handle) {
      state.handle = handle;
    },
    async getPrivateKey() {
      return state.privateKey;
    },
    async setPrivateKey(/** @type {Uint8Array} */ key) {
      state.privateKey = key;
    },
    async getKemKey() {
      return state.kemKey ?? null;
    },
    async setKemKey(/** @type {Uint8Array} */ key) {
      state.kemKey = key;
    },
    async getAccountStatus() {
      return state.status;
    },
    async setAccountStatus(/** @type {string} */ status) {
      state.status = status;
    },
    async getEmail() {
      return state.email;
    },
    async setEmail(/** @type {any} */ email) {
      state.email = email;
    },
  };
}

/** @param {ReturnType<typeof createEmptyStorage>} storage */
function createServer(storage) {
  return new PersonalDataServer({
    actorStorage: /** @type {any} */ (storage),
    sharedStorage: /** @type {any} */ ({}),
    blobs: /** @type {any} */ ({}),
    jwtSecret: JWT_SECRET,
    hostname: HOSTNAME,
    password: PASSWORD,
  });
}

/**
 * Encode a JWT with the given claims. The signature is a placeholder; the
 * server validates claims but not the signature for migration tokens.
 * @param {Record<string, any>} claims
 */
function encodeToken(claims) {
  const b64url = (/** @type {object} */ obj) =>
    Buffer.from(JSON.stringify(obj))
      .toString('base64')
      .replace(/\+/g, '-')
      .replace(/\//g, '_')
      .replace(/=+$/, '');
  return `${b64url({ alg: 'ES256', typ: 'JWT' })}.${b64url(claims)}.sig`;
}

function serviceToken(overrides = {}) {
  return encodeToken({
    iss: MIGRATING_DID,
    aud: `did:web:${HOSTNAME}`,
    lxm: 'com.atproto.server.createAccount',
    exp: Math.floor(Date.now() / 1000) + 300,
    ...overrides,
  });
}

/**
 * @param {PersonalDataServer} pds
 * @param {{body: Record<string, unknown>, headers?: Record<string, string>}} opts
 */
function createAccount(pds, { body, headers = {} }) {
  return pds.fetch(
    new Request(`https://${HOSTNAME}/xrpc/com.atproto.server.createAccount`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json', ...headers },
      body: JSON.stringify(body),
    }),
  );
}

describe('createAccount (migration destination)', () => {
  /** @type {ReturnType<typeof createEmptyStorage>} */
  let storage;
  /** @type {PersonalDataServer} */
  let pds;

  beforeEach(() => {
    storage = createEmptyStorage();
    pds = createServer(storage);
  });

  it('adopts an existing DID authorized by a service token', async () => {
    const res = await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: HANDLE },
      headers: { Authorization: `Bearer ${serviceToken()}` },
    });
    expect(res.status).toBe(200);
    const data = await res.json();
    expect(data.did).toBe(MIGRATING_DID);
    expect(data.handle).toBe(HANDLE);
    expect(data.accessJwt).toBeTruthy();
    expect(data.refreshJwt).toBeTruthy();

    // The access token names the adopted DID.
    const payload = await verifyAccessJwt(data.accessJwt, JWT_SECRET);
    expect(payload.sub).toBe(MIGRATING_DID);

    // The account lands deactivated with a freshly minted signing key.
    expect(storage._state.status).toBe('deactivated');
    expect(storage._state.did).toBe(MIGRATING_DID);
    expect(storage._state.handle).toBe(HANDLE);
    expect(storage._state.privateKey).toBeInstanceOf(Uint8Array);
  });

  it('authorizes with the configured password as a manual fallback', async () => {
    const res = await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: HANDLE, password: PASSWORD },
    });
    expect(res.status).toBe(200);
    expect(storage._state.did).toBe(MIGRATING_DID);
  });

  it('rejects a service token whose audience is another server', async () => {
    const res = await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: HANDLE },
      headers: {
        Authorization: `Bearer ${serviceToken({ aud: 'did:web:other.example' })}`,
      },
    });
    expect(res.status).toBe(401);
    expect(storage._state.did).toBeNull();
  });

  it('rejects a service token issued by a different DID', async () => {
    const res = await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: HANDLE },
      headers: {
        Authorization: `Bearer ${serviceToken({ iss: 'did:plc:someoneelse00000000000000' })}`,
      },
    });
    expect(res.status).toBe(401);
  });

  it('rejects an expired service token', async () => {
    const res = await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: HANDLE },
      headers: {
        Authorization: `Bearer ${serviceToken({ exp: Math.floor(Date.now() / 1000) - 1 })}`,
      },
    });
    expect(res.status).toBe(401);
  });

  it('rejects a service token scoped to a different method', async () => {
    const res = await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: HANDLE },
      headers: {
        Authorization: `Bearer ${serviceToken({ lxm: 'com.atproto.repo.createRecord' })}`,
      },
    });
    expect(res.status).toBe(401);
  });

  it('refuses without any authorization', async () => {
    const res = await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: HANDLE },
    });
    expect(res.status).toBe(401);
    expect(storage._state.did).toBeNull();
  });

  it('requires an existing DID (migration only)', async () => {
    const res = await createAccount(pds, {
      body: { handle: HANDLE, password: PASSWORD },
    });
    expect(res.status).toBe(400);
  });

  it('requires a valid handle', async () => {
    const res = await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: 'not a handle', password: PASSWORD },
    });
    expect(res.status).toBe(400);
  });

  it('refuses when the server already hosts an account', async () => {
    storage._state.did = 'did:plc:alreadyhereaccount000000000';
    const res = await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: HANDLE },
      headers: { Authorization: `Bearer ${serviceToken()}` },
    });
    expect(res.status).toBe(400);
  });

  it('the returned session can sign in via createSession', async () => {
    await createAccount(pds, {
      body: { did: MIGRATING_DID, handle: HANDLE, password: PASSWORD },
    });
    const res = await pds.fetch(
      new Request(`https://${HOSTNAME}/xrpc/com.atproto.server.createSession`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ identifier: HANDLE, password: PASSWORD }),
      }),
    );
    expect(res.status).toBe(200);
    const data = await res.json();
    expect(data.did).toBe(MIGRATING_DID);
  });
});

describe('describeServer service identity', () => {
  it('falls back to a did:web identity before an account exists', async () => {
    const pds = createServer(createEmptyStorage());
    const res = await pds.fetch(
      new Request(`https://${HOSTNAME}/xrpc/com.atproto.server.describeServer`),
    );
    const data = await res.json();
    expect(data.did).toBe(`did:web:${HOSTNAME}`);
  });
});

describe('createAccount with signature verification', () => {
  /**
   * Build a server that verifies token signatures, resolving the issuer's DID
   * to the given document. A P-256 verifier suffices for these tests; the
   * secp256k1 path is covered where @noble is available (see the node package).
   * @param {ReturnType<typeof createEmptyStorage>} storage
   * @param {import('../packages/core/src/ports.js').DidResolverPort} didResolver
   */
  function createVerifyingServer(storage, didResolver) {
    return new PersonalDataServer({
      actorStorage: /** @type {any} */ (storage),
      sharedStorage: /** @type {any} */ ({}),
      blobs: /** @type {any} */ ({}),
      jwtSecret: JWT_SECRET,
      hostname: HOSTNAME,
      password: PASSWORD,
      verifier: createVerifier(),
      didResolver,
    });
  }

  /**
   * A DID document publishing `didKey` as the account's atproto key.
   * @param {string} did
   * @param {string} didKey
   */
  function didDoc(did, didKey) {
    return {
      id: did,
      verificationMethod: [
        {
          id: `${did}#atproto`,
          type: 'Multikey',
          controller: did,
          publicKeyMultibase: didKey.replace('did:key:', ''),
        },
      ],
    };
  }

  /** Generate a P-256 keypair and its did:key. */
  async function makeKey() {
    const { privateKey, publicKey } = await generateKeyPair();
    const signingKey = await importPrivateKey(privateKey);
    return { signingKey, didKey: publicKeyToDidKey(publicKey) };
  }

  /**
   * @param {PersonalDataServer} pds
   * @param {string} token
   */
  function post(pds, token) {
    return pds.fetch(
      new Request(`https://${HOSTNAME}/xrpc/com.atproto.server.createAccount`, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          Authorization: `Bearer ${token}`,
        },
        body: JSON.stringify({ did: MIGRATING_DID, handle: HANDLE }),
      }),
    );
  }

  it('accepts a token whose signature matches the published atproto key', async () => {
    const { signingKey, didKey } = await makeKey();
    const token = await createServiceJwt({
      iss: MIGRATING_DID,
      aud: `did:web:${HOSTNAME}`,
      lxm: 'com.atproto.server.createAccount',
      signingKey,
    });
    const storage = createEmptyStorage();
    const pds = createVerifyingServer(storage, async () =>
      didDoc(MIGRATING_DID, didKey),
    );

    const res = await post(pds, token);
    expect(res.status).toBe(200);
    expect(storage._state.did).toBe(MIGRATING_DID);
  });

  it('reads past a held DID document', async () => {
    // Adopting a DID is worth a fresh read: a held document can name a key the
    // account has rotated away from, and no longer controls.
    const { signingKey, didKey } = await makeKey();
    const token = await createServiceJwt({
      iss: MIGRATING_DID,
      aud: `did:web:${HOSTNAME}`,
      lxm: 'com.atproto.server.createAccount',
      signingKey,
    });
    /** @type {Array<{forceRefresh?: boolean}|undefined>} */
    const reads = [];
    const pds = createVerifyingServer(
      createEmptyStorage(),
      async (_did, opts) => {
        reads.push(opts);
        return didDoc(MIGRATING_DID, didKey);
      },
    );

    expect((await post(pds, token)).status).toBe(200);
    expect(reads).toEqual([{ forceRefresh: true }]);
  });

  it('rejects a token signed by a key other than the published one', async () => {
    // Signed by an attacker key; the DID document advertises a different key.
    const attacker = await makeKey();
    const published = await makeKey();
    const token = await createServiceJwt({
      iss: MIGRATING_DID,
      aud: `did:web:${HOSTNAME}`,
      lxm: 'com.atproto.server.createAccount',
      signingKey: attacker.signingKey,
    });
    const storage = createEmptyStorage();
    const pds = createVerifyingServer(storage, async () =>
      didDoc(MIGRATING_DID, published.didKey),
    );

    const res = await post(pds, token);
    expect(res.status).toBe(401);
    expect(storage._state.did).toBeNull();
  });

  /**
   * A service token with the header written by hand, so a test can stamp an
   * `alg` the signing key does not produce and still sign it correctly.
   * @param {string} alg
   * @param {import('../packages/core/src/crypto.js').SigningKey} signingKey
   */
  async function serviceJwtWithAlg(alg, signingKey) {
    const encode = (/** @type {object} */ o) =>
      base64UrlEncode(new TextEncoder().encode(JSON.stringify(o)));
    const now = Math.floor(Date.now() / 1000);
    const signingInput = `${encode({ typ: 'JWT', alg })}.${encode({
      iss: MIGRATING_DID,
      aud: `did:web:${HOSTNAME}`,
      lxm: 'com.atproto.server.createAccount',
      iat: now,
      exp: now + 60,
    })}`;
    const sig = await sign(signingKey, new TextEncoder().encode(signingInput));
    return `${signingInput}.${base64UrlEncode(sig)}`;
  }

  it('accepts a hand-built token stamped ES256 by a P-256 key', async () => {
    const { signingKey, didKey } = await makeKey();
    const storage = createEmptyStorage();
    const pds = createVerifyingServer(storage, async () =>
      didDoc(MIGRATING_DID, didKey),
    );

    const res = await post(pds, await serviceJwtWithAlg('ES256', signingKey));
    expect(res.status).toBe(200);
    expect(storage._state.did).toBe(MIGRATING_DID);
  });

  it('rejects a token whose alg the published key cannot produce', async () => {
    // The same signature, stamped ES256K. The key is P-256, and the reference
    // implementation refuses the mismatch before it compares bytes.
    const { signingKey, didKey } = await makeKey();
    const storage = createEmptyStorage();
    const pds = createVerifyingServer(storage, async () =>
      didDoc(MIGRATING_DID, didKey),
    );

    const res = await post(pds, await serviceJwtWithAlg('ES256K', signingKey));
    expect(res.status).toBe(401);
    expect(storage._state.did).toBeNull();
  });

  it('rejects when the issuer DID cannot be resolved', async () => {
    const { signingKey } = await makeKey();
    const token = await createServiceJwt({
      iss: MIGRATING_DID,
      aud: `did:web:${HOSTNAME}`,
      lxm: 'com.atproto.server.createAccount',
      signingKey,
    });
    const storage = createEmptyStorage();
    const pds = createVerifyingServer(storage, async () => null);

    const res = await post(pds, token);
    expect(res.status).toBe(401);
  });

  it('rejects when the DID document has no atproto key', async () => {
    const { signingKey } = await makeKey();
    const token = await createServiceJwt({
      iss: MIGRATING_DID,
      aud: `did:web:${HOSTNAME}`,
      lxm: 'com.atproto.server.createAccount',
      signingKey,
    });
    const storage = createEmptyStorage();
    const pds = createVerifyingServer(storage, async () => ({
      id: MIGRATING_DID,
      verificationMethod: [],
    }));

    const res = await post(pds, token);
    expect(res.status).toBe(401);
  });

  it('rejects a tampered signature segment', async () => {
    const { signingKey, didKey } = await makeKey();
    const token = await createServiceJwt({
      iss: MIGRATING_DID,
      aud: `did:web:${HOSTNAME}`,
      lxm: 'com.atproto.server.createAccount',
      signingKey,
    });
    // Corrupt the first byte of the signature. Flipping a base64url *character*
    // is unreliable — the last char carries padding bits a flip may not change —
    // so decode, mutate a real byte, and re-encode.
    const parts = token.split('.');
    const sigBytes = base64UrlDecode(parts[2]);
    sigBytes[0] ^= 0xff;
    parts[2] = base64UrlEncode(sigBytes);
    const storage = createEmptyStorage();
    const pds = createVerifyingServer(storage, async () =>
      didDoc(MIGRATING_DID, didKey),
    );

    const res = await post(pds, parts.join('.'));
    expect(res.status).toBe(401);
  });
});