import { test, expect, beforeAll } from "bun:test";
import { alignmentFit, autoAlign, bandedWithRetry, barTimes, beatTimes, constantOffsetBaseline, fitFromCues, meterGrid, recordingCues, rejectOffsetOutliers } from "../../src/align/auto.ts";
import { makeWarp, type Anchor, type Warp } from "../../src/align/warp.ts";
import type { RefTrack, Smf } from "../../src/align/smf.ts";
import type { Audio, Role, TruthTrack } from "../../src/core/types.ts";
import { buildSong, loadSongDir } from "../../src/bench/dataset.ts";

const SF2 = new Uint8Array(await Bun.file("node_modules/@coderline/alphatab/dist/soundfont/sonivox.sf2").arrayBuffer());
const BIG = 60_000;

let song: Awaited<ReturnType<typeof loadSongDir>>;
beforeAll(async () => { song = await loadSongDir(await buildSong("t1c", 1)); }, BIG);

/** In-memory Smf from truth tracks: `map` turns truth (= recording) time into MIDI time; `reps` tiles the song. */
function smfFrom(truth: TruthTrack[], programs: number[], bpm: number, map: (t: number) => number, reps = 1, period = 0): Smf {
  const tracks: RefTrack[] = truth.map((t, i) => ({
    index: i + 1, name: t.name, program: programs[i]!, isDrums: t.role === "drums",
    notes: Array.from({ length: reps }, (_, r) => t.notes.map((n) => ({
      onset: map(n.onset + r * period), offset: map(n.offset + r * period), pitch: n.pitch, velocity: n.velocity ?? 96,
      channel: t.role === "drums" ? 9 : i,
    }))).flat(),
  }));
  const durationSec = Math.max(...tracks.flatMap((t) => t.notes.map((n) => n.offset)));
  return { durationSec, tracks, tempoMap: [{ sec: 0, bpm }], timeSigs: [{ sec: 0, num: 4, den: 4 }] };
}
const rolesOf = (truth: TruthTrack[]): Record<string, Role | null> => Object.fromEntries(truth.map((t) => [t.name, t.role]));
const durOf = (a: Audio) => a.channels[0]!.length / a.sampleRate;
const median = (xs: number[]) => { const s = [...xs].sort((a, b) => a - b); return s[s.length >> 1]!; };
const p90 = (xs: number[]) => { const s = [...xs].sort((a, b) => a - b); return s[Math.floor(s.length * 0.9)]!; };
const debug = (msg: string) => { if (process.env.ALIGN_DEBUG) console.log(msg); };
const strictlyMonotone = (A: Anchor[]) => A.every((a, k) => k === 0 || (a.midiSec > A[k - 1]!.midiSec && a.audioSec > A[k - 1]!.audioSec));

/** Known distortion recording time -> MIDI time: +0.4 s, then bars 4–10 at ×1.06 tempo, then ≈ −0.8 s (continuous). */
function distortion(bpm: number): Warp {
  const bar = (4 * 60) / bpm, b1 = 4 * bar, b2 = 10 * bar;
  const m2 = b1 + 0.4 + (b2 - b1) / 1.06;
  return makeWarp([
    { midiSec: 0, audioSec: 0.4, kind: "manual" },
    { midiSec: b1, audioSec: b1 + 0.4, kind: "manual" },
    { midiSec: b2, audioSec: m2, kind: "manual" },
    { midiSec: b2 + 100, audioSec: m2 + 100, kind: "manual" },
  ]); // toAudio: recording -> MIDI, toMidi: MIDI -> recording (the truth the aligner must recover)
}

function recoveryErrors(smf: Smf, got: Anchor[], trueMidiToAudio: (m: number) => number, audioDur: number): number[] {
  const w = makeWarp(got);
  const onsets = [...new Set(smf.tracks.flatMap((t) => t.notes.map((n) => n.onset)))];
  return onsets.filter((m) => { const a = trueMidiToAudio(m); return a >= 0 && a <= audioDur; })
    .map((m) => Math.abs(w.toAudio(m) - trueMidiToAudio(m)));
}

test("recovers a known piecewise warp: median error < 40 ms, fit precision ≥ 0.8, DTW beats the constant baseline", () => {
  const D = distortion(song.song.bpm);
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, D.toAudio);
  const t0 = performance.now();
  const { anchors, fit, baseline, chosen } = autoAlign(song.mix, smf, rolesOf(song.truth.tracks), [], SF2);
  const ms = performance.now() - t0;
  const err = recoveryErrors(smf, anchors, D.toMidi, durOf(song.mix));
  debug(`synthetic warp: ${anchors.length} anchors, median err ${(median(err) * 1000).toFixed(1)} ms, p90 ${(p90(err) * 1000).toFixed(1)} ms, fit ${fit.precision?.toFixed(3)} lift ${fit.lift?.toFixed(3)} / ${fit.medianResidualMs?.toFixed(1)} ms, baseline ${baseline.offsetSec} s ${baseline.precision?.toFixed(3)}, ${ms.toFixed(0)} ms`);
  expect(chosen).toBe("dtw");
  expect(anchors.every((a) => a.kind === "auto")).toBe(true);
  expect(median(err)).toBeLessThan(0.04);
  expect(p90(err)).toBeLessThan(0.08);
  expect(fit.precision!).toBeGreaterThanOrEqual(0.8);
  expect(fit.precision!).toBeGreaterThanOrEqual(baseline.precision! + 0.02);
  expect(fit.windows.some((w) => w.flagged)).toBe(false);
  expect(fit.stretches!.total).toBe(anchors.length + 1);
  expect(fit.windows.length).toBe(2);
  expect(anchors.every((a) => a.audioSec > 0 || a.midiSec === 0)).toBe(true);
}, BIG);

test("a correct manual anchor mid-song is honoured exactly", () => {
  const D = distortion(song.song.bpm);
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, D.toAudio);
  const manual: Anchor = { midiSec: D.toAudio(30), audioSec: 30, kind: "manual" };
  const { anchors } = autoAlign(song.mix, smf, rolesOf(song.truth.tracks), [manual], SF2);
  expect(anchors.filter((a) => a.kind === "manual")).toEqual([manual]);
  expect(makeWarp(anchors).toAudio(manual.midiSec)).toBe(30);
  // no auto anchor within 1 s (MIDI time) of the manual one
  expect(anchors.some((a) => a.kind === "auto" && Math.abs(a.midiSec - manual.midiSec) < 1)).toBe(false);
  expect(anchors.length).toBeGreaterThan(5);
  const err = recoveryErrors(smf, anchors, D.toMidi, durOf(song.mix));
  expect(median(err)).toBeLessThan(0.04);
}, BIG);

test("a silent intro in the recording does not drag the path", () => {
  const pad = 4 * song.mix.sampleRate;
  const mix: Audio = { sampleRate: song.mix.sampleRate, channels: song.mix.channels.map((c) => { const o = new Float32Array(c.length + pad); o.set(c, pad); return o; }) };
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, (t) => t);
  const { anchors, fit } = autoAlign(mix, smf, rolesOf(song.truth.tracks), [], SF2);
  const err = recoveryErrors(smf, anchors, (m) => m + 4, durOf(mix));
  expect(median(err)).toBeLessThan(0.04);
  expect(fit.precision!).toBeGreaterThanOrEqual(0.8);
}, BIG);

test("5-minute input aligns in < 15 s", () => {
  const reps = 5, period = durOf(song.mix);
  const mix: Audio = { sampleRate: song.mix.sampleRate, channels: song.mix.channels.map((c) => { const o = new Float32Array(c.length * reps); for (let r = 0; r < reps; r++) o.set(c, r * c.length); return o; }) };
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, (t) => t + 0.3, reps, period);
  const t0 = performance.now();
  const { anchors, fit, baseline, chosen } = autoAlign(mix, smf, rolesOf(song.truth.tracks), [], SF2);
  const sec = (performance.now() - t0) / 1000;
  debug(`5-min (${durOf(mix).toFixed(0)} s) input: ${sec.toFixed(2)} s, ${anchors.length} anchors, fit ${fit.precision?.toFixed(3)}, ${chosen}, baseline ${baseline.offsetSec} ${baseline.precision?.toFixed(3)}`);
  expect(durOf(mix)).toBeGreaterThan(290);
  expect(sec).toBeLessThan(15);
  expect(fit.windows.length).toBe(10);
  // a constant offset: DTW cannot beat the baseline by 0.02, so the single-anchor baseline is returned
  expect(baseline.offsetSec).toBeCloseTo(-0.3, 1);
  expect(chosen).toBe("baseline");
  expect(anchors).toHaveLength(1);
  expect(anchors[0]!.kind).toBe("auto");
  expect(anchors[0]!.audioSec - anchors[0]!.midiSec).toBeCloseTo(baseline.offsetSec, 6);
}, BIG);

test("barTimes follows tempo and time-signature changes", () => {
  const smf: Smf = {
    durationSec: 10, tracks: [],
    tempoMap: [{ sec: 0, bpm: 120 }, { sec: 4, bpm: 60 }],
    timeSigs: [{ sec: 0, num: 4, den: 4 }, { sec: 4, num: 3, den: 4 }],
  };
  // 4/4 @120: bars of 2 s at 0, 2; then 3/4 @60: bars of 3 s at 4, 7
  expect(barTimes(smf)).toEqual([0, 2, 4, 7]);
  expect(barTimes({ durationSec: 3, tracks: [], tempoMap: [{ sec: 0, bpm: 120 }], timeSigs: [] })).toEqual([0, 2]); // default 4/4
  expect(barTimes({ durationSec: 4, tracks: [], tempoMap: [{ sec: 0, bpm: 120 }], timeSigs: [{ sec: 0, num: 6, den: 8 }] })).toEqual([0, 1.5, 3]);
});

test("fine pass widens the band when the end cell is unreachable, then gives up with null", () => {
  const cost = (i: number, j: number) => Math.abs(i - j);
  // centre jumps by 120 frames mid-way: ±25 and ±50 leave a gap between rows, ±100 bridges it
  const jump = (i: number) => (i < 200 ? i : i + 120);
  const p = bandedWithRetry(cost, 400, 600, jump);
  expect(p).not.toBeNull();
  expect(p![p!.length - 1]).toEqual([399, 599]);
  // a 400-frame jump defeats even ±2 s (±100 frames)
  expect(bandedWithRetry(cost, 400, 1000, (i) => (i < 200 ? i : i + 400))).toBeNull();
});

test("fit exposes a bad alignment instead of hiding it", () => {
  const D = distortion(song.song.bpm);
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, D.toAudio);
  const roles = rolesOf(song.truth.tracks);
  const bad = alignmentFit(song.mix, smf, roles, []); // identity: off by 0.4–0.8 s throughout
  expect(bad.precision!).toBeLessThan(0.6);
  // the −0.8 s section is exposed; the +0.4 s section is not, as 0.4 s is ≈ 40 ms from two 16ths at 68 bpm (grid alias)
  expect(bad.windows[1]!.flagged).toBe(true);
  const good = alignmentFit(song.mix, smf, roles, autoAlign(song.mix, smf, roles, [], SF2).anchors);
  expect(good.windows.some((w) => w.flagged)).toBe(false);
  // unmapped tracks are excluded: with no scored roles there is nothing to measure
  expect(alignmentFit(song.mix, smf, Object.fromEntries(Object.keys(roles).map((k) => [k, null])), []).precision).toBeNull();
}, BIG);

test("a deliberately wrong manual anchor (+0.3 s) is kept exactly and the result stays monotonic", () => {
  const D = distortion(song.song.bpm);
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, D.toAudio);
  const manual: Anchor = { midiSec: D.toAudio(30), audioSec: 30.3, kind: "manual" };
  const { anchors } = autoAlign(song.mix, smf, rolesOf(song.truth.tracks), [manual], SF2);
  expect(anchors.filter((a) => a.kind === "manual")).toEqual([manual]);
  expect(makeWarp(anchors).toAudio(manual.midiSec)).toBe(30.3);
  expect(strictlyMonotone(anchors)).toBe(true);
  expect(anchors.every((a) => a.audioSec > 0 || a.midiSec === 0)).toBe(true);
  // away from the bad anchor (> 4 s), the autos still follow the truth
  const far = anchors.filter((a) => a.kind === "auto" && Math.abs(a.audioSec - 30) > 4);
  expect(far.length).toBeGreaterThan(5);
  expect(median(far.map((a) => Math.abs(a.audioSec - D.toMidi(a.midiSec))))).toBeLessThan(0.04);
}, BIG);

test("when every fine band fails, the coarse-path fallback still yields auto anchors", () => {
  const D = distortion(song.song.bpm);
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, D.toAudio);
  const manual: Anchor = { midiSec: D.toAudio(30), audioSec: 30, kind: "manual" };
  const { anchors, fit, chosen } = autoAlign(song.mix, smf, rolesOf(song.truth.tracks), [manual], SF2, { fineBandsSec: [] });
  expect(chosen).toBe("dtw"); // manual anchors present: never substituted by the baseline
  const autos = anchors.filter((a) => a.kind === "auto");
  // both segments (before and after the manual anchor) get anchors
  expect(autos.filter((a) => a.audioSec < 30).length).toBeGreaterThan(3);
  expect(autos.filter((a) => a.audioSec > 30).length).toBeGreaterThan(3);
  expect(strictlyMonotone(anchors)).toBe(true);
  const err = recoveryErrors(smf, anchors, D.toMidi, durOf(song.mix));
  debug(`coarse fallback: ${autos.length} autos, median err ${(median(err) * 1000).toFixed(1)} ms, p90 ${(p90(err) * 1000).toFixed(1)} ms, fit ${fit.precision?.toFixed(3)}`);
  expect(median(err)).toBeLessThan(0.06); // 10 Hz resolution: coarser than the fine pass, still hand-correctable
  expect(p90(err)).toBeLessThan(0.12);
}, BIG);

test("fitFromCues with cached recordingCues equals alignmentFit", () => {
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, (t) => t);
  const roles = rolesOf(song.truth.tracks);
  const anchors: Anchor[] = [{ midiSec: 1, audioSec: 1.05, kind: "manual" }, { midiSec: 20, audioSec: 20.1, kind: "manual" }];
  const cues = recordingCues(song.mix);
  expect(cues.onsets.length).toBeGreaterThan(10);
  expect(cues.chroma.length).toBeGreaterThan(cues.durSec * cues.chromaRate - 2);
  expect(fitFromCues(cues, smf, roles, anchors)).toEqual(alignmentFit(song.mix, smf, roles, anchors));
});

// ---- ruling A7: calibrated fit, beat-shift check, constant-offset baseline ----

test("fit precision: correct alignment ≥ 0.8, a 0.12 s off-grid shift < 0.5 (and flagged)", () => {
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, (t) => t);
  const roles = rolesOf(song.truth.tracks);
  const good = alignmentFit(song.mix, smf, roles, []);
  const off = alignmentFit(song.mix, smf, roles, [{ midiSec: 10, audioSec: 10.12, kind: "manual" }]);
  debug(`precision: correct ${good.precision?.toFixed(3)} (lift ${good.lift?.toFixed(3)}, audible ${good.audibleRate?.toFixed(3)}), +0.12 s ${off.precision?.toFixed(3)} (lift ${off.lift?.toFixed(3)})`);
  expect(good.precision!).toBeGreaterThanOrEqual(0.8);
  expect(good.lift!).toBeGreaterThanOrEqual(0.25);
  expect(good.windows.some((w) => w.flagged)).toBe(false);
  expect(off.precision!).toBeLessThan(0.5);
  expect(off.windows.filter((w) => w.precision !== null).every((w) => w.flagged)).toBe(true);
});

test("beat-shift check flags an alignment that is one beat late, though onset precision cannot see it", () => {
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, (t) => t);
  const roles = rolesOf(song.truth.tracks);
  const beat = 60 / song.song.bpm;
  const late = alignmentFit(song.mix, smf, roles, [{ midiSec: 10, audioSec: 10 + beat, kind: "manual" }]);
  const good = alignmentFit(song.mix, smf, roles, []);
  debug(`beat shift: late precision ${late.precision?.toFixed(3)}, windows ${JSON.stringify(late.windows.map((w) => [w.precision?.toFixed(2), w.lift?.toFixed(2), w.beatShift]))}`);
  expect(good.windows.some((w) => w.beatShift)).toBe(false);
  const shifted = late.windows.filter((w) => w.beatShift);
  expect(shifted.length).toBeGreaterThan(0);
  expect(shifted.every((w) => w.flagged && w.reasons.some((r) => r.startsWith("beat shift suspected: −")))).toBe(true);
});

test("constantOffsetBaseline grid-searches −2..+2 s for the offset with the best precision", () => {
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, (t) => t + 0.27);
  const b = constantOffsetBaseline(recordingCues(song.mix), smf, rolesOf(song.truth.tracks));
  expect(b.offsetSec).toBeCloseTo(-0.27, 1);
  expect(Math.round(b.offsetSec * 100)).toBe(b.offsetSec * 100); // 10 ms grid
  expect(b.precision!).toBeGreaterThanOrEqual(0.8);
});

// ---- ruling A8: aligner guards ----

test("rejectOffsetOutliers drops bars > ¼ beat from the robust median of ±3 neighbours (slope-transported), iterating until stable", () => {
  const bars = Array.from({ length: 20 }, (_, k) => ({ midiSec: k * 2, off: 0.1 + k * 0.004, beatSec: 0.5 }));
  bars[5]!.off = 0.6; // a beat out
  bars[6]!.off = 0.55; // its neighbour, also out
  bars[15]!.off = -0.3;
  const kept = rejectOffsetOutliers(bars).map((b) => b.midiSec / 2);
  expect(kept).toEqual(Array.from({ length: 20 }, (_, k) => k).filter((k) => k !== 5 && k !== 6 && k !== 15));
  // a smooth ramp (genuine tempo change) is kept whole
  const ramp = Array.from({ length: 20 }, (_, k) => ({ midiSec: k * 2, off: k * 0.05, beatSec: 0.5 }));
  expect(rejectOffsetOutliers(ramp)).toHaveLength(20);
  // ... also a steep one (10% tempo difference: 0.2 s per 2 s bar), whose ends a plain one-sided median would erode
  // bar by bar; neighbours are transported along their median (Theil–Sen) slope
  const steep = Array.from({ length: 20 }, (_, k) => ({ midiSec: k * 2, off: k * 0.2, beatSec: 0.5 }));
  expect(rejectOffsetOutliers(steep)).toHaveLength(20);
  // a beat jump inside the steep ramp is still rejected
  steep[10]!.off += 0.5;
  expect(rejectOffsetOutliers(steep).map((b) => b.midiSec / 2)).toEqual(Array.from({ length: 20 }, (_, k) => k).filter((k) => k !== 10));
});

test("meterGrid: beats and bars from tempo/time signatures; the final bar is partial when the MIDI ends mid-bar", () => {
  const smf: Smf = { durationSec: 9, tracks: [], tempoMap: [{ sec: 0, bpm: 120 }, { sec: 4, bpm: 60 }], timeSigs: [{ sec: 0, num: 4, den: 4 }, { sec: 4, num: 3, den: 4 }] };
  const g = meterGrid(smf);
  expect(g.bars).toEqual([0, 2, 4, 7]);
  expect(g.beats).toEqual([0, 0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4, 5, 6, 7, 8, 9]);
  expect(g.lastBarPartial).toBe(true); // bar 7..10 ends after 9
  expect(meterGrid({ ...smf, durationSec: 10 }).lastBarPartial).toBe(false);
  expect(beatTimes({ durationSec: 3, tracks: [], tempoMap: [{ sec: 0, bpm: 120 }], timeSigs: [{ sec: 0, num: 6, den: 8 }] })).toEqual([0, 0.25, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.25, 2.5, 2.75]);
});

test("no auto anchor in the final partial bar", () => {
  const D = distortion(song.song.bpm);
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, D.toAudio);
  const g = meterGrid(smf);
  const last = g.bars[g.bars.length - 1]!;
  const cut: Smf = { ...smf, durationSec: last + 0.5 * (last - g.bars[g.bars.length - 2]!) }; // ends half-way into the last bar
  expect(meterGrid(cut).lastBarPartial).toBe(true);
  const { anchors } = autoAlign(song.mix, cut, rolesOf(song.truth.tracks), [], SF2);
  expect(anchors.some((a) => a.midiSec >= last - 1e-9)).toBe(false);
}, BIG);

test("12 s of extra outro in the recording: open-end DTW, no ramp at the end", () => {
  const D = distortion(song.song.bpm);
  const smf = smfFrom(song.truth.tracks, song.song.tracks.map((t) => t.program), song.song.bpm, D.toAudio);
  const sr = song.mix.sampleRate, extra = 12 * sr;
  // the outro repeats the first 12 s of the song (material that DTW could be tempted to match)
  const mix: Audio = { sampleRate: sr, channels: song.mix.channels.map((c) => { const o = new Float32Array(c.length + extra); o.set(c); o.set(c.subarray(0, extra), c.length); return o; }) };
  const { anchors, chosen } = autoAlign(mix, smf, rolesOf(song.truth.tracks), [], SF2);
  expect(chosen).toBe("dtw");
  const autos = anchors.filter((a) => a.kind === "auto");
  const tail = autos.slice(-4).map((a) => a.audioSec - D.toMidi(a.midiSec));
  debug(`outro: last anchor offsets vs truth ${tail.map((e) => (e * 1000).toFixed(0)).join(", ")} ms`);
  expect(tail.every((e) => Math.abs(e) < 0.06)).toBe(true);
  const w = makeWarp(anchors);
  const endErr = Math.abs(w.toAudio(smf.durationSec) - D.toMidi(smf.durationSec));
  expect(endErr).toBeLessThan(0.1);
  const err = recoveryErrors(smf, anchors, D.toMidi, durOf(song.mix));
  expect(median(err)).toBeLessThan(0.04);
}, BIG);