import { describe, expect, test } from "bun:test";
import {
  ParakeetVocabTokenizer,
  parakeetFileName,
  parakeetFileUrl,
  parseVocab,
  planAudioChunks,
  planChunks,
  transcribeWindows,
  wordsToChunks,
  type ParakeetWord,
} from "../src/lib/parakeet";

// Pure Parakeet engine logic (D-WS7a) — no model, no runtime.

const RATE = 16_000;

function word(text: string, start: number, end: number): ParakeetWord {
  return { text, start_time: start, end_time: end };
}

function expectContiguous(windows: Array<[number, number]>, length: number) {
  expect(windows[0]![0]).toBe(0);
  expect(windows[windows.length - 1]![1]).toBe(length);
  for (let i = 1; i < windows.length; i++) {
    expect(windows[i]![0]).toBe(windows[i - 1]![1]);
  }
}

describe("planChunks", () => {
  const flat = () => 1;

  test("audio up to the window length is a single window", () => {
    expect(planChunks(10 * RATE, flat)).toEqual([[0, 10 * RATE]]);
    expect(planChunks(30 * RATE, flat)).toEqual([[0, 30 * RATE]]);
    expect(planChunks(0, flat)).toEqual([]);
  });

  test("windows never exceed 30 s, are contiguous, and cover everything", () => {
    const length = 95 * RATE + 123;
    const windows = planChunks(length, flat);
    expectContiguous(windows, length);
    for (const [start, end] of windows) {
      expect(end - start).toBeLessThanOrEqual(30 * RATE);
      expect(end).toBeGreaterThan(start);
    }
  });

  test("a 90-minute file plans ≥ 180 bounded windows (no full-length copies needed)", () => {
    const length = 90 * 60 * RATE;
    const windows = planChunks(length, flat);
    expect(windows.length).toBeGreaterThanOrEqual(180);
    expectContiguous(windows, length);
    for (const [start, end] of windows) expect(end - start).toBeLessThanOrEqual(30 * RATE);
  });

  test("cuts in the quietest 100 ms of the last 6 s before the limit", () => {
    // Loud everywhere except a pause at 27.0–27.1 s.
    const quietStart = 27 * RATE;
    const energy = (start: number, end: number) =>
      start >= quietStart && end <= quietStart + 0.1 * RATE ? 0 : 1;
    const windows = planChunks(60 * RATE, energy);
    const firstCut = windows[0]![1];
    expect(firstCut).toBeGreaterThanOrEqual(quietStart);
    expect(firstCut).toBeLessThanOrEqual(quietStart + 0.1 * RATE);
  });

  test("planAudioChunks finds a real silent gap in PCM", () => {
    const audio = new Float32Array(50 * RATE);
    for (let i = 0; i < audio.length; i++) audio[i] = Math.sin(i / 7) * 0.5;
    audio.fill(0, 26 * RATE, 26.5 * RATE);
    const windows = planAudioChunks(audio);
    expect(windows[0]![1]).toBeGreaterThanOrEqual(26 * RATE);
    expect(windows[0]![1]).toBeLessThanOrEqual(26.5 * RATE);
    expectContiguous(windows, audio.length);
  });
});

describe("wordsToChunks", () => {
  test("splits after sentence-ending punctuation, keeping absolute timestamps", () => {
    const chunks = wordsToChunks([
      word("Hello", 0.5, 0.8),
      word("there.", 0.9, 1.2),
      word("How", 1.5, 1.7),
      word("are", 1.7, 1.8),
      word("you?", 1.8, 2.1),
    ]);
    expect(chunks).toEqual([
      { text: " Hello there.", timestamp: [0.5, 1.2] },
      { text: " How are you?", timestamp: [1.5, 2.1] },
    ]);
  });

  test("closing quotes after the full stop still end the sentence", () => {
    const chunks = wordsToChunks([word('"Yes."', 0, 0.4), word("Then", 0.5, 0.7)]);
    expect(chunks.map((chunk) => chunk.text)).toEqual([' "Yes."', " Then"]);
  });

  test("a pause of at least 1 s starts a new chunk", () => {
    const chunks = wordsToChunks([word("one", 0, 0.3), word("two", 1.4, 1.6)]);
    expect(chunks).toHaveLength(2);
  });

  test("a run-on sentence is split once it reaches 12 s", () => {
    const words = Array.from({ length: 40 }, (_, i) => word(`w${i}`, i * 0.5, i * 0.5 + 0.4));
    const chunks = wordsToChunks(words);
    expect(chunks.length).toBeGreaterThan(1);
    for (const chunk of chunks) {
      expect((chunk.timestamp[1] ?? 0) - chunk.timestamp[0]).toBeLessThanOrEqual(12.5);
    }
    expect(chunks.map((c) => c.text.trim()).join(" ").split(" ")).toHaveLength(40);
  });

  test("no words → no chunks", () => {
    expect(wordsToChunks([])).toEqual([]);
  });
});

describe("vocab tokenizer", () => {
  const vocab = "<unk> 0\n▁the 1\n▁cat 2\ns 3\n. 4\n<blk> 5\n";

  test("parseVocab indexes tokens by id", () => {
    expect(parseVocab(vocab)).toEqual(["<unk>", "▁the", "▁cat", "s", ".", "<blk>"]);
  });

  test("decode joins pieces, skips blank, and attaches punctuation", () => {
    const tokenizer = new ParakeetVocabTokenizer(parseVocab(vocab));
    expect(tokenizer.blankId).toBe(5);
    expect(tokenizer.blankToken).toBe("<blk>");
    expect(tokenizer.decode([1, 5, 2, 3, 4])).toBe("the cats.");
  });
});

describe("Parakeet files", () => {
  test("int8 and fp16 file names; other dtypes are not offered", () => {
    expect(parakeetFileName("encoder", "q8")).toBe("encoder-model.int8.onnx");
    expect(parakeetFileName("decoder", "q8")).toBe("decoder_joint-model.int8.onnx");
    expect(parakeetFileName("encoder", "fp16")).toBe("encoder-model.fp16.onnx");
    expect(() => parakeetFileName("encoder", "q4")).toThrow();
  });

  test("URLs are pinned to the measured revision under the model id", () => {
    expect(parakeetFileUrl("vocab.txt")).toBe(
      "https://huggingface.co/ysdede/parakeet-tdt-0.6b-v3-onnx/resolve/" +
        "e5356bf6e24bdb7a269536b0e9e7869bc034b893/vocab.txt",
    );
  });
});

describe("transcribeWindows", () => {
  test("recognizes each window with its offset and reports progress + new chunks", async () => {
    const audio = new Float32Array(70 * RATE).fill(0.1);
    const offsets: number[] = [];
    const sizes: number[] = [];
    const progress: number[] = [];
    const streamed: string[] = [];
    const result = await transcribeWindows(
      audio,
      async (window, offset) => {
        offsets.push(offset);
        sizes.push(window.length);
        return [word(`at${Math.round(offset)}.`, offset + 0.1, offset + 0.5)];
      },
      {
        onWindow: (chunks, processed, total) => {
          progress.push(processed);
          expect(total).toBe(70);
          streamed.push(...chunks.map((chunk) => chunk.text));
        },
      },
    );
    expect(offsets[0]).toBe(0);
    expect(offsets.length).toBe(3);
    for (const size of sizes) expect(size).toBeLessThanOrEqual(30 * RATE);
    expect(progress[progress.length - 1]).toBe(70);
    expect(progress).toEqual([...progress].sort((a, b) => a - b));
    expect(result.chunks.map((c) => c.text)).toEqual(streamed);
    expect(result.text).toBe(streamed.join("").trim());
    expect(result.chunks[1]!.timestamp[0]).toBeCloseTo(offsets[1]! + 0.1, 5);
  });

  test("silence-only audio: empty transcript, progress still reaches the end", async () => {
    const progress: number[] = [];
    const result = await transcribeWindows(
      new Float32Array(40 * RATE),
      async () => [],
      { onWindow: (_chunks, processed) => progress.push(processed) },
    );
    expect(result).toEqual({ text: "", chunks: [] });
    expect(progress[progress.length - 1]).toBe(40);
  });
});