import { afterEach, describe, expect, it, vi } from 'vitest';
import {
  type AccessKey,
  decodeAccessKey,
  decodeAccessKeyBody,
  encodeAccessKey,
  marshalAccessKeyBody,
  wipeAccessKey,
} from './accesskey.ts';
import { sha256 } from './bytes.ts';
import { Decoder } from './cbor.ts';
import { MAX_DKK_BODY_LEN } from './framing.ts';
import { arr, b, bn, ext, map, t, u } from './testing/cborhex.ts';
import { expectCode, h, hx, readBytes, readJSON } from './testing/testdata.ts';

const NC = 'ERR_NON_CANONICAL_CBOR';

interface DkkJSON {
  file: string;
  sha256: string;
  credential_id: string;
  capsule_id: string;
  access_type: string;
  access_material: string;
  capsule_digest?: string;
  capsule: string;
  extensions?: { critical: boolean; id: string; version: number; data?: string }[];
}

function body(o: { type?: string; material?: string; cred?: string; ver?: string; crit?: string; non?: string; version?: string; extra?: [number, string][] } = {}): Uint8Array {
  const entries: [number, string][] = [
    [0, t('datekeys-access-key')],
    [1, o.version ?? u(1)],
    [2, o.cred ?? bn(16, 1)],
    [3, bn(16, 2)],
    [4, o.type ?? t('x25519')],
    [5, o.material ?? bn(32, 3)],
  ];
  if (o.ver) entries.push([6, o.ver]);
  if (o.crit) entries.push([7, o.crit]);
  if (o.non) entries.push([8, o.non]);
  entries.push(...(o.extra ?? []));
  return h(map(...entries));
}

function framed(bodyBytes: Uint8Array): Uint8Array {
  const out = new Uint8Array(12 + bodyBytes.length);
  out.set([0x44, 0x4b, 0x4b, 0x31, 1]);
  new DataView(out.buffer).setUint32(8, bodyBytes.length);
  out.set(bodyBytes, 12);
  return out;
}

describe('.dkk fixtures', () => {
  for (const name of ['time_and_key_portable.dkk', 'time_and_key_recipients.dkk', 'time_and_key_portable_extension.dkk']) {
    it(`decodes ${name} and re-encodes it byte for byte`, async () => {
      const fx = readJSON<DkkJSON>(`fixtures/${name}.json`);
      const raw = readBytes(`fixtures/${name}`);
      expect(hx(await sha256(raw))).toBe(fx.sha256);
      const k = decodeAccessKey(raw);
      expect(hx(k.credentialId)).toBe(fx.credential_id);
      expect(hx(k.capsuleId)).toBe(fx.capsule_id);
      expect(k.type).toBe(fx.access_type);
      expect(hx(k.material)).toBe(fx.access_material);
      expect(k.verification && hx(k.verification.capsuleDigest)).toBe(fx.capsule_digest);
      // capsule_digest is the SHA-256 of the exact .dkc bytes.
      expect(hx(await sha256(readBytes(`fixtures/${fx.capsule}`)))).toBe(fx.capsule_digest);
      const exts = [...k.critical.map((e) => ({ ...e, critical: true })), ...k.noncritical.map((e) => ({ ...e, critical: false }))];
      expect(exts.map((e) => ({ critical: e.critical, id: e.id, version: e.version, data: e.data && hx(e.data) }))).toEqual(
        (fx.extensions ?? []).map((e) => ({ critical: e.critical, id: e.id, version: e.version, data: e.data })),
      );
      expect(hx(encodeAccessKey(k))).toBe(hx(raw));
      wipeAccessKey(k);
      expect(k.material.every((x) => x === 0)).toBe(true);
    });
  }
});

describe('DKK1 framing', () => {
  const ok = readBytes('fixtures/time_and_key_portable.dkk');
  const with_ = (i: number, v: number): Uint8Array => {
    const c = ok.slice();
    c[i] = v;
    return c;
  };

  it('checks magic, prelude, limits and length in order', () => {
    expectCode(() => decodeAccessKey(new Uint8Array(0)), 'ERR_INVALID_MAGIC', /^accesskey: ERR_INVALID_MAGIC$/);
    expectCode(() => decodeAccessKey(ok.subarray(0, 3)), 'ERR_INVALID_MAGIC');
    expectCode(() => decodeAccessKey(with_(3, 0x32)), 'ERR_INVALID_MAGIC');
    expectCode(() => decodeAccessKey(ok.subarray(0, 11)), 'ERR_INTEGRITY', /truncated prelude/);
    expectCode(() => decodeAccessKey(with_(4, 2)), 'ERR_UNSUPPORTED_VERSION', /framing version 2/);
    expectCode(() => decodeAccessKey(with_(5, 0x80)), 'ERR_INVALID_FLAGS', /flags 0x80, reserved 0x000/);
    expectCode(() => decodeAccessKey(with_(6, 0xab)), 'ERR_INVALID_FLAGS', /flags 0x0, reserved 0xab00/);
    expectCode(() => decodeAccessKey(with_(7, 1)), 'ERR_INVALID_FLAGS', /reserved 0x001/);
    const big = ok.slice();
    new DataView(big.buffer).setUint32(8, MAX_DKK_BODY_LEN + 1);
    expectCode(() => decodeAccessKey(big), 'ERR_INTEGRITY', /^accesskey: BODY_LEN 16777217 outside 1..16777216: ERR_INTEGRITY$/);
    // Spec §40, §57: no empty frame holds a valid body, so BODY_LEN 0 is a
    // framing error, after FLAGS and RESERVED.
    const empty = ok.slice(0, 12);
    new DataView(empty.buffer).setUint32(8, 0);
    expectCode(() => decodeAccessKey(empty), 'ERR_INTEGRITY', /^accesskey: BODY_LEN 0 outside 1..16777216: ERR_INTEGRITY$/);
    empty[5] = 1;
    expectCode(() => decodeAccessKey(empty), 'ERR_INVALID_FLAGS', /^accesskey: flags 0x1, reserved 0x000: ERR_INVALID_FLAGS$/);
    expectCode(() => decodeAccessKey(ok.subarray(0, ok.length - 1)), 'ERR_INTEGRITY', /truncated body/);
    expectCode(() => decodeAccessKey(new Uint8Array([...ok, 0])), 'ERR_INTEGRITY', /data after BODY_CBOR/);
  });
});

describe('BODY_CBOR', () => {
  it('accepts a minimal body and the optional fields', () => {
    const k = decodeAccessKey(framed(body()));
    expect(k.verification).toBeUndefined();
    expect(decodeAccessKeyBody(body({ ver: map([0, bn(32, 9)]) })).verification?.capsuleDigest).toEqual(new Uint8Array(32).fill(9));
    expect(decodeAccessKeyBody(body({ crit: arr(ext('a')), non: arr(ext('b', 2, b('01'))) })).noncritical).toEqual([
      { id: 'b', version: 2, data: h('01') },
    ]);
  });

  it('checks the layers of spec §69.1 in order: schema head, CDDL, then access_type and access_material', () => {
    expectCode(() => decodeAccessKeyBody(new Uint8Array(MAX_DKK_BODY_LEN + 1)), 'ERR_INTEGRITY');
    const cases: [string, Uint8Array, string, RegExp][] = [
      ['schema version 2', body({ version: u(2) }), 'ERR_UNSUPPORTED_VERSION', /^accesskey: codec: datekeys-access-key schema version 2, want 1: ERR_UNSUPPORTED_VERSION$/],
      ['unknown key 9', body({ extra: [[9, u(0)]] }), NC, /^accesskey: key 9 is not defined: ERR/],
      ['credential_id of 15 bytes', body({ cred: bn(15) }), NC, /^accesskey: key 2: codec: offset 26: a byte string of 15 bytes outside 16\.\.16: ERR/],
      ['credential_id of 15 bytes, access_type y', body({ cred: bn(15), type: t('y') }), NC, /key 2: codec: offset 26/],
      ['access_type as a byte string', body({ type: b('783235353139') }), NC, /^accesskey: key 4: codec: offset 61: a byte string where a text string was expected: ERR/],
      ['empty verification_metadata', body({ ver: map() }), NC, /^accesskey: key 6: empty verification_metadata; an absent one omits key 6: ERR/],
      ['empty verification_metadata, access_type y', body({ ver: map(), type: t('y') }), NC, /empty verification_metadata/],
      ['empty capsule_digest', body({ ver: map([0, b('')]) }), NC, /^accesskey: key 6: capsule_digest: codec: offset 107: a byte string of 0 bytes outside 32\.\.32: ERR/],
      ['capsule_digest of 31 bytes', body({ ver: map([0, bn(31)]) }), NC, /capsule_digest: codec: offset 108: a byte string of 31 bytes outside 32\.\.32/],
      ['verification_metadata key 1', body({ ver: map([1, bn(32)]) }), NC, /^accesskey: key 6: verification_metadata key 1 is not defined: ERR/],
      ['verification_metadata of two keys', body({ ver: map([0, bn(32)], [1, u(0)]) }), NC, /key 6: codec: offset 105: map of 2 entries, at most 1/],
      ['null verification_metadata', body({ ver: 'f6' }), NC, /key 6: codec: offset 104: a float or simple value \(initial byte 0xf6\) is outside the CBOR profile/],
      ['out of order, access_type y', body({ type: t('y'), non: arr(ext('b'), ext('a')) }), NC, /^accesskey: key 8: extension a: array is not in canonical order: ERR/],
      ['empty extension_id', body({ crit: arr(ext('')) }), NC, /^accesskey: key 7: extension: invalid extension_id "": ERR/],
      ['id in both arrays', body({ crit: arr(ext('a')), non: arr(ext('a')) }), NC, /^accesskey: extension a: both critical and noncritical: ERR/],
      ['access_type X25519', body({ type: t('X25519') }), 'ERR_ACCESS_INVALID', /^accesskey: access_type "X25519" is not supported by V1: ERR_ACCESS_INVALID$/],
      ['access_material of 31 bytes', body({ material: bn(31) }), 'ERR_ACCESS_INVALID', /^accesskey: x25519 access_material is 31 bytes, want 32: ERR_ACCESS_INVALID$/],
      // access_type (key 4) before access_material (key 5).
      ['access_type y, material of 31 bytes', body({ type: t('y'), material: bn(31) }), 'ERR_ACCESS_INVALID', /access_type "y" is not supported/],
    ];
    for (const [name, bytes, code, msg] of cases) expectCode(() => decodeAccessKeyBody(bytes), code, msg, name);
  });
});

describe('access_material wiping (the reference clears it on every path)', () => {
  afterEach(() => {
    vi.restoreAllMocks();
  });

  // Every byte string the decoder hands out, to find the copies of the material.
  function spyBstr(): Uint8Array[] {
    const out: Uint8Array[] = [];
    const orig = Decoder.prototype.bstr;
    vi.spyOn(Decoder.prototype, 'bstr').mockImplementation(function (this: Decoder, min: number, max: number) {
      const b = orig.call(this, min, max);
      out.push(b);
      return b;
    });
    return out;
  }
  const materials = (copies: Uint8Array[]): Uint8Array[] => copies.filter((c) => c.length === 32 && c.every((x) => x === 3 || x === 0));

  it('wipes the material when decoding fails after key 5', () => {
    for (const bad of [
      body({ crit: arr() }),
      body({ non: arr(ext('a', 1, b(''))) }),
      body({ ver: map([0, b('')]) }),
      body({ extra: [[9, u(0)]] }),
    ]) {
      const copies = spyBstr();
      expectCode(() => decodeAccessKeyBody(bad), NC);
      const m = materials(copies);
      expect(m).toHaveLength(1);
      expect(m[0]!.every((x) => x === 0)).toBe(true);
      vi.restoreAllMocks();
    }
  });

  it('wipes the material when unmarshal rejects the decoded value', () => {
    const copies = spyBstr();
    vi.spyOn(Decoder.prototype, 'done').mockImplementation(() => {
      throw new Error('rejected');
    });
    expect(() => decodeAccessKeyBody(body())).toThrow('rejected');
    const m = materials(copies);
    expect(m).toHaveLength(1);
    expect(m[0]!.every((x) => x === 0)).toBe(true);
  });

  it('wipes its own copy after a later check fails, and after success', () => {
    for (const input of [body({ cred: bn(15) }), body({ type: t('y') }), body()]) {
      const copies = spyBstr();
      try {
        wipeAccessKey(decodeAccessKeyBody(input));
      } catch {
        // The failures are tested above; only the wiping matters here.
      }
      expect(materials(copies).every((c) => c.every((x) => x === 0))).toBe(true);
      vi.restoreAllMocks();
    }
  });
});

describe('BODY_CBOR size limit', () => {
  const key = (n: number): AccessKey => ({
    credentialId: new Uint8Array(16).fill(1),
    capsuleId: new Uint8Array(16).fill(2),
    type: 'x25519',
    material: new Uint8Array(32).fill(3),
    verification: undefined,
    critical: [],
    noncritical: [{ id: 'a', version: 1, data: new Uint8Array(n).fill(4) }],
  });

  it('accepts a body of exactly MAX_DKK_BODY_LEN bytes and rejects one more', () => {
    const n = 100_000 + MAX_DKK_BODY_LEN - marshalAccessKeyBody(key(100_000)).length;
    const exact = marshalAccessKeyBody(key(n));
    expect(exact.length).toBe(MAX_DKK_BODY_LEN);
    expect(decodeAccessKeyBody(exact).noncritical[0]!.data!.length).toBe(n);
    expect(decodeAccessKey(framed(exact)).noncritical[0]!.data!.length).toBe(n);
    expectCode(() => marshalAccessKeyBody(key(n + 1)), 'ERR_INTEGRITY', /BODY_CBOR of 16777217 bytes exceeds 16777216/);
    expectCode(() => decodeAccessKeyBody(new Uint8Array(MAX_DKK_BODY_LEN)), NC, /offset 0: an unsigned integer where a map was expected/);
    expectCode(() => decodeAccessKeyBody(new Uint8Array(MAX_DKK_BODY_LEN + 1)), 'ERR_INTEGRITY', /BODY_CBOR of 16777217 bytes exceeds/);
    expectCode(() => decodeAccessKey(framed(new Uint8Array(MAX_DKK_BODY_LEN))), NC, /offset 0: an unsigned integer where a map was expected/);
    expectCode(() => decodeAccessKey(framed(new Uint8Array(MAX_DKK_BODY_LEN + 1))), 'ERR_INTEGRITY', /BODY_LEN 16777217 outside 1..16777216/);
  });
});

describe('marshalAccessKeyBody', () => {
  const key: AccessKey = {
    credentialId: new Uint8Array(16).fill(1),
    capsuleId: new Uint8Array(16).fill(2),
    type: 'x25519',
    material: new Uint8Array(32).fill(3),
    verification: { capsuleDigest: new Uint8Array(32).fill(4) },
    critical: [],
    noncritical: [{ id: 'z', version: 1, data: h('00') }],
  };

  it('writes what the reader accepts', () => {
    const k = decodeAccessKeyBody(marshalAccessKeyBody(key));
    expect(k).toEqual(key);
    expect(hx(marshalAccessKeyBody({ ...key, verification: undefined, noncritical: [] }))).toBe(hx(body()));
  });

  it('rejects invalid keys', () => {
    expectCode(() => marshalAccessKeyBody({ ...key, type: 'y' }), 'ERR_ACCESS_INVALID');
    expectCode(() => marshalAccessKeyBody({ ...key, material: new Uint8Array(31) }), 'ERR_ACCESS_INVALID');
    expectCode(() => marshalAccessKeyBody({ ...key, capsuleId: new Uint8Array(15) }), NC);
    expectCode(() => marshalAccessKeyBody({ ...key, verification: { capsuleDigest: new Uint8Array(0) } }), NC, /capsule_digest must be 32 bytes/);
    expectCode(() => marshalAccessKeyBody({ ...key, critical: [{ id: 'z', version: 1, data: undefined }] }), NC, /both critical/);
    expectCode(() => marshalAccessKeyBody({ ...key, critical: [{ id: 'y', version: -1, data: undefined }] }), NC);
    const big = [{ id: 'a', version: 1, data: new Uint8Array(MAX_DKK_BODY_LEN) }];
    expectCode(() => marshalAccessKeyBody({ ...key, noncritical: big }), 'ERR_INTEGRITY', /exceeds 16777216/);
  });

  // Spec §72, extension.CheckWrite in Go's MarshalBody: datekeys.capsule goes only in the noncritical array of a
  // .dkk, with data, and datekeys.note never in a .dkk. The texts of Go on the branch v0.12.
  it('writes the extensions of the specification only where §72 registers them', () => {
    const capsule = { id: 'datekeys.capsule', version: 1, data: h('a0') };
    const misplaced = (id: string, arr: string): string =>
      `accesskey: extension: ${id} version 1 is not registered for ${arr} of .dkk: an encoder must not write it there (spec §72)`;
    expect(() => marshalAccessKeyBody({ ...key, noncritical: [{ id: 'datekeys.note', version: 1, data: h('41') }] })).toThrow(
      new Error(misplaced('datekeys.note', 'noncritical_extensions')),
    );
    expect(() => marshalAccessKeyBody({ ...key, critical: [capsule], noncritical: [] })).toThrow(new Error(misplaced('datekeys.capsule', 'critical_extensions')));
    expectCode(
      () => marshalAccessKeyBody({ ...key, noncritical: [{ ...capsule, data: undefined }] }),
      'ERR_EXTENSION_DATA_INVALID',
      /^accesskey: extension: datekeys\.capsule version 1: datekeys\.capsule without data: ERR_EXTENSION_DATA_INVALID$/,
    );
    // Its data is the locator's, which this library does not check: one byte is enough here.
    expect(decodeAccessKeyBody(marshalAccessKeyBody({ ...key, noncritical: [capsule] })).noncritical).toEqual([capsule]);
    // Another version is the application's own.
    expect(() => marshalAccessKeyBody({ ...key, noncritical: [{ id: 'datekeys.note', version: 2, data: h('41') }] })).not.toThrow();
  });
});