import { test, expect, beforeAll } from "bun:test";
import { join } from "node:path";
import { buildSong, loadSongDir } from "../../src/bench/dataset.ts";
import { createPipeline, copyFingering, sumAudio } from "../../src/pipeline/run.ts";
import { createTranscriber } from "../../src/pipeline/transcribe.ts";
import { repeatVote } from "../../src/pipeline/repeat.ts";
import { snapTicks, quantize } from "../../src/pipeline/quantize.ts";
import type { TimedNote } from "../../src/core/types.ts";
import { truthFromGp } from "../../src/score/gp.ts";
import { scoreSong } from "../../src/eval/score.ts";

const SLOW = 600_000;
const root = join(import.meta.dir, "../..");
let run: Awaited<ReturnType<typeof createPipeline>>;

beforeAll(async () => {
  const bytes = async (p: string) => new Uint8Array(await Bun.file(join(root, p)).arrayBuffer());
  run = await createPipeline({ demucs: await bytes("models/demucs/htdemucs_6s.onnx"), basicPitch: await bytes("models/basic-pitch/nmp.onnx") });
  await buildSong("t1a", 1);
  await buildSong("t1c", 1);
}, SLOW);

test("t1a/1 oracle stems+tempo: guitar track, GP re-imports", async () => {
  const d = await loadSongDir(join(root, "data/t1a/1"));
  const flags = { oracleStems: true, oracleTempo: true, oracleNotes: false };
  const { output, gp } = await run({ mix: d.mix, oracle: { stems: d.stems, truth: d.truth } }, flags);
  const s = scoreSong(d.truth, output, flags);
  expect(s.guitarMerged.f1!).toBeGreaterThan(0.7);
  const g = output.tracks.filter((t) => t.role === "guitar");
  expect(g.length).toBe(1);
  expect(g[0]!.notes.length).toBeGreaterThan(0);
  expect(truthFromGp(gp).tracks.some((t) => t.role === "guitar")).toBe(true);
}, SLOW);

test("t1a/1 oracle notes: merged guitar F1 === 1, stringAgreement is a number", async () => {
  const d = await loadSongDir(join(root, "data/t1a/1"));
  const flags = { oracleStems: true, oracleTempo: true, oracleNotes: true };
  const { output } = await run({ mix: d.mix, oracle: { stems: d.stems, truth: d.truth } }, flags);
  const s = scoreSong(d.truth, output, flags);
  expect(s.guitarMerged.f1).toBe(1);
  expect(typeof s.stringAgreement).toBe("number");
}, SLOW);

test("t1c/1 full pipeline yields all roles", async () => {
  const d = await loadSongDir(join(root, "data/t1c/1"));
  const flags = { oracleStems: false, oracleTempo: false, oracleNotes: false };
  const { output, gp, warnings } = await run({ mix: d.mix }, flags);
  const s = scoreSong(d.truth, output, flags);
  const roles = new Set(output.tracks.map((t) => t.role));
  for (const r of ["guitar", "bass", "keys", "vocals", "drums"]) expect(roles.has(r as never)).toBe(true);
  expect(gp.length).toBeGreaterThan(0);
  expect(warnings.some((w) => w.includes("other"))).toBe(false);
  expect(Math.abs(output.grid.bpm - d.truth.bpm)).toBeLessThan(2);
  expect(s.guitarMerged.f1!).toBeGreaterThan(0.6); // guitar from guitar+piano+other stems (was 0.57 from the guitar stem alone)
}, SLOW);

test("t1d/1 all oracle flags: two guitar tracks, partSplit and merged F1 === 1", async () => {
  await buildSong("t1d", 1);
  const d = await loadSongDir(join(root, "data/t1d/1"));
  const flags = { oracleStems: true, oracleTempo: true, oracleNotes: true };
  const { output } = await run({ mix: d.mix, oracle: { stems: d.stems, truth: d.truth } }, flags);
  const s = scoreSong(d.truth, output, flags);
  const names = output.tracks.filter((t) => t.role === "guitar").map((t) => t.name).sort();
  expect(names).toEqual(d.truth.tracks.filter((t) => t.role === "guitar").map((t) => t.name).sort());
  expect(names.length).toBe(2);
  expect(s.partSplit).not.toBeNull();
  expect(s.partSplit!.f1).toBe(1);
  expect(s.guitarMerged.f1).toBe(1);
}, SLOW);

test("oracle flag without input.oracle throws", async () => {
  const d = await loadSongDir(join(root, "data/t1a/1"));
  await expect(run({ mix: d.mix }, { oracleStems: false, oracleTempo: true, oracleNotes: false })).rejects.toThrow(/oracle/);
});

test("silent mix with oracle stems: estimateGrid fallback warns and does not crash", async () => {
  const d = await loadSongDir(join(root, "data/t1a/1"));
  const silent = { sampleRate: 44100, channels: [new Float32Array(44100 * 3)] };
  const flags = { oracleStems: true, oracleTempo: false, oracleNotes: true };
  const { output, warnings } = await run({ mix: silent, oracle: { stems: d.stems, truth: d.truth } }, flags);
  expect(output.grid).toEqual({ bpm: 120, firstDownbeat: 0, beatsPerBar: 4 });
  expect(warnings.some((w) => w.includes("estimateGrid failed"))).toBe(true);
});

test("copyFingering matches by snapped tick + pitch (merged cluster, adjacent repeats)", () => {
  const grid = { bpm: 120, firstDownbeat: 0.1, beatsPerBar: 4 };
  const n = (onset: number, pitch: number): TimedNote => ({ onset, offset: onset + 0.1, pitch });
  // a,b same pitch 20 ms apart (merge into one tick 0); c same pitch on the adjacent 16th (tick 240); d another pitch at 0
  const timed = [n(0.1, 60), n(0.12, 60), n(0.225, 60), n(0.1, 64)];
  expect(snapTicks(timed, grid)).toEqual([0, 0, 240, 0]);
  const assigned = quantize(timed, grid).map((q) => ({ ...q, string: q.pitch === 60 ? (q.tick === 0 ? 2 : 3) : 1, fret: q.tick }));
  copyFingering(timed, assigned, grid);
  expect(timed.map((x) => [x.string, x.fret])).toEqual([[2, 0], [2, 0], [3, 240], [1, 0]]);
});

test("sumAudio adds buffers sample-wise", () => {
  const a = { sampleRate: 10, channels: [new Float32Array([1, 2]), new Float32Array([0, 1])] };
  const b = { sampleRate: 10, channels: [new Float32Array([3, -2]), new Float32Array([1, 1])] };
  const s = sumAudio(a, b);
  expect([...s.channels[0]!]).toEqual([4, 0]);
  expect([...s.channels[1]!]).toEqual([1, 2]);
  expect(s.sampleRate).toBe(10);
});

test("repeatVote option: off = raw transcription, on (default) = repeat-voted", async () => {
  // Cheap fixture: oracle stems + oracle tempo (no separation), t1a/1 guitar stem.
  const bytes = async (p: string) => new Uint8Array(await Bun.file(join(root, p)).arrayBuffer());
  const off = await createPipeline({ demucs: await bytes("models/demucs/htdemucs_6s.onnx"), basicPitch: await bytes("models/basic-pitch/nmp.onnx") }, { repeatVote: false });
  const transcribe = await createTranscriber(await bytes("models/basic-pitch/nmp.onnx"));
  const d = await loadSongDir(join(root, "data/t1a/1"));
  const flags = { oracleStems: true, oracleTempo: true, oracleNotes: false };
  const input = { mix: d.mix, oracle: { stems: d.stems, truth: d.truth } };
  const gt = d.truth.tracks.find((t) => t.role === "guitar")!;
  const raw = await transcribe(d.stems.get(gt.name)!, { minPitch: 40, maxPitch: 88 });
  const strip = (ns: TimedNote[]) => ns.map((n) => [n.onset, n.offset, n.pitch]);
  const offNotes = (await off(input, flags)).output.tracks.find((t) => t.role === "guitar")!.notes;
  expect(strip(offNotes)).toEqual(strip(raw));
  const onRes = await run(input, flags);
  const onNotes = onRes.output.tracks.find((t) => t.role === "guitar")!.notes;
  expect(strip(onNotes)).toEqual(strip(repeatVote(raw, onRes.output.grid)));
}, SLOW);