/**
 * COSE Merkle Tree Proofs-style batch receipts for signed Remote ID logs.
 *
 * Reference: draft-ietf-cose-merkle-tree-proofs-18 and RFC 9052 (COSE).
 * We implement the core tree mechanics the draft standardises — leaf
 * hashing, bottom-up pairwise SHA-256, odd-node duplication, and Merkle
 * inclusion proofs — plus an optional COSE_Sign1 envelope per RFC 9052
 * for the tree head + inclusion proof bundle. Consumers today get a
 * root that's byte-identical to the draft's tree construction for the
 * same leaves, and can additionally produce a signed receipt.
 *
 * Terminology (draft-18):
 *   leaf_hash   — SHA-256(signature_bytes || keyId_ascii || uint64BE timestampMs).
 *                 Binds each leaf to the signer, timestamp, and full ed25519
 *                 signature so leaf equality is independent of the decoded
 *                 message representation.
 *   tree_head   — { tree_size, root_hash, timestamp }. The signed summary.
 *   audit_path  — sibling hashes (with direction) walking the leaf up to
 *                 the root. Also referred to as `siblings` in the legacy
 *                 API — both views are exposed for back-compat.
 *
 * Legacy names (InclusionProof.index, .leafCount, .siblings, .root) remain
 * exported and continue to work. IETF-aligned aliases are added additively.
 */

import { createHash, sign as nodeSign, verify as nodeVerify, type KeyObject } from 'node:crypto';
import { Encoder } from 'cbor-x';
import type { SignedRemoteIdLog } from './types.js';

/** Hex-encoded SHA-256 digest. */
export type HexDigest = string;

/** Output of buildMerkleRoot — root + inclusion-proof generator. */
export interface MerkleBatch {
  /** Hex SHA-256 root. */
  root: HexDigest;
  /** Number of leaves in the tree. */
  leafCount: number;
  /** Leaves in canonical hex form — same order as the input. */
  leaves: readonly HexDigest[];
  /**
   * Build an inclusion proof for the leaf at `index`. Returns the
   * sibling hashes walking up to the root.
   */
  inclusion_proof: (index: number) => InclusionProof;
}

/**
 * Audit-path step — one sibling hash plus the side it sits on.
 * `side === 'L'` means the sibling is to the left of the current node
 * at that level; `side === 'R'` means to the right.
 */
export interface AuditPathStep {
  hash: HexDigest;
  side: 'L' | 'R';
}

/**
 * Per-leaf inclusion proof.
 *
 * Both the legacy field names (`index`, `leafCount`, `siblings`, `root`)
 * and IETF draft-18 aliases (`leaf_index`, `tree_size`, `audit_path`,
 * `root_hash`) are present on every instance. Writers mutate the legacy
 * names and the IETF aliases are kept consistent at construction time;
 * do not mutate fields on a returned proof.
 */
export interface InclusionProof {
  /** 0-based index of the leaf in the original input order. */
  index: number;
  /** IETF draft-18 alias for {@link InclusionProof.index}. */
  leaf_index: number;
  /** Total leaf count (needed for verification path reconstruction). */
  leafCount: number;
  /** IETF draft-18 alias for {@link InclusionProof.leafCount}. */
  tree_size: number;
  /**
   * Sibling hashes from the leaf's level up to (but not including) the
   * root. Each entry is paired with its side: 'L' means the sibling sits
   * to the left of the current node; 'R' means to the right.
   */
  siblings: readonly AuditPathStep[];
  /** IETF draft-18 alias for {@link InclusionProof.siblings}. */
  audit_path: readonly AuditPathStep[];
  /** Convenience copy of the root. */
  root: HexDigest;
  /** IETF draft-18 alias for {@link InclusionProof.root}. */
  root_hash: HexDigest;
}

/**
 * @deprecated Use {@link InclusionProof} (same shape, IETF-aligned aliases added).
 *             Retained for back-compat; will remain a type alias indefinitely.
 */
export type COSE_MerkleInclusionProof = InclusionProof;

/**
 * IETF draft-18 `tree_head` — the signed summary of a Merkle tree.
 * `tree_size` is the number of leaves; `root_hash` is the SHA-256 root
 * (bytes, not hex, to match RFC 9052 conventions when wrapped in COSE);
 * `timestamp` is Unix epoch milliseconds at the moment the head was built.
 */
export interface TreeHead {
  tree_size: number;
  root_hash: Uint8Array;
  timestamp: number;
}

// ============================================================
// HASHING
// ============================================================

/** SHA-256 of the given buffers concatenated, hex-encoded. */
function sha256Hex(...parts: (Buffer | Uint8Array)[]): HexDigest {
  const h = createHash('sha256');
  for (const p of parts) h.update(p);
  return h.digest('hex');
}

/** Compute the canonical leaf hash for a signed Remote ID log. */
export function leafHash(log: SignedRemoteIdLog): HexDigest {
  const sig = Buffer.from(log.signature, 'hex');
  const kid = Buffer.from(log.keyId, 'ascii');
  const ts = Buffer.alloc(8);
  ts.writeBigUInt64BE(BigInt(log.timestampMs), 0);
  return sha256Hex(sig, kid, ts);
}

/** Combine two child hashes into their parent hash. */
function parentHash(left: HexDigest, right: HexDigest): HexDigest {
  return sha256Hex(Buffer.from(left, 'hex'), Buffer.from(right, 'hex'));
}

// ============================================================
// TREE CONSTRUCTION
// ============================================================

/**
 * Build a bottom-up Merkle tree of SHA-256 leaves. Odd nodes at any
 * level are duplicated (RFC 6962 / draft-ietf-cose-merkle-tree-proofs
 * convention) so the root is always well-defined.
 *
 * Returns a MerkleBatch with an inclusion-proof generator. The root is
 * deterministic for the same input order.
 */
export function buildMerkleRoot(logs: readonly SignedRemoteIdLog[]): MerkleBatch {
  if (logs.length === 0) {
    throw new Error('buildMerkleRoot requires at least one signed log');
  }
  const leaves = logs.map(leafHash);
  return buildMerkleRootFromLeaves(leaves);
}

/**
 * Lower-level primitive: build a tree from pre-computed leaf hashes.
 * Useful when you want to batch heterogeneous artefacts together.
 */
export function buildMerkleRootFromLeaves(
  leaves: readonly HexDigest[],
): MerkleBatch {
  if (leaves.length === 0) {
    throw new Error('buildMerkleRootFromLeaves requires at least one leaf');
  }

  // Build every level so we can answer inclusion queries without recomputing.
  const levels: HexDigest[][] = [leaves.slice()];
  while (levels[levels.length - 1]!.length > 1) {
    const prev = levels[levels.length - 1]!;
    const next: HexDigest[] = [];
    for (let i = 0; i < prev.length; i += 2) {
      const left = prev[i]!;
      const right = i + 1 < prev.length ? prev[i + 1]! : left; // duplicate odd
      next.push(parentHash(left, right));
    }
    levels.push(next);
  }

  const root = levels[levels.length - 1]![0]!;
  const leafCount = leaves.length;

  return {
    root,
    leafCount,
    leaves: leaves.slice(),
    inclusion_proof(index: number): InclusionProof {
      if (index < 0 || index >= leafCount) {
        throw new Error(`Leaf index ${index} out of range [0, ${leafCount})`);
      }
      const siblings: AuditPathStep[] = [];
      let idx = index;
      for (let level = 0; level < levels.length - 1; level++) {
        const nodes = levels[level]!;
        const isRightChild = idx % 2 === 1;
        let siblingIdx: number;
        let side: 'L' | 'R';
        if (isRightChild) {
          siblingIdx = idx - 1;
          side = 'L';
        } else {
          // left child — sibling is to the right, or the node itself when odd
          siblingIdx = idx + 1 < nodes.length ? idx + 1 : idx;
          side = 'R';
        }
        siblings.push({ hash: nodes[siblingIdx]!, side });
        idx = Math.floor(idx / 2);
      }
      return {
        index,
        leaf_index: index,
        leafCount,
        tree_size: leafCount,
        siblings,
        audit_path: siblings,
        root,
        root_hash: root,
      };
    },
  };
}

/**
 * Verify an inclusion proof for a given leaf hash against an expected
 * root. Returns true iff walking the siblings reproduces the root.
 */
export function verifyInclusionProof(
  leaf: HexDigest,
  proof: InclusionProof,
  expectedRoot: HexDigest,
): boolean {
  if (proof.index < 0 || proof.index >= proof.leafCount) return false;
  let cur = leaf;
  for (const step of proof.siblings) {
    if (step.side === 'L') {
      cur = parentHash(step.hash, cur);
    } else {
      cur = parentHash(cur, step.hash);
    }
  }
  return cur === expectedRoot && cur === proof.root;
}

// ============================================================
// TREE HEAD + COSE_Sign1 ENVELOPE (RFC 9052)
// ============================================================

/**
 * Build an IETF draft-18 `tree_head` from a Merkle batch.
 * `timestamp` defaults to the current Unix epoch milliseconds.
 */
export function buildTreeHead(batch: MerkleBatch, now?: number): TreeHead {
  return {
    tree_size: batch.leafCount,
    root_hash: Buffer.from(batch.root, 'hex'),
    timestamp: now ?? Date.now(),
  };
}

/** COSE `alg` header parameter value for EdDSA (RFC 9053 §2.2). */
export const COSE_ALG_EDDSA = -8;

/** COSE standard header parameter labels (RFC 9052 §3.1). */
export const COSE_HEADER_ALG = 1;

/** Reusable cbor-x encoder configured for COSE-friendly output. */
const cborEncoder = new Encoder({
  // Keep canonical output stable. cbor-x defaults are fine for typed arrays;
  // we stick to byte strings + plain objects/arrays below.
  useRecords: false,
  tagUint8Array: false,
});

/**
 * CBOR-encode a value using the shared encoder. Returns a Buffer so the
 * result composes cleanly with Node's crypto APIs.
 */
function cborEncode(value: unknown): Buffer {
  const out = cborEncoder.encode(value);
  return Buffer.isBuffer(out) ? out : Buffer.from(out);
}

/**
 * Build the Sig_structure that COSE_Sign1 signs over (RFC 9052 §4.4):
 *   Sig_structure = [ "Signature1", body_protected, external_aad, payload ]
 *
 * Exported for testing; most callers want {@link wrapInCOSESign1}.
 */
export function buildSig_structure(
  protectedHeader: Buffer,
  payload: Buffer,
  externalAad: Buffer = Buffer.alloc(0),
): Buffer {
  return cborEncode(['Signature1', protectedHeader, externalAad, payload]);
}

/**
 * Build the CBOR-encoded payload for a tree-head + inclusion-proof bundle.
 * This is what COSE_Sign1 signs (inside the `Sig_structure` wrapper).
 *
 * Layout (plain CBOR map — RFC 9052 leaves payload format to the
 * application):
 *   {
 *     "tree_head":        { tree_size, root_hash (bstr), timestamp },
 *     "inclusion_proof":  { leaf_index, tree_size, audit_path },
 *   }
 *
 * `audit_path` is serialised as an array of `[hash_bstr, side_u8]` pairs,
 * where `side_u8` is 0 for 'L' (sibling on the left) and 1 for 'R'.
 */
export function encodeCoseMerklePayload(head: TreeHead, proof: InclusionProof): Buffer {
  const auditPath = proof.audit_path.map(step => [
    Buffer.from(step.hash, 'hex'),
    step.side === 'L' ? 0 : 1,
  ]);
  return cborEncode({
    tree_head: {
      tree_size: head.tree_size,
      root_hash: Buffer.from(head.root_hash),
      timestamp: head.timestamp,
    },
    inclusion_proof: {
      leaf_index: proof.leaf_index,
      tree_size: proof.tree_size,
      audit_path: auditPath,
    },
  });
}

/** COSE_Sign1 envelope bytes + the components used to build them. */
export interface CoseSign1Envelope {
  /** Canonical CBOR bytes of the COSE_Sign1 4-tuple. */
  bytes: Buffer;
  /** Protected header as a CBOR byte string (bstr). Empty protected map is `0x40` (bstr len 0). */
  protectedHeader: Buffer;
  /** Payload bytes that were signed (CBOR map per {@link encodeCoseMerklePayload}). */
  payload: Buffer;
  /** 64-byte Ed25519 signature. */
  signature: Buffer;
}

/**
 * Wrap a tree head + inclusion proof in a COSE_Sign1 envelope per RFC 9052 §4.2.
 *
 * Envelope layout (CBOR-encoded 4-tuple):
 *   [ protected: bstr, unprotected: {}, payload: bstr, signature: bstr ]
 *
 * The protected header encodes `{ alg: -8 }` (EdDSA, RFC 9053 §2.2). The
 * unprotected header is an empty map. The payload is the CBOR-encoded
 * `{ tree_head, inclusion_proof }` bundle produced by
 * {@link encodeCoseMerklePayload}. The signature is Ed25519 over the
 * `Sig_structure` from RFC 9052 §4.4.
 *
 * `privateKey` must be the same Ed25519 KeyObject used to sign the
 * underlying Remote ID logs (see {@link signRemoteId}); there is no
 * separate COSE key. This is additive — it does not replace the per-log
 * signatures.
 */
export function wrapInCOSESign1(
  head: TreeHead,
  proof: InclusionProof,
  privateKey: KeyObject,
): CoseSign1Envelope {
  // Protected header: encoded as a CBOR byte string wrapping a CBOR map.
  const protectedMap = cborEncode(new Map<number, number>([[COSE_HEADER_ALG, COSE_ALG_EDDSA]]));
  const protectedHeader = protectedMap;

  const payload = encodeCoseMerklePayload(head, proof);
  const sigStructure = buildSig_structure(protectedHeader, payload);

  // Ed25519 does not use a digest algorithm — pass null.
  const signature = nodeSign(null, sigStructure, privateKey);

  const envelope = cborEncode([
    protectedHeader,
    {}, // empty unprotected map
    payload,
    signature,
  ]);

  return {
    bytes: envelope,
    protectedHeader,
    payload,
    signature,
  };
}

/**
 * Verify a COSE_Sign1 envelope produced by {@link wrapInCOSESign1}.
 * Returns true iff the envelope decodes to a well-formed 4-tuple and
 * the signature verifies against `publicKey`.
 */
export function verifyCOSESign1(
  envelopeBytes: Buffer | Uint8Array,
  publicKey: KeyObject,
): boolean {
  let decoded: unknown;
  try {
    decoded = cborEncoder.decode(Buffer.from(envelopeBytes));
  } catch {
    return false;
  }
  if (!Array.isArray(decoded) || decoded.length !== 4) return false;
  const [protectedHeader, , payload, signature] = decoded as [
    Buffer | Uint8Array,
    unknown,
    Buffer | Uint8Array,
    Buffer | Uint8Array,
  ];
  if (!protectedHeader || !payload || !signature) return false;
  const sig = Buffer.from(signature);
  if (sig.length !== 64) return false;
  const sigStructure = buildSig_structure(Buffer.from(protectedHeader), Buffer.from(payload));
  try {
    return nodeVerify(null, sigStructure, publicKey, sig);
  } catch {
    return false;
  }
}