/**
 * Binary payload layout for a CJK phrase (spec §2).
 *
 *   [31 bytes: header]
 *     [1]  version          = 0x02
 *     [2]  entry_count      (uint16 LE)
 *     [1]  flags            = 0x01 (zstd + AES-256-GCM)
 *     [4]  argon2_memory    (uint32 LE, KiB)
 *     [4]  argon2_time      (uint32 LE)
 *     [1]  argon2_parallel  (uint8)
 *     [2]  remaining_bits   (uint16 LE)
 *     [16] salt             (CSPRNG, per-export)
 *
 *   [12 bytes: nonce]       (CSPRNG, per-export)
 *   [N bytes: ciphertext]   (AES-256-GCM)
 *   [16 bytes: auth_tag]    (AES-256-GCM tag)
 *   [32 bytes: checksum]    (BLAKE3 of everything above)
 *
 * The checksum covers every byte of the payload except itself. It is
 * verified before the user is prompted for the passphrase, so a corrupt or
 * truncated phrase fails fast with a clear error.
 */

import { blake3 } from "hash-wasm";
import { remainingBits } from "./cjk";

export const PAYLOAD_VERSION = 0x03;
export const PAYLOAD_FLAGS = 0x01; // zstd + AES-256-GCM

export const HEADER_LEN = 31;
export const NONCE_LEN = 12;
export const AUTH_TAG_LEN = 16;
export const CHECKSUM_LEN = 32;
export const SALT_LEN = 16;

/** Minimum valid payload: header + nonce + 0-byte ciphertext + tag + checksum. */
export const MIN_PAYLOAD_LEN = HEADER_LEN + NONCE_LEN + AUTH_TAG_LEN + CHECKSUM_LEN;

// ---------------------------------------------------------------------------
// Header-driven field positions, so the rest of the module has a single
// source of truth for where things live.
// ---------------------------------------------------------------------------

const OFFSET_VERSION = 0;
const OFFSET_ENTRY_COUNT = 1;
const OFFSET_FLAGS = 3;
const OFFSET_ARGON2_MEMORY = 4;
const OFFSET_ARGON2_TIME = 8;
const OFFSET_ARGON2_PARALLEL = 12;
const OFFSET_REMAINING_BITS = 13;
const OFFSET_SALT = 15;

// ---------------------------------------------------------------------------
// Assembly
// ---------------------------------------------------------------------------

export interface HeaderFields {
  entryCount: number;
  argon2Memory: number;
  argon2Time: number;
  argon2Parallel: number;
  salt: Uint8Array;
  remainingBitsValue: number;
}

/**
 * Build the 31-byte header exactly as specified in §2.1.
 */
export function buildHeader(fields: HeaderFields): Uint8Array {
  if (fields.salt.length !== SALT_LEN) {
    throw new Error(`SALT_LEN must be ${SALT_LEN}, got ${fields.salt.length}`);
  }
  if (fields.entryCount > 0xffff) {
    throw new Error("entryCount exceeds uint16 range");
  }

  const header = new Uint8Array(HEADER_LEN);
  const view = new DataView(header.buffer);

  header[OFFSET_VERSION] = PAYLOAD_VERSION;
  view.setUint16(OFFSET_ENTRY_COUNT, fields.entryCount, true);
  header[OFFSET_FLAGS] = PAYLOAD_FLAGS;
  view.setUint32(OFFSET_ARGON2_MEMORY, fields.argon2Memory, true);
  view.setUint32(OFFSET_ARGON2_TIME, fields.argon2Time, true);
  header[OFFSET_ARGON2_PARALLEL] = fields.argon2Parallel;
  view.setUint16(OFFSET_REMAINING_BITS, fields.remainingBitsValue, true);
  header.set(fields.salt, OFFSET_SALT);

  return header;
}

/**
 * Compute `remaining_bits` for the fully-assembled binary given just the
 * ciphertext length. This resolves the circular dependency called out in
 * spec §4.1 Step 10: remaining_bits depends on the final byte count, but is
 * itself stored in the header. The total byte count is deterministic once
 * the ciphertext length is known.
 */
export function remainingBitsForCiphertextLen(ciphertextLen: number): number {
  return remainingBits(HEADER_LEN + NONCE_LEN + ciphertextLen + AUTH_TAG_LEN + CHECKSUM_LEN);
}

/**
 * Concatenate header | nonce | ciphertext | auth_tag, and append the BLAKE3
 * checksum. Returns the full binary that the CJK encoder will turn into
 * characters.
 */
export async function finalisePayload(
  header: Uint8Array,
  nonce: Uint8Array,
  ciphertext: Uint8Array,
  authTag: Uint8Array
): Promise<Uint8Array> {
  if (header.length !== HEADER_LEN) throw new Error("header wrong length");
  if (nonce.length !== NONCE_LEN) throw new Error("nonce wrong length");
  if (authTag.length !== AUTH_TAG_LEN) throw new Error("authTag wrong length");

  const withoutChecksum = concatBytes([header, nonce, ciphertext, authTag]);
  const checksum = await computeBlake3(withoutChecksum);

  return concatBytes([withoutChecksum, checksum]);
}

// ---------------------------------------------------------------------------
// Extraction
// ---------------------------------------------------------------------------

export interface ExtractedFields extends HeaderFields {
  nonce: Uint8Array;
  ciphertext: Uint8Array;
  authTag: Uint8Array;
  storedChecksum: Uint8Array;
  /** Bytes over which BLAKE3 should be recomputed (everything except the checksum). */
  payloadForChecksum: Uint8Array;
}

/**
 * Parse the binary blob that came out of `cjkDecode`. Throws for any
 * structural problem — unknown version, unsupported flags, truncated
 * payload — using the exact error codes the import UI surfaces.
 */
export function extractFields(binary: Uint8Array): ExtractedFields {
  if (binary.length < MIN_PAYLOAD_LEN) {
    throw new Error("PAYLOAD_TOO_SHORT");
  }

  const view = new DataView(binary.buffer, binary.byteOffset, binary.byteLength);

  const version = binary[OFFSET_VERSION];
  if (version !== PAYLOAD_VERSION) {
    throw new Error(`UNSUPPORTED_VERSION: ${version}`);
  }

  const flags = binary[OFFSET_FLAGS];
  if (flags !== PAYLOAD_FLAGS) {
    throw new Error(`UNSUPPORTED_FLAGS: ${flags}`);
  }

  const entryCount = view.getUint16(OFFSET_ENTRY_COUNT, true);
  const argon2Memory = view.getUint32(OFFSET_ARGON2_MEMORY, true);
  const argon2Time = view.getUint32(OFFSET_ARGON2_TIME, true);
  const argon2Parallel = binary[OFFSET_ARGON2_PARALLEL];
  const remainingBitsValue = view.getUint16(OFFSET_REMAINING_BITS, true);
  const salt = binary.slice(OFFSET_SALT, OFFSET_SALT + SALT_LEN);

  const nonceStart = HEADER_LEN;
  const nonceEnd = nonceStart + NONCE_LEN;
  const nonce = binary.slice(nonceStart, nonceEnd);

  const checksumStart = binary.length - CHECKSUM_LEN;
  const authTagStart = checksumStart - AUTH_TAG_LEN;

  if (authTagStart < nonceEnd) {
    throw new Error("PAYLOAD_TOO_SHORT");
  }

  const ciphertext = binary.slice(nonceEnd, authTagStart);
  const authTag = binary.slice(authTagStart, checksumStart);
  const storedChecksum = binary.slice(checksumStart);
  const payloadForChecksum = binary.slice(0, checksumStart);

  return {
    entryCount,
    argon2Memory,
    argon2Time,
    argon2Parallel,
    salt,
    remainingBitsValue,
    nonce,
    ciphertext,
    authTag,
    storedChecksum,
    payloadForChecksum,
  };
}

// ---------------------------------------------------------------------------
// BLAKE3
// ---------------------------------------------------------------------------

export async function computeBlake3(data: Uint8Array): Promise<Uint8Array> {
  const hex = await blake3(data);
  return hexToBytes(hex);
}

/**
 * Constant-time equality for Uint8Arrays. Required for checksum verification
 * — an early-exit comparison leaks timing information about how many leading
 * bytes of a tampered checksum matched the computed one.
 */
export function constantTimeEqual(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) return false;
  let diff = 0;
  for (let i = 0; i < a.length; i++) {
    diff |= a[i] ^ b[i];
  }
  return diff === 0;
}

// ---------------------------------------------------------------------------
// Byte utilities
// ---------------------------------------------------------------------------

export function concatBytes(chunks: Uint8Array[]): Uint8Array {
  let total = 0;
  for (const c of chunks) total += c.length;
  const out = new Uint8Array(total);
  let offset = 0;
  for (const c of chunks) {
    out.set(c, offset);
    offset += c.length;
  }
  return out;
}

function hexToBytes(hex: string): Uint8Array {
  const out = new Uint8Array(hex.length / 2);
  for (let i = 0; i < hex.length; i += 2) {
    out[i / 2] = parseInt(hex.slice(i, i + 2), 16);
  }
  return out;
}