/**
 * Binary payload assembly and extraction (spec §2, §10.2).
 */

import { describe, expect, it } from "vitest";

import {
  AUTH_TAG_LEN,
  buildHeader,
  CHECKSUM_LEN,
  computeBlake3,
  concatBytes,
  constantTimeEqual,
  extractFields,
  finalisePayload,
  HEADER_LEN,
  MIN_PAYLOAD_LEN,
  NONCE_LEN,
  PAYLOAD_FLAGS,
  PAYLOAD_VERSION,
  remainingBitsForCiphertextLen,
  SALT_LEN,
  type HeaderFields,
} from "./cjkPayload";

function makeFields(overrides: Partial<HeaderFields> = {}): HeaderFields {
  return {
    entryCount: 1,
    argon2Memory: 65536,
    argon2Time: 3,
    argon2Parallel: 4,
    salt: new Uint8Array(SALT_LEN).fill(0x42),
    remainingBitsValue: 0,
    ...overrides,
  };
}

describe("buildHeader", () => {
  it("writes the exact byte layout from §2.1", () => {
    const ciphertextLen = 128;
    const rb = remainingBitsForCiphertextLen(ciphertextLen);

    const header = buildHeader(
      makeFields({ remainingBitsValue: rb, entryCount: 1 })
    );

    expect(header.length).toBe(HEADER_LEN);
    const view = new DataView(
      header.buffer,
      header.byteOffset,
      header.byteLength
    );

    expect(header[0]).toBe(PAYLOAD_VERSION);
    expect(view.getUint16(1, true)).toBe(1); // entry_count
    expect(header[3]).toBe(PAYLOAD_FLAGS);
    expect(view.getUint32(4, true)).toBe(65536); // argon2_memory
    expect(view.getUint32(8, true)).toBe(3); // argon2_time
    expect(header[12]).toBe(4); // argon2_parallel
    expect(view.getUint16(13, true)).toBe(rb);
    for (let i = 0; i < SALT_LEN; i++) {
      expect(header[15 + i]).toBe(0x42);
    }
  });

  it("rejects wrongly-sized salts", () => {
    expect(() =>
      buildHeader(makeFields({ salt: new Uint8Array(8) }))
    ).toThrow(/SALT_LEN/);
  });

  it("rejects entryCounts that overflow uint16", () => {
    expect(() =>
      buildHeader(makeFields({ entryCount: 70000 }))
    ).toThrow(/uint16/);
  });
});

describe("finalisePayload + extractFields round-trip", () => {
  it("assembles header|nonce|ct|tag|checksum and extracts each field", async () => {
    const ciphertext = new Uint8Array(64);
    crypto.getRandomValues(ciphertext);
    const authTag = new Uint8Array(AUTH_TAG_LEN).fill(0xAA);
    const nonce = new Uint8Array(NONCE_LEN).fill(0x13);
    const rb = remainingBitsForCiphertextLen(ciphertext.length);

    const header = buildHeader(
      makeFields({ entryCount: 2, remainingBitsValue: rb })
    );

    const full = await finalisePayload(header, nonce, ciphertext, authTag);

    // Length sanity
    expect(full.length).toBe(
      HEADER_LEN + NONCE_LEN + ciphertext.length + AUTH_TAG_LEN + CHECKSUM_LEN
    );

    const extracted = extractFields(full);

    expect(extracted.entryCount).toBe(2);
    expect(extracted.argon2Memory).toBe(65536);
    expect(extracted.argon2Time).toBe(3);
    expect(extracted.argon2Parallel).toBe(4);
    expect(extracted.remainingBitsValue).toBe(rb);
    expect(extracted.salt).toEqual(new Uint8Array(SALT_LEN).fill(0x42));
    expect(extracted.nonce).toEqual(nonce);
    expect(extracted.ciphertext).toEqual(ciphertext);
    expect(extracted.authTag).toEqual(authTag);

    // Stored checksum must match a fresh BLAKE3 over the payload-minus-checksum
    const recomputed = await computeBlake3(extracted.payloadForChecksum);
    expect(constantTimeEqual(recomputed, extracted.storedChecksum)).toBe(true);
  });

  it("rejects a truncated payload", () => {
    const tiny = new Uint8Array(MIN_PAYLOAD_LEN - 1);
    expect(() => extractFields(tiny)).toThrow(/PAYLOAD_TOO_SHORT/);
  });

  it("rejects an unsupported version byte", async () => {
    const ciphertext = new Uint8Array(16);
    const authTag = new Uint8Array(AUTH_TAG_LEN);
    const nonce = new Uint8Array(NONCE_LEN);
    const rb = remainingBitsForCiphertextLen(ciphertext.length);
    const header = buildHeader(makeFields({ remainingBitsValue: rb }));

    const full = await finalisePayload(header, nonce, ciphertext, authTag);
    full[0] = 0x99; // corrupt version byte
    expect(() => extractFields(full)).toThrow(/UNSUPPORTED_VERSION/);
  });

  it("rejects unsupported flags", async () => {
    const ciphertext = new Uint8Array(16);
    const authTag = new Uint8Array(AUTH_TAG_LEN);
    const nonce = new Uint8Array(NONCE_LEN);
    const rb = remainingBitsForCiphertextLen(ciphertext.length);
    const header = buildHeader(makeFields({ remainingBitsValue: rb }));

    const full = await finalisePayload(header, nonce, ciphertext, authTag);
    full[3] = 0x05; // corrupt flags byte
    expect(() => extractFields(full)).toThrow(/UNSUPPORTED_FLAGS/);
  });
});

describe("constantTimeEqual", () => {
  it("returns true for identical buffers", () => {
    const a = new Uint8Array([1, 2, 3, 4]);
    const b = new Uint8Array([1, 2, 3, 4]);
    expect(constantTimeEqual(a, b)).toBe(true);
  });

  it("returns false when any byte differs", () => {
    const a = new Uint8Array([1, 2, 3, 4]);
    const b = new Uint8Array([1, 2, 3, 5]);
    expect(constantTimeEqual(a, b)).toBe(false);
  });

  it("returns false for different lengths", () => {
    expect(
      constantTimeEqual(new Uint8Array([1, 2]), new Uint8Array([1, 2, 3]))
    ).toBe(false);
  });
});

describe("concatBytes", () => {
  it("concatenates in order", () => {
    const out = concatBytes([
      new Uint8Array([1, 2]),
      new Uint8Array([3]),
      new Uint8Array([4, 5, 6]),
    ]);
    expect(out).toEqual(new Uint8Array([1, 2, 3, 4, 5, 6]));
  });
});