import { describe, expect, it } from 'vitest';
import { AVC_CODEC, colourOf, type MatroskaHead, MatroskaWriter } from '../../src/core/recording/matroska';

/** One element as read back: its id in hex, its bytes, and what it holds where it is a master. */
interface Read {
  id: string;
  data: Uint8Array;
  children: Read[];
}

const MASTERS = new Set([
  '1a45dfa3',
  '18538067',
  '1549a966',
  '1654ae6b',
  'ae',
  'e0',
  '55b0',
  'e1',
  '1f43b675',
]);

function vint(bytes: Uint8Array, at: number, keepMarker: boolean): { value: number; length: number } {
  const first = bytes[at]!;
  let length = 1;
  while (length <= 8 && !(first & (0x80 >> (length - 1)))) length++;
  let value = keepMarker ? first : first & (0xff >> length);
  for (let i = 1; i < length; i++) value = value * 256 + bytes[at + i]!;
  return { value, length };
}

function readAll(bytes: Uint8Array, from = 0, to = bytes.length): Read[] {
  const out: Read[] = [];
  let at = from;
  while (at < to) {
    const id = vint(bytes, at, true);
    const size = vint(bytes, at + id.length, false);
    const start = at + id.length + size.length;
    const data = bytes.subarray(start, start + size.value);
    const hex = id.value.toString(16);
    out.push({ id: hex, data, children: MASTERS.has(hex) ? readAll(bytes, start, start + size.value) : [] });
    at = start + size.value;
  }
  return out;
}

const find = (list: Read[], id: string) => list.find((e) => e.id === id);
const uintOf = (e: Read | undefined) => (e ? e.data.reduce((v, b) => v * 256 + b, 0) : NaN);
const textOf = (e: Read | undefined) => new TextDecoder().decode(e?.data);

/** Track, time relative to its cluster, key flag and payload of each block, with the cluster's time. */
function blocksOf(
  file: Uint8Array,
): { cluster: number; track: number; ms: number; key: boolean; data: number[] }[] {
  const segment = find(readAll(file), '18538067')!;
  return segment.children
    .filter((e) => e.id === '1f43b675')
    .flatMap((cluster) => {
      const start = uintOf(find(cluster.children, 'e7'));
      return cluster.children
        .filter((e) => e.id === 'a3')
        .map((block) => ({
          cluster: start,
          track: block.data[0]! & 0x7f,
          ms: start + new DataView(block.data.buffer, block.data.byteOffset + 1, 2).getInt16(0),
          key: (block.data[3]! & 0x80) !== 0,
          data: [...block.data.subarray(4)],
        }));
    });
}

const HEAD: MatroskaHead = {
  video: {
    codec: AVC_CODEC,
    width: 1280,
    height: 720,
    private: Uint8Array.from([1, 100, 0, 31]),
    fps: 30,
    colour: colourOf({ matrix: 'bt709', transfer: 'iec61966-2-1', primaries: 'bt709', fullRange: false }),
  },
  audio: { channels: 2, private: new TextEncoder().encode('OpusHead'), preSkip: 312 },
  durationMs: 1234.5,
  app: 'feuerwerk',
};

/** The whole file: the head, then the clusters handed out before (`earlier`) and at the end. */
function fileOf(writer: MatroskaWriter, head: MatroskaHead = HEAD, earlier: Uint8Array[] = []): Uint8Array {
  const taken = writer.take();
  const { head: bytes, rest } = writer.finish(head);
  const parts = [bytes, ...earlier, ...taken, ...rest];
  const out = new Uint8Array(parts.reduce((n, p) => n + p.length, 0));
  let at = 0;
  for (const part of parts) {
    out.set(part, at);
    at += part.length;
  }
  return out;
}

describe('MatroskaWriter (#405)', () => {
  it('writes a head that says what the tracks are and how long the file is', () => {
    const writer = new MatroskaWriter(2);
    writer.add(1, 0, true, Uint8Array.from([9]));
    const file = fileOf(writer);
    const [ebml, segment] = readAll(file);
    expect(textOf(find(ebml!.children, '4282'))).toBe('matroska');
    // The segment's size covers exactly the rest of the file.
    expect(segment!.data.byteOffset + segment!.data.length).toBe(file.length);
    const info = find(segment!.children, '1549a966')!;
    expect(uintOf(find(info.children, '2ad7b1'))).toBe(1_000_000);
    const duration = find(info.children, '4489')!;
    expect(new DataView(duration.data.buffer, duration.data.byteOffset, 8).getFloat64(0)).toBe(1234.5);
    const [video, audio] = find(segment!.children, '1654ae6b')!.children;
    expect(textOf(find(video!.children, '86'))).toBe('V_MPEG4/ISO/AVC');
    expect([...find(video!.children, '63a2')!.data]).toEqual([1, 100, 0, 31]);
    const picture = find(video!.children, 'e0')!;
    expect([uintOf(find(picture.children, 'b0')), uintOf(find(picture.children, 'ba'))]).toEqual([1280, 720]);
    // A frame lasts a 30th of a second, in nanoseconds, so nothing has to guess the rate.
    expect(uintOf(find(video!.children, '23e383'))).toBe(33_333_333);
    // BT.709 pixels in the broadcast range, sRGB's transfer.
    const colour = find(picture.children, '55b0')!.children;
    expect(['55b1', '55b9', '55ba', '55bb'].map((id) => uintOf(find(colour, id)))).toEqual([1, 1, 13, 1]);
    expect(textOf(find(audio!.children, '86'))).toBe('A_OPUS');
    expect(textOf(find(audio!.children, '63a2'))).toBe('OpusHead');
    // 312 samples at 48 kHz, in nanoseconds.
    expect(uintOf(find(audio!.children, '56aa'))).toBe(6_500_000);
    expect(uintOf(find(find(audio!.children, 'e1')!.children, '9f'))).toBe(2);
  });

  it('takes the colour from what the encoder says, BT.709 where it says nothing', () => {
    expect(colourOf(null)).toEqual({ matrix: 1, transfer: 1, primaries: 1, fullRange: false });
    expect(colourOf({ matrix: 'smpte170m', fullRange: true })).toEqual({
      matrix: 6,
      transfer: 1,
      primaries: 1,
      fullRange: true,
    });
  });

  it('writes the blocks of both tracks in the order of their time, whatever order they come in', () => {
    const writer = new MatroskaWriter(2);
    writer.add(1, 0, true, Uint8Array.from([1]));
    writer.add(1, 16_667, false, Uint8Array.from([2]));
    writer.add(1, 33_333, false, Uint8Array.from([3]));
    writer.add(2, 0, true, Uint8Array.from([10]));
    writer.add(2, 20_000, true, Uint8Array.from([11]));
    writer.add(2, 40_000, true, Uint8Array.from([12]));
    const blocks = blocksOf(fileOf(writer));
    expect(blocks.map((b) => [b.track, b.ms])).toEqual([
      [1, 0],
      [2, 0],
      [1, 17],
      [2, 20],
      [1, 33],
      [2, 40],
    ]);
    expect(blocks.map((b) => b.data[0])).toEqual([1, 10, 2, 11, 3, 12]);
  });

  it('hands out a cluster once the next key frame starts another', () => {
    const writer = new MatroskaWriter(1);
    writer.add(1, 0, true, Uint8Array.from([1]));
    writer.add(1, 1_000_000, false, Uint8Array.from([2]));
    expect(writer.take()).toEqual([]);
    writer.add(1, 2_000_000, true, Uint8Array.from([3]));
    const first = writer.take();
    expect(first).toHaveLength(1);
    const blocks = blocksOf(fileOf(writer, HEAD, first));
    expect(blocks.map((b) => [b.cluster, b.ms, b.key])).toEqual([
      [0, 0, true],
      [0, 1000, false],
      [2000, 2000, true],
    ]);
  });

  it('starts a new cluster before a block would be too far from its start to be told', () => {
    const writer = new MatroskaWriter(1);
    writer.add(1, 0, true, Uint8Array.from([1]));
    writer.add(1, 40_000_000, false, Uint8Array.from([2]));
    const blocks = blocksOf(fileOf(writer));
    expect(blocks.map((b) => [b.cluster, b.ms])).toEqual([
      [0, 0],
      [40_000, 40_000],
    ]);
  });

  it('does not hold the video back for ever when no sound comes', () => {
    const writer = new MatroskaWriter(2);
    for (let i = 0; i < 90; i++) writer.add(1, i * 33_333, i % 30 === 0, Uint8Array.from([i]));
    // Three seconds of video, no sound: what is more than a second behind has gone into clusters.
    expect(writer.take().length).toBeGreaterThan(0);
  });

  it('counts its bytes, the open cluster included, and leaves out a missing sound track', () => {
    const writer = new MatroskaWriter(1);
    writer.add(1, 0, true, new Uint8Array(100));
    expect(writer.bytes).toBeGreaterThan(100);
    const file = fileOf(writer, { ...HEAD, audio: null });
    const segment = find(readAll(file), '18538067')!;
    expect(find(segment.children, '1654ae6b')!.children).toHaveLength(1);
  });
});