import { describe, expect, test } from "bun:test";
import {
  transcribeLongForm,
  type DecodeRequest,
  type DecodeResult,
  type LongFormDeps,
} from "../src/worker/longform";
import { SAMPLE_RATE, VAD_WINDOW_SAMPLES } from "../src/worker/longform/vad";

const MAX_CHUNK_SAMPLES = 30 * SAMPLE_RATE;

/** Speech probabilities for audio of `seconds` with speech in [start, end) ranges. */
function maskProbs(seconds: number, speech: [number, number][]): Float32Array {
  const windows = Math.ceil((seconds * SAMPLE_RATE) / VAD_WINDOW_SAMPLES);
  const probs = new Float32Array(windows).fill(0.02);
  for (const [start, end] of speech) {
    const from = Math.floor((start * SAMPLE_RATE) / VAD_WINDOW_SAMPLES);
    const to = Math.min(windows, Math.floor((end * SAMPLE_RATE) / VAD_WINDOW_SAMPLES));
    for (let i = from; i < to; i++) probs[i] = 0.95;
  }
  return probs;
}

interface Recorder {
  deps: LongFormDeps;
  decodes: { audio: Float32Array; request: DecodeRequest }[];
  vadInputs: Float32Array[];
  detectCalls: Float32Array[];
}

function fakeDeps(
  probs: Float32Array,
  decode: (audio: Float32Array, request: DecodeRequest, call: number) => DecodeResult = (
    audio,
  ) => ({
    text: ` words ${audio.length}`,
    chunks: [{ text: ` words ${audio.length}`, timestamp: [0, audio.length / SAMPLE_RATE] }],
    generatedTokens: 5,
    avgLogProb: -0.2,
  }),
  detect: string | null = null,
): Recorder {
  const recorder: Recorder = { decodes: [], vadInputs: [], detectCalls: [], deps: null as never };
  recorder.deps = {
    speechProbabilities: async (audio) => {
      recorder.vadInputs.push(audio);
      return probs;
    },
    decode: async (audio, request) => {
      recorder.decodes.push({ audio, request });
      request.onPartial?.("x");
      request.onTimestamp?.(1);
      return decode(audio, request, recorder.decodes.length);
    },
    ...(detect !== null && {
      detectLanguage: async (audio: Float32Array) => {
        recorder.detectCalls.push(audio);
        return detect;
      },
    }),
  };
  return recorder;
}

describe("transcribeLongForm", () => {
  test("no speech → empty result, decode never called", async () => {
    const audio = new Float32Array(60 * SAMPLE_RATE);
    const { deps, decodes } = fakeDeps(maskProbs(60, []));
    const progress: number[] = [];
    const result = await transcribeLongForm(audio, deps, {
      onProgress: (processed) => progress.push(processed),
    });
    expect(decodes).toHaveLength(0);
    expect(result.text).toBe("");
    expect(result.chunks).toEqual([]);
    expect(result.noSpeech).toBe(true);
    expect(progress.at(-1)).toBe(60);
  });

  test("speech regions are decoded per chunk and timestamps mapped to the file", async () => {
    // Two 5 s regions 15 s apart → one ≤ 30 s chunk holding both spans.
    const audio = new Float32Array(40 * SAMPLE_RATE);
    const { deps, decodes } = fakeDeps(maskProbs(40, [[5, 10], [25, 30]]), (chunk) => ({
      text: " hello there",
      chunks: [
        { text: " hello", timestamp: [0.5, 2] },
        { text: " there", timestamp: [6.5, 8] },
      ],
      generatedTokens: 4,
      avgLogProb: -0.1,
    }));
    const result = await transcribeLongForm(audio, deps, {});
    expect(decodes).toHaveLength(1);
    // Window-aligned spans padded by 400 ms: [73 472, 166 144) and
    // [393 472, 486 144) → 2 × 92 672 samples, silence between them dropped.
    expect(decodes[0]!.audio.length).toBe(2 * 92_672);
    expect(result.text).toBe("hello there");
    expect(result.noSpeech).toBe(false);
    // The chunk starts at 4.592 s; 6.5 s into the chunk is past the first
    // span (5.792 s long) → 0.708 s into the second, which starts at 24.592 s.
    expect(result.chunks[0]!.timestamp[0]).toBeCloseTo(4.592 + 0.5, 3);
    expect(result.chunks[1]!.timestamp[0]).toBeCloseTo(24.592 + 0.708, 3);
    expect(result.chunks[1]!.timestamp[0]).toBeGreaterThan(24);
    expect(result.chunks[1]!.timestamp[1]!).toBeLessThan(30.5);
  });

  test("auto language: detected once on the first speech chunk, then pinned", async () => {
    const audio = new Float32Array(100 * SAMPLE_RATE);
    const recorder = fakeDeps(maskProbs(100, [[0, 25], [40, 65], [75, 99]]), undefined, "da");
    const result = await transcribeLongForm(audio, recorder.deps, { language: "auto" });
    expect(recorder.detectCalls).toHaveLength(1);
    expect(recorder.detectCalls[0]).toBe(recorder.decodes[0]!.audio);
    expect(recorder.decodes.length).toBeGreaterThanOrEqual(3);
    for (const { request } of recorder.decodes) expect(request.language).toBe("da");
    expect(result.language).toBe("da");
  });

  test("an explicit language skips detection; English-only models send none", async () => {
    const audio = new Float32Array(20 * SAMPLE_RATE);
    const pinned = fakeDeps(maskProbs(20, [[2, 10]]), undefined, "de");
    await transcribeLongForm(audio, pinned.deps, { language: "en" });
    expect(pinned.detectCalls).toHaveLength(0);
    expect(pinned.decodes[0]!.request.language).toBe("en");

    const englishOnly = fakeDeps(maskProbs(20, [[2, 10]]));
    await transcribeLongForm(audio, englishOnly.deps, {});
    expect(englishOnly.decodes[0]!.request.language).toBeUndefined();
  });

  test("task, prompt and the repetition guard reach every decode", async () => {
    const audio = new Float32Array(100 * SAMPLE_RATE);
    const recorder = fakeDeps(maskProbs(100, [[0, 25], [40, 65], [75, 99]]));
    await transcribeLongForm(audio, recorder.deps, {
      task: "translate",
      prompt: "Voices of the Void",
      guard: { noRepeatNgramSize: 5, repetitionPenalty: 1.1 },
    });
    expect(recorder.decodes.length).toBeGreaterThanOrEqual(3);
    for (const { request } of recorder.decodes) {
      expect(request.task).toBe("translate");
      expect(request.prompt).toBe("Voices of the Void");
      expect(request.noRepeatNgramSize).toBe(5);
      expect(request.repetitionPenalty).toBe(1.1);
      expect(request.temperature).toBe(0);
    }
  });

  test("fallback: a looping decode is retried at the next temperature", async () => {
    const audio = new Float32Array(20 * SAMPLE_RATE);
    const loop = " the main clar,".repeat(40);
    const recorder = fakeDeps(maskProbs(20, [[2, 10]]), (_audio, request) => ({
      text: request.temperature === 0 ? loop : " a fine sentence about festivals",
      chunks: [],
      generatedTokens: request.temperature === 0 ? 400 : 8,
      avgLogProb: -0.2,
    }));
    const result = await transcribeLongForm(audio, recorder.deps, { fallback: {} });
    expect(recorder.decodes.map((d) => d.request.temperature)).toEqual([0, 0.2]);
    expect(result.text).toBe("a fine sentence about festivals");
    expect(result.generatedTokens).toBe(408);
  });

  test("fallback: when every temperature fails, the best attempt is kept", async () => {
    const audio = new Float32Array(20 * SAMPLE_RATE);
    const recorder = fakeDeps(maskProbs(20, [[2, 10]]), (_audio, request) => ({
      text: request.temperature === 0.6 ? " ok ok ok ok ok ok ok ok ok ok ok ok ok" : " la".repeat(80),
      chunks: [],
      generatedTokens: 10,
      avgLogProb: null,
    }));
    const result = await transcribeLongForm(audio, recorder.deps, { fallback: {} });
    expect(recorder.decodes).toHaveLength(6);
    expect(result.text).toBe("ok ok ok ok ok ok ok ok ok ok ok ok ok");
  });

  test("fallback off → exactly one decode per chunk", async () => {
    const audio = new Float32Array(20 * SAMPLE_RATE);
    const recorder = fakeDeps(maskProbs(20, [[2, 10]]), () => ({
      text: " la".repeat(80),
      chunks: [],
      generatedTokens: 10,
      avgLogProb: -3,
    }));
    await transcribeLongForm(audio, recorder.deps, { fallback: false });
    expect(recorder.decodes).toHaveLength(1);
  });

  test("progress is monotone and ends at the total; partials are forwarded", async () => {
    const audio = new Float32Array(100 * SAMPLE_RATE);
    const recorder = fakeDeps(maskProbs(100, [[0, 25], [40, 65], [75, 90]]));
    const progress: number[] = [];
    const partials: string[] = [];
    await transcribeLongForm(audio, recorder.deps, {
      onProgress: (processed, total) => {
        expect(total).toBe(100);
        progress.push(processed);
      },
      onPartial: (text) => partials.push(text),
    });
    expect(progress[0]).toBe(0);
    expect(progress.at(-1)).toBe(100);
    for (let i = 1; i < progress.length; i++) {
      expect(progress[i]!).toBeGreaterThan(progress[i - 1]!);
    }
    expect(partials.length).toBe(recorder.decodes.length);
  });

  test("90-minute input: VAD sees the caller's buffer, decode sees ≤ 30 s chunks", async () => {
    const seconds = 90 * 60;
    const audio = new Float32Array(seconds * SAMPLE_RATE);
    const speech: [number, number][] = [];
    for (let t = 0; t < seconds; t += 25) speech.push([t, t + 20]);
    const recorder = fakeDeps(maskProbs(seconds, speech));
    const started = performance.now();
    const result = await transcribeLongForm(audio, recorder.deps, {});
    expect(performance.now() - started).toBeLessThan(5_000);
    expect(recorder.vadInputs).toHaveLength(1);
    expect(recorder.vadInputs[0]).toBe(audio); // no copy before VAD
    let decoded = 0;
    for (const { audio: chunk } of recorder.decodes) {
      expect(chunk.length).toBeLessThanOrEqual(MAX_CHUNK_SAMPLES);
      expect(chunk.buffer).not.toBe(audio.buffer);
      decoded += chunk.length;
    }
    // Every speech second is decoded once (padding included), silence never.
    expect(decoded).toBeLessThan(audio.length);
    expect(result.speechSeconds).toBeCloseTo(decoded / SAMPLE_RATE, 3);
    expect(result.decodeCalls).toBe(recorder.decodes.length);
  });
});