// @pdsjs/spaces/token - the three space JWTs.
//
// A delegation token proves the user delegated to an app; a client attestation
// proves which app is acting; a space credential is what the authority issues in
// exchange, and is what repo hosts actually accept.
//
// They share a wire shape and differ only in who signs, who they address, and
// how long they live — so they are data, not three implementations.

import {
  base64UrlDecode,
  base64UrlEncode,
  bytesToHex,
} from '@pdsjs/core/crypto';
import { parseDidKey } from '@pdsjs/core/verify';

export const SPACE_TOKEN_TYPES = {
  delegation: {
    typ: 'atproto-space-delegation+jwt',
    kid: '#atproto',
    expiresInSec: 60,
    requireAud: true,
    requireCnf: false,
    singleUse: true,
  },
  credential: {
    typ: 'atproto-space-credential+jwt',
    kid: '#atproto',
    // Multi-use across repo hosts until it expires, so it carries no aud. It is
    // bound to the holder's key instead, and each use names its audience in
    // a signature. See http-signature.js.
    expiresInSec: 600,
    requireAud: false,
    requireCnf: true,
    singleUse: false,
  },
  clientAttestation: {
    typ: 'atproto-client-attestation+jwt',
    kid: undefined,
    expiresInSec: 60,
    requireAud: true,
    requireCnf: false,
    singleUse: true,
  },
};

export const CLOCK_SKEW_SEC = 5;

/** The longest a space credential may live, from issue to expiry. */
export const MAX_CREDENTIAL_LIFETIME_SEC = 3600;

/**
 * The account key that signs space tokens and repo commits.
 * @typedef {Object} SpaceSigner
 * @property {(bytes: Uint8Array) => Promise<Uint8Array>} sign
 * @property {'p256'|'secp256k1'} [curve] - p256 when absent
 */

/** The JWT `alg` each curve signs under. */
const JWT_ALG = { p256: 'ES256', secp256k1: 'ES256K' };

/**
 * The JWT `alg` for a signer's curve. A verifier resolves the issuer's key,
 * reads its type, and refuses a header that names another algorithm before it
 * checks the signature. So a secp256k1 key must stamp ES256K.
 * @param {SpaceSigner} signer
 * @returns {string}
 */
export function signerAlg(signer) {
  return JWT_ALG[signer.curve ?? 'p256'];
}

export class SpaceTokenError extends Error {
  /**
   * @param {string} message
   * @param {string} [code] - machine-readable error name for the XRPC response
   */
  constructor(message, code = 'InvalidToken') {
    super(message);
    this.name = 'SpaceTokenError';
    this.code = code;
  }
}

/**
 * Refuse a token whose header names an algorithm the issuer's key cannot
 * produce. The signature would verify anyway, since the curve comes from the
 * resolved key rather than the header. The reference implementation rejects
 * the mismatch, so a token this server accepts has to be one the rest of the
 * ecosystem accepts too.
 *
 * @param {string} didKey
 * @param {string} alg - the token header's `alg`
 * @throws {SpaceTokenError}
 */
export function assertKeyAlg(didKey, alg) {
  let expected;
  try {
    expected = JWT_ALG[parseDidKey(didKey).curve];
  } catch (err) {
    throw new SpaceTokenError(
      `could not read the issuer key: ${err instanceof Error ? err.message : String(err)}`,
      'BadJwtSignature',
    );
  }
  if (alg !== expected) {
    throw new SpaceTokenError(
      `token alg ${alg} does not match the issuer key, which signs ${expected}`,
      'BadJwtSignature',
    );
  }
}

/**
 * @typedef {Object} SpaceTokenPayload
 * @property {string} iss
 * @property {string} sub - the space URI, or the client_id for an attestation
 * @property {string} [aud]
 * @property {number} iat
 * @property {number} exp
 * @property {string} jti
 * @property {{kid: string}} [cnf] - the did:key whose signatures present it
 */

/**
 * @typedef {Object} SpaceTokenHeader
 * @property {string} alg
 * @property {string} typ
 * @property {string} [kid]
 */

/**
 * @param {Record<string, unknown>} json
 * @returns {string}
 */
function jsonToB64Url(json) {
  return base64UrlEncode(new TextEncoder().encode(JSON.stringify(json)));
}

/**
 * @param {string} b64
 * @param {string} part
 * @returns {any}
 */
function decodeJsonPart(b64, part) {
  try {
    return JSON.parse(new TextDecoder().decode(base64UrlDecode(b64)));
  } catch (err) {
    throw new SpaceTokenError(
      `could not parse token ${part}: ${err instanceof Error ? err.message : String(err)}`,
      'BadJwt',
    );
  }
}

/**
 * @param {keyof typeof SPACE_TOKEN_TYPES} type
 * @param {{iss: string, sub: string, aud?: string, keyId?: string, expiresInSec?: number, kid?: string, alg?: string}} opts
 * @param {SpaceSigner} signer
 * @returns {Promise<string>}
 */
export async function createSpaceToken(type, opts, signer) {
  const spec = SPACE_TOKEN_TYPES[type];
  if (spec.requireAud && !opts.aud) {
    throw new SpaceTokenError(`a ${type} token requires an "aud"`);
  }
  if (spec.requireCnf && !opts.keyId) {
    throw new SpaceTokenError(`a ${type} token requires a "keyId"`);
  }
  if (spec.kid && opts.kid) {
    throw new SpaceTokenError(
      `a ${type} token's "kid" is always "${spec.kid}"`,
    );
  }
  const expiresInSec = opts.expiresInSec ?? spec.expiresInSec;
  if (
    type === 'credential' &&
    (!Number.isFinite(expiresInSec) ||
      expiresInSec <= 0 ||
      expiresInSec > MAX_CREDENTIAL_LIFETIME_SEC)
  ) {
    throw new SpaceTokenError('invalid space credential lifetime');
  }

  const iat = Math.floor(Date.now() / 1000);
  /** @type {SpaceTokenHeader} */
  const header = { alg: opts.alg ?? signerAlg(signer), typ: spec.typ };
  const kid = opts.kid ?? spec.kid;
  if (kid) header.kid = kid;

  /** @type {SpaceTokenPayload} */
  const payload = {
    iss: opts.iss,
    sub: opts.sub,
    ...(opts.aud ? { aud: opts.aud } : undefined),
    ...(opts.keyId ? { cnf: { kid: opts.keyId } } : undefined),
    iat,
    exp: iat + expiresInSec,
    jti: bytesToHex(crypto.getRandomValues(new Uint8Array(16))),
  };

  const signingInput = `${jsonToB64Url(header)}.${jsonToB64Url(payload)}`;
  const sig = await signer.sign(new TextEncoder().encode(signingInput));
  return `${signingInput}.${base64UrlEncode(sig)}`;
}

/**
 * Structural validation only, no signature check. A client attestation's key
 * comes from the client's own JWKS rather than a DID document, so attestation.js
 * takes it from here and checks the signature against that.
 *
 * @param {keyof typeof SPACE_TOKEN_TYPES} type
 * @param {string} jwt
 * @returns {{header: SpaceTokenHeader, payload: SpaceTokenPayload, signingInput: Uint8Array, sig: Uint8Array}}
 */
export function parseSpaceToken(type, jwt) {
  const spec = SPACE_TOKEN_TYPES[type];

  const parts = typeof jwt === 'string' ? jwt.split('.') : [];
  if (parts.length !== 3) {
    throw new SpaceTokenError('malformed token: expected 3 parts', 'BadJwt');
  }
  const [headerB64, payloadB64, sigB64] = parts;

  const header = decodeJsonPart(headerB64, 'header');
  const payload = decodeJsonPart(payloadB64, 'payload');

  if (header.typ !== spec.typ) {
    throw new SpaceTokenError(
      `wrong token type: expected "${spec.typ}", got "${header.typ}"`,
      'BadJwtType',
    );
  }
  if (typeof header.alg !== 'string' || !header.alg) {
    throw new SpaceTokenError('missing token "alg"', 'BadJwt');
  }
  if (spec.kid && header.kid !== spec.kid) {
    throw new SpaceTokenError(
      `wrong token "kid": expected "${spec.kid}", got "${header.kid}"`,
      'BadJwt',
    );
  }
  if (!payload.iss)
    throw new SpaceTokenError('missing token "iss"', 'BadJwtIss');
  if (!payload.sub)
    throw new SpaceTokenError('missing token "sub"', 'BadJwtSub');
  if (!Number.isFinite(payload.exp)) {
    throw new SpaceTokenError('missing token "exp"', 'BadJwt');
  }
  if (spec.requireAud && !payload.aud) {
    throw new SpaceTokenError('missing token "aud"', 'BadJwtAudience');
  }
  if (
    spec.requireCnf &&
    (typeof payload.cnf?.kid !== 'string' ||
      !payload.cnf.kid.startsWith('did:'))
  ) {
    throw new SpaceTokenError('missing token "cnf.kid"', 'BadJwtCnf');
  }
  // A single-use token is spent by its jti, and a credential is revoked by
  // it, so a token without one could be neither.
  if (typeof payload.jti !== 'string' || !payload.jti) {
    throw new SpaceTokenError(`a ${type} token requires a "jti"`, 'BadJwt');
  }
  if (
    type === 'credential' &&
    (!Number.isFinite(payload.iat) ||
      payload.exp <= payload.iat ||
      payload.exp - payload.iat > MAX_CREDENTIAL_LIFETIME_SEC)
  ) {
    throw new SpaceTokenError('invalid space credential lifetime', 'BadJwt');
  }
  if (type === 'clientAttestation' && payload.iss !== payload.sub) {
    throw new SpaceTokenError(
      'client attestation "iss" and "sub" must both be the client_id',
      'BadJwtIss',
    );
  }

  return {
    header,
    payload,
    signingInput: new TextEncoder().encode(`${headerB64}.${payloadB64}`),
    sig: base64UrlDecode(sigB64),
  };
}

/**
 * @param {keyof typeof SPACE_TOKEN_TYPES} type
 * @param {string} jwt
 * @param {Object} opts
 * @param {(iss: string, forceRefresh?: boolean) => Promise<string>} opts.getSigningKey
 *   - resolves the issuer to its `#atproto` did:key. Given `forceRefresh` to
 *   read past a held answer after a signature fails.
 * @param {import('@pdsjs/core/ports').SignatureVerifierPort} opts.verifier
 * @param {string} [opts.aud] - required audience, when the caller knows it
 * @param {string} [opts.sub] - required subject, when the caller knows it
 * @returns {Promise<{header: SpaceTokenHeader, payload: SpaceTokenPayload}>}
 */
export async function verifySpaceToken(type, jwt, opts) {
  const { header, payload, signingInput, sig } = parseSpaceToken(type, jwt);

  const now = Math.floor(Date.now() / 1000);
  if (type === 'credential' && payload.iat > now + CLOCK_SKEW_SEC) {
    throw new SpaceTokenError(
      'space credential issued in the future',
      'BadJwt',
    );
  }
  if (now - CLOCK_SKEW_SEC >= payload.exp) {
    throw new SpaceTokenError('token expired', 'JwtExpired');
  }
  if (opts.aud !== undefined && payload.aud !== opts.aud) {
    throw new SpaceTokenError(
      'token audience does not match this service',
      'BadJwtAudience',
    );
  }
  if (opts.sub !== undefined && payload.sub !== opts.sub) {
    throw new SpaceTokenError(
      'token subject does not match the requested space',
      'BadJwtSub',
    );
  }

  /**
   * @param {string} didKey
   * @returns {Promise<boolean>}
   */
  const matchesSignature = async (didKey) => {
    assertKeyAlg(didKey, header.alg);
    try {
      return await opts.verifier.verify(didKey, signingInput, sig);
    } catch (err) {
      throw new SpaceTokenError(
        `could not verify token signature: ${err instanceof Error ? err.message : String(err)}`,
        'BadJwtSignature',
      );
    }
  };

  const didKey = await opts.getSigningKey(payload.iss);
  if (await matchesSignature(didKey)) return { header, payload };

  // A resolver that holds its answers can name a key the issuer has rotated
  // away from, and the token is signed by the current one.
  const freshDidKey = await opts.getSigningKey(payload.iss, true);
  if (freshDidKey !== didKey && (await matchesSignature(freshDidKey))) {
    return { header, payload };
  }

  throw new SpaceTokenError('invalid token signature', 'BadJwtSignature');
}