import { describe, expect, test } from "bun:test";
import {
  chunkAudio,
  chunkTimeToOriginal,
  collectChunks,
  getSpeechTimestamps,
  SAMPLE_RATE,
  VAD_WINDOW_SAMPLES,
} from "../src/worker/longform/vad";

const W = VAD_WINDOW_SAMPLES;
const PAD = (400 / 1000) * SAMPLE_RATE; // speech_pad_ms 400 → 6400 samples

/** Probabilities for `seconds` of audio with speech in the given [start, end) second ranges. */
function probsFor(seconds: number, speech: [number, number][]): Float32Array {
  const windows = Math.ceil((seconds * SAMPLE_RATE) / W);
  const probs = new Float32Array(windows).fill(0.05);
  for (const [start, end] of speech) {
    const from = Math.floor((start * SAMPLE_RATE) / W);
    const to = Math.floor((end * SAMPLE_RATE) / W);
    for (let i = from; i < to; i++) probs[i] = 0.9;
  }
  return probs;
}

describe("getSpeechTimestamps (port of faster-whisper vad.py)", () => {
  test("no window above the threshold → no speech", () => {
    const probs = new Float32Array(500).fill(0.3);
    expect(getSpeechTimestamps(probs, 500 * W)).toEqual([]);
  });

  test("one burst is padded by speech_pad_ms on both sides", () => {
    const probs = new Float32Array(200).fill(0.05);
    for (let i = 50; i < 80; i++) probs[i] = 0.9;
    expect(getSpeechTimestamps(probs, 200 * W)).toEqual([
      { start: 50 * W - PAD, end: 80 * W + PAD },
    ]);
  });

  test("padding is clamped to the audio", () => {
    const probs = new Float32Array(40).fill(0.9);
    expect(getSpeechTimestamps(probs, 40 * W)).toEqual([{ start: 0, end: 40 * W }]);
  });

  test("a pause shorter than min_silence_duration_ms does not split", () => {
    const spans = getSpeechTimestamps(probsFor(20, [[2, 6], [7, 11]]), 20 * SAMPLE_RATE);
    expect(spans).toHaveLength(1);
  });

  test("a pause longer than min_silence_duration_ms splits; close spans share the gap", () => {
    const spans = getSpeechTimestamps(probsFor(30, [[2, 6], [9, 12]]), 30 * SAMPLE_RATE);
    expect(spans).toHaveLength(2);
    // 3 s gap ≥ 2·pad → each side padded by the full 400 ms.
    expect(spans[0]!.end - Math.floor((6 * SAMPLE_RATE) / W) * W).toBe(PAD);
    const narrow = getSpeechTimestamps(probsFor(30, [[2, 6], [8.5, 12]]), 30 * SAMPLE_RATE, {
      minSilenceDurationMs: 2000,
      speechPadMs: 1500,
    });
    // 2.5 s gap < 2·1.5 s pad → the gap is split in half instead.
    expect(narrow).toHaveLength(2);
    expect(narrow[0]!.end).toBe(narrow[1]!.start);
  });

  test("continuous speech is cut at max_speech_duration_s", () => {
    const spans = getSpeechTimestamps(probsFor(70, [[0, 70]]), 70 * SAMPLE_RATE, {
      maxSpeechDurationS: 30,
    });
    expect(spans.length).toBeGreaterThanOrEqual(3);
    for (const span of spans) {
      expect(span.end - span.start).toBeLessThanOrEqual(30 * SAMPLE_RATE);
    }
  });
});

describe("collectChunks / chunk time mapping", () => {
  const s = (sec: number) => sec * SAMPLE_RATE;
  const spans = [
    { start: s(0), end: s(12) },
    { start: s(20), end: s(32) },
    { start: s(40), end: s(52) },
  ];

  test("packs spans into chunks no longer than the maximum", () => {
    const chunks = collectChunks(spans, s(30));
    expect(chunks.map((chunk) => chunk.samples)).toEqual([s(24), s(12)]);
    expect(chunks[0]!.spans).toEqual([spans[0]!, spans[1]!]);
  });

  test("no spans → no chunks", () => {
    expect(collectChunks([], s(30))).toEqual([]);
  });

  test("chunk-relative seconds map back to the original timeline", () => {
    const [chunk] = collectChunks(spans, s(30));
    expect(chunkTimeToOriginal(chunk!, 0)).toBe(0);
    expect(chunkTimeToOriginal(chunk!, 5)).toBe(5);
    expect(chunkTimeToOriginal(chunk!, 12)).toBe(12); // splice → earlier span's end
    expect(chunkTimeToOriginal(chunk!, 15)).toBe(23);
    expect(chunkTimeToOriginal(chunk!, 99)).toBe(32); // clamps to the last span
  });

  test("chunkAudio copies exactly the chunk's spans", () => {
    const audio = new Float32Array(s(60)).map((_, i) => i);
    const [chunk] = collectChunks(spans, s(30));
    const out = chunkAudio(audio, chunk!);
    expect(out.length).toBe(chunk!.samples);
    expect(out[0]).toBe(0);
    expect(out[s(12)]).toBe(s(20));
  });
});