// Media for a WAI track (wai/extensions/moq-streaming-format.md §3): every object is one
// whole WAI1 envelope, whose payload is in the framing SPEC §5 pins for its capability:
//
//   wai.video.vp9    VP9 frames in the IVF frame container (FourCC VP90)
//   wai.video.av1    five u32 (n, h, w, fps_num, fps_den), then each temporal unit as a
//                    u32 length and its low-overhead bitstream (every OBU with its size)
//   wai.audio.opus   an Ogg Opus stream (RFC 7845)
//
// A live track carries one object per group (§4.1), so each object here is one encoded
// frame: an IVF file of one frame, an AV1 payload of one temporal unit, an Ogg Opus stream
// of one packet. The manifest's `target` (informational, SPEC §3) states the frame's
// dimensions, its timestamp and whether it is a key frame, which a decoder is told.

const le32 = (v, at, b) => new DataView(b.buffer, b.byteOffset).setUint32(at, v >>> 0, true);
const le16 = (v, at, b) => new DataView(b.buffer, b.byteOffset).setUint16(at, v, true);

export const CAPABILITY_OF = (codec) =>
  codec.startsWith("vp09") ? "wai.video.vp9" : codec.startsWith("av01") ? "wai.video.av1" : codec === "opus" ? "wai.audio.opus" : null;

// ---------------------------------------------------------------------------------- IVF

export function ivfOne(frame, width, height, timestampUs) {
  const b = new Uint8Array(32 + 12 + frame.length);
  b.set([0x44, 0x4b, 0x49, 0x46], 0); // DKIF
  le16(0, 4, b); // version
  le16(32, 6, b); // header size
  b.set([0x56, 0x50, 0x39, 0x30], 8); // VP90
  le16(width, 12, b);
  le16(height, 14, b);
  le32(1_000_000, 16, b); // time base: microseconds
  le32(1, 20, b);
  le32(1, 24, b); // one frame
  le32(frame.length, 32, b);
  const dv = new DataView(b.buffer);
  dv.setBigUint64(36, BigInt(Math.max(0, Math.round(timestampUs))), true);
  b.set(frame, 44);
  return b;
}

export function ivfFrames(b) {
  const dv = new DataView(b.buffer, b.byteOffset, b.byteLength);
  if (b.length < 32 || String.fromCharCode(...b.subarray(0, 4)) !== "DKIF") throw new Error("not an IVF file");
  const hl = dv.getUint16(6, true);
  const out = [];
  for (let at = hl; at + 12 <= b.length; ) {
    const n = dv.getUint32(at, true);
    const ts = Number(dv.getBigUint64(at + 4, true));
    if (at + 12 + n > b.length) throw new Error("an IVF frame past its file");
    out.push({ data: b.subarray(at + 12, at + 12 + n), timestamp: ts });
    at += 12 + n;
  }
  return out;
}

// ---------------------------------------------------------------------------- AV1 (SPEC §5)

export function av1One(tu, width, height, fpsNum, fpsDen) {
  const b = new Uint8Array(20 + 4 + tu.length);
  le32(1, 0, b);
  le32(height, 4, b);
  le32(width, 8, b);
  le32(fpsNum, 12, b);
  le32(fpsDen, 16, b);
  le32(tu.length, 20, b);
  b.set(tu, 24);
  return b;
}

export function av1Units(b) {
  const dv = new DataView(b.buffer, b.byteOffset, b.byteLength);
  if (b.length < 20) throw new Error("an AV1 payload shorter than its framing header");
  const n = dv.getUint32(0, true);
  const out = [];
  let at = 20;
  for (let i = 0; i < n; i++) {
    const len = dv.getUint32(at, true);
    if (at + 4 + len > b.length) throw new Error("an AV1 temporal unit past its payload");
    out.push(b.subarray(at + 4, at + 4 + len));
    at += 4 + len;
  }
  if (at !== b.length) throw new Error("an AV1 payload holding more temporal units than it declares");
  return out;
}

// ------------------------------------------------------------------------- Ogg Opus (RFC 7845)

const CRC = (() => {
  const t = new Uint32Array(256);
  for (let i = 0; i < 256; i++) {
    let r = i << 24;
    for (let j = 0; j < 8; j++) r = r & 0x80000000 ? (r << 1) ^ 0x04c11db7 : r << 1;
    t[i] = r >>> 0;
  }
  return t;
})();

function oggCrc(b) {
  let c = 0;
  for (let i = 0; i < b.length; i++) c = ((c << 8) ^ CRC[((c >>> 24) ^ b[i]) & 0xff]) >>> 0;
  return c;
}

function oggPage(packet, { bos = false, eos = false, granule = 0n, serial, seq }) {
  const segs = [];
  let n = packet.length;
  while (n >= 255) { segs.push(255); n -= 255; }
  segs.push(n);
  const b = new Uint8Array(27 + segs.length + packet.length);
  b.set([0x4f, 0x67, 0x67, 0x53], 0); // OggS
  b[4] = 0;
  b[5] = (bos ? 2 : 0) | (eos ? 4 : 0);
  const dv = new DataView(b.buffer);
  dv.setBigInt64(6, granule, true);
  dv.setUint32(14, serial, true);
  dv.setUint32(18, seq, true);
  b[26] = segs.length;
  b.set(segs, 27);
  b.set(packet, 27 + segs.length);
  dv.setUint32(22, oggCrc(b), true);
  return b;
}

// One Opus packet as a whole Ogg Opus stream: OpusHead, OpusTags, the packet. Pre-skip is
// 0: the object continues a live stream whose decoder is already running, and a decoder
// that starts here starts at the packet.
export function oggOpusOne(packet, { channels = 1, inputRate = 48000, samples, serial = 0x77616901 }) {
  const head = new Uint8Array(19);
  head.set([0x4f, 0x70, 0x75, 0x73, 0x48, 0x65, 0x61, 0x64], 0); // OpusHead
  head[8] = 1;
  head[9] = channels;
  le16(0, 10, head);
  le32(inputRate, 12, head);
  le16(0, 16, head);
  head[18] = 0;
  const vendor = new TextEncoder().encode("wai");
  const tags = new Uint8Array(8 + 4 + vendor.length + 4);
  tags.set([0x4f, 0x70, 0x75, 0x73, 0x54, 0x61, 0x67, 0x73], 0); // OpusTags
  le32(vendor.length, 8, tags);
  tags.set(vendor, 12);
  le32(0, 12 + vendor.length, tags);
  const pages = [
    oggPage(head, { bos: true, serial, seq: 0 }),
    oggPage(tags, { serial, seq: 1 }),
    oggPage(packet, { eos: true, granule: BigInt(samples), serial, seq: 2 }),
  ];
  const out = new Uint8Array(pages.reduce((n, p) => n + p.length, 0));
  let at = 0;
  for (const p of pages) { out.set(p, at); at += p.length; }
  return out;
}

// The audio packets of an Ogg Opus stream (its two header packets passed over), each page
// checked against its CRC.
export function oggOpusPackets(b) {
  const out = [];
  const dv = new DataView(b.buffer, b.byteOffset, b.byteLength);
  let pending = [];
  let index = 0;
  for (let at = 0; at < b.length; ) {
    if (String.fromCharCode(...b.subarray(at, at + 4)) !== "OggS") throw new Error("not an Ogg page");
    const nseg = b[at + 26];
    const segs = b.subarray(at + 27, at + 27 + nseg);
    const body = segs.reduce((n, x) => n + x, 0);
    const page = b.slice(at, at + 27 + nseg + body);
    const want = dv.getUint32(at + 22, true);
    new DataView(page.buffer).setUint32(22, 0, true);
    if (oggCrc(page) !== want) throw new Error("an Ogg page whose CRC does not match");
    let p = at + 27 + nseg;
    for (const s of segs) {
      pending.push(b.subarray(p, p + s));
      p += s;
      if (s < 255) {
        const n = pending.reduce((k, x) => k + x.length, 0);
        const pk = new Uint8Array(n);
        let o = 0;
        for (const x of pending) { pk.set(x, o); o += x.length; }
        if (index >= 2) out.push(pk);
        index++;
        pending = [];
      }
    }
    at += 27 + nseg + body;
  }
  return out;
}

// ------------------------------------------------------------------------------ envelopes

export function manifestFor(capability, target) {
  const media = capability.startsWith("wai.audio") ? "audio" : "video";
  const kind = capability.split(".").pop();
  // Member order is the writer's and is kept (SPEC §2).
  return JSON.stringify({ wai: "1.0", media, intent: "replicate", model_requirement: { capability, fallback: null }, conditioning: { kind }, target });
}