//! Learned vocal separation for the lyric aligner — the "honest fix" the
//! vocals.rs reducer comment promises: a real separator lifting the voice
//! out of the band before CTC alignment, instead of band-pass pre-emphasis.
//!
//! Model: UVR-MDX-NET Voc_FT (KUIELab MDX-Net / TFC-TDF v2, the Ultimate
//! Vocal Remover project's vocal fine-tune; MIT with credit). Chosen
//! 2026-08-19 from the verified table in
//! research/lyrics-vocal-separation/candidates.md: 63.7 MiB, ONNX opset 13 —
//! under the inbox-WinML ceiling this machine PROVED by running the MMS
//! aligner (opset 14) — static `input`/`output` [1,4,3072,256] f32, vocals
//! SDR 9.64 on the MVSEP multisong benchmark (+1.5 dB over htdemucs at a
//! fifth of the size; the RoFormer/Demucs ONNX exports are opset 17–23 and
//! cannot load here as published).
//!
//! Posture is lyric_align.rs's exactly: WinML (no onnxruntime, no child
//! processes), download only from the user's explicit Settings click, the
//! session cached in a static — and NOTHING depends on this module. When
//! the model is absent or anything fails, the aligner keeps its band-passed
//! mixture input (the ctc-v4 behavior).
//!
//! The inference recipe is UVR's own non-overlap framing, kept bit-faithful
//! so the model sees what it was trained on: stereo 44.1 kHz → per-chunk
//! centered STFT (n_fft 7680, hop 1024, periodic Hann, reflect pad) → crop
//! to 3072 bins → model (channels [L_re, L_im, R_re, R_im]) → vocal
//! spectrogram → zero-pad the cropped bins back → WOLA iSTFT → trim
//! n_fft/2 from each end → butt-join. Each chunk contributes 253,440 fresh
//! samples (~5.75 s), so a four-minute track is ~42 model evaluations.

use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::Instant;

use realfft::num_complex::Complex;
use realfft::{ComplexToReal, RealFftPlanner, RealToComplex};
use tauri::Emitter;
use windows::{
    core::{Interface, HSTRING},
    AI::MachineLearning::{
        LearningModel, LearningModelBinding, LearningModelDevice, LearningModelDeviceKind,
        LearningModelSession, TensorFloat,
    },
    Win32::System::WinRT::{RoInitialize, RO_INIT_MULTITHREADED},
};

const MODEL_URL: &str =
    "https://github.com/TRvlvr/model_repo/releases/download/all_public_uvr_models/UVR-MDX-NET-Voc_FT.onnx";
/// Exact pinned size, shown in Settings and enforced after download — if the
/// asset is ever republished with different bytes, refuse rather than trust.
/// (Verified against the served bytes AND the file's own protobuf tail:
/// opset 13, `input`/`output` [batch,4,3072,256]. Kim_Vocal_2 is 66,759,214
/// — a first pin mistakenly used that sibling's size, and this guard caught
/// it exactly as designed.)
const MODEL_BYTES: u64 = 66_762_490;

// UVR MDX family constants (their model_data.json for this checkpoint:
// dim_f 3072, dim_t 2^8, n_fft 7680; hop is the architecture's fixed 1024).
const N_FFT: usize = 7_680;
const HOP: usize = 1_024;
const BINS: usize = N_FFT / 2 + 1; // 3841
const DIM_F: usize = 3_072;
const DIM_T: usize = 256;
/// Samples one model window covers: 256 centered STFT frames at hop 1024.
const CHUNK: usize = HOP * (DIM_T - 1); // 261,120
/// STFT edge effects live inside half a window; UVR trims this from each
/// chunk end and butt-joins the rest — no cross-fade, nothing to tune, and
/// the kept region sits in full window-overlap coverage by construction.
const TRIM: usize = N_FFT / 2; // 3840
/// Fresh samples each chunk contributes (~5.75 s at 44.1 kHz).
const GEN: usize = CHUNK - 2 * TRIM; // 253,440
/// Flat f32 count of the model's input and output: [1, 4, 3072, 256].
const TENSOR_LEN: usize = 4 * DIM_F * DIM_T;

struct Engine {
    session: LearningModelSession,
    /// Which device kind actually carried the smoke eval — for the log line.
    device: &'static str,
}

static ENGINE: Mutex<Option<Engine>> = Mutex::new(None);

fn model_file(app: &tauri::AppHandle) -> Result<PathBuf, String> {
    Ok(crate::lyric_align::models_dir(app)?.join("uvr-mdx-voc-ft.onnx"))
}

pub fn model_present(app: &tauri::AppHandle) -> bool {
    matches!(model_file(app), Ok(m) if m.is_file())
}

#[tauri::command]
pub fn vocal_sep_status(app: tauri::AppHandle) -> serde_json::Value {
    let size = model_file(&app)
        .ok()
        .and_then(|p| std::fs::metadata(p).ok())
        .map(|m| m.len())
        .unwrap_or(0);
    serde_json::json!({
        "present": model_present(&app),
        "bytes": size,
        "downloadBytes": MODEL_BYTES,
    })
}

/// Download the separation model. Runs ONLY from the user's explicit
/// Settings click — the consent-before-download charter rule. Streams to a
/// temp file, verifies the exact pinned byte count, renames atomically,
/// emits `lyrics:sep-model-progress` along the way.
#[tauri::command]
pub async fn vocal_sep_download(app: tauri::AppHandle) -> Result<(), String> {
    let client = reqwest::Client::builder()
        .user_agent("SoulStream/0.1 (vocal separation)")
        .build()
        .map_err(|e| e.to_string())?;
    let dest = model_file(&app)?;
    let tmp = dest.with_extension("onnx.downloading");
    let mut response = client
        .get(MODEL_URL)
        .send()
        .await
        .map_err(|e| e.to_string())?
        .error_for_status()
        .map_err(|e| e.to_string())?;
    let total = response.content_length().unwrap_or(MODEL_BYTES);
    let mut file = std::fs::File::create(&tmp).map_err(|e| e.to_string())?;
    let mut received: u64 = 0;
    let mut last_emit: u64 = 0;
    use std::io::Write;
    while let Some(chunk) = response.chunk().await.map_err(|e| e.to_string())? {
        file.write_all(&chunk).map_err(|e| e.to_string())?;
        received += chunk.len() as u64;
        if received - last_emit > 2 * 1024 * 1024 {
            last_emit = received;
            let _ = app.emit(
                "lyrics:sep-model-progress",
                serde_json::json!({ "received": received, "total": total }),
            );
        }
    }
    drop(file);
    if received != MODEL_BYTES {
        let _ = std::fs::remove_file(&tmp);
        return Err(format!(
            "download size mismatch: {received} bytes, expected {MODEL_BYTES}"
        ));
    }
    std::fs::rename(&tmp, &dest).map_err(|e| e.to_string())?;
    let _ = app.emit(
        "lyrics:sep-model-progress",
        serde_json::json!({ "received": received, "total": total, "done": true }),
    );
    Ok(())
}

fn ensure_engine(app: &tauri::AppHandle) -> Result<(), String> {
    let mut guard = ENGINE.lock().unwrap();
    if guard.is_some() {
        return Ok(());
    }
    let model_path = model_file(app)?;
    if !model_path.is_file() {
        return Err("separation model not downloaded".into());
    }
    // Same WinRT posture as the aligner: initialize, tolerate an existing
    // apartment, never uninitialize someone else's.
    match unsafe { RoInitialize(RO_INIT_MULTITHREADED) } {
        Ok(()) => {}
        Err(e) if e.code().0 as u32 == 0x8001_0106 => {}
        Err(e) => return Err(format!("WinRT init failed: {e}")),
    }
    let model =
        LearningModel::LoadFromFilePath(&HSTRING::from(model_path.to_string_lossy().as_ref()))
            .map_err(|e| format!("could not load separation model: {e}"))?;
    // DirectX first (DirectML — the discrete GPU when there is one): a track
    // is ~42 conv-net evaluations and the GPU turns minutes into seconds.
    // The zeros smoke eval proves the device can actually RUN the graph —
    // "loads but cannot execute" is a real DML failure mode — then Cpu is
    // the quiet fallback, so this works everywhere the aligner does.
    for (kind, name) in [
        (LearningModelDeviceKind::DirectX, "DirectX"),
        (LearningModelDeviceKind::Cpu, "Cpu"),
    ] {
        let Ok(device) = LearningModelDevice::Create(kind) else {
            continue;
        };
        let Ok(session) = LearningModelSession::CreateFromModelOnDevice(&model, &device) else {
            continue;
        };
        let zeros = vec![0.0f32; TENSOR_LEN];
        if eval_chunk(&session, &zeros).is_ok() {
            *guard = Some(Engine {
                session,
                device: name,
            });
            return Ok(());
        }
    }
    Err("no WinML device could run the separation model".into())
}

fn eval_chunk(session: &LearningModelSession, flat: &[f32]) -> Result<Vec<f32>, String> {
    let input = TensorFloat::CreateFromShapeArrayAndDataArray(
        &[1, 4, DIM_F as i64, DIM_T as i64],
        flat,
    )
    .map_err(|e| format!("could not create input tensor: {e}"))?;
    let binding = LearningModelBinding::CreateFromSession(session)
        .map_err(|e| format!("could not create binding: {e}"))?;
    binding
        .Bind(&HSTRING::from("input"), &input)
        .map_err(|e| format!("could not bind input: {e}"))?;
    let result = session
        .Evaluate(&binding, &HSTRING::from("vocal-sep"))
        .map_err(|e| format!("inference failed: {e}"))?;
    if !result.Succeeded().unwrap_or(false) {
        return Err("inference did not succeed".into());
    }
    let out: TensorFloat = result
        .Outputs()
        .and_then(|o| o.Lookup(&HSTRING::from("output")))
        .map_err(|e| format!("output missing: {e}"))?
        .cast()
        .map_err(|e| format!("output wrong type: {e}"))?;
    let values: Vec<f32> = out
        .GetAsVectorView()
        .map_err(|e| format!("could not read output: {e}"))?
        .into_iter()
        .collect();
    if values.len() != TENSOR_LEN {
        return Err(format!("unexpected output length {}", values.len()));
    }
    Ok(values)
}

// ---- STFT plumbing (pure DSP — no WinML below this line) -------------------

struct Stft {
    r2c: Arc<dyn RealToComplex<f32>>,
    c2r: Arc<dyn ComplexToReal<f32>>,
    /// Periodic Hann — torch.hann_window's default, what the model trained on.
    window: Vec<f32>,
    /// WOLA denominator Σ w² over the padded chunk layout. The 256 frames
    /// tile the padded buffer exactly, so it is identical for every chunk.
    envelope: Vec<f32>,
}

impl Stft {
    fn new() -> Self {
        let mut planner = RealFftPlanner::<f32>::new();
        let r2c = planner.plan_fft_forward(N_FFT);
        let c2r = planner.plan_fft_inverse(N_FFT);
        let window: Vec<f32> = (0..N_FFT)
            .map(|n| 0.5 - 0.5 * (2.0 * std::f32::consts::PI * n as f32 / N_FFT as f32).cos())
            .collect();
        let mut envelope = vec![1e-8f32; CHUNK + 2 * TRIM];
        for t in 0..DIM_T {
            for (n, w) in window.iter().enumerate() {
                envelope[t * HOP + n] += w * w;
            }
        }
        Stft {
            r2c,
            c2r,
            window,
            envelope,
        }
    }

    /// Centered STFT of one CHUNK-sized segment: reflect-pad n_fft/2 each
    /// side (torch.stft center=True), frames at the hop — frame-major
    /// [t * BINS + f].
    fn forward(&self, chunk: &[f32]) -> Result<Vec<Complex<f32>>, String> {
        debug_assert_eq!(chunk.len(), CHUNK);
        let mut padded = vec![0.0f32; CHUNK + 2 * TRIM];
        padded[TRIM..TRIM + CHUNK].copy_from_slice(chunk);
        for i in 1..=TRIM {
            padded[TRIM - i] = chunk[i];
            padded[TRIM + CHUNK - 1 + i] = chunk[CHUNK - 1 - i];
        }
        let mut spec = vec![Complex::default(); DIM_T * BINS];
        let mut frame = vec![0.0f32; N_FFT];
        for t in 0..DIM_T {
            for (n, w) in self.window.iter().enumerate() {
                frame[n] = padded[t * HOP + n] * w;
            }
            self.r2c
                .process(&mut frame, &mut spec[t * BINS..(t + 1) * BINS])
                .map_err(|e| format!("stft failed: {e}"))?;
        }
        Ok(spec)
    }

    /// WOLA inverse of a frame-major spectrogram back to CHUNK samples with
    /// the centered pad trimmed off — torch.istft's window-envelope
    /// normalization. Consumes the spectrogram (the FFT scratches it).
    fn inverse(&self, spec: &mut [Complex<f32>]) -> Result<Vec<f32>, String> {
        let mut acc = vec![0.0f32; CHUNK + 2 * TRIM];
        let mut frame = vec![0.0f32; N_FFT];
        let scale = 1.0 / N_FFT as f32; // realfft's c2r is unnormalized
        for t in 0..DIM_T {
            let row = &mut spec[t * BINS..(t + 1) * BINS];
            // The model emits free re/im everywhere; a real inverse needs
            // pure-real DC and Nyquist. Nyquist sits in the cropped-zero
            // region already, DC gets clamped here.
            row[0].im = 0.0;
            row[BINS - 1].im = 0.0;
            self.c2r
                .process(row, &mut frame)
                .map_err(|e| format!("istft failed: {e}"))?;
            for (n, w) in self.window.iter().enumerate() {
                acc[t * HOP + n] += frame[n] * scale * w;
            }
        }
        Ok((TRIM..TRIM + CHUNK)
            .map(|i| acc[i] / self.envelope[i])
            .collect())
    }
}

/// Pack two frame-major spectrograms into the model's flat input — channel
/// order [L_re, L_im, R_re, R_im] (UVR's reshape/permute order), frequency
/// cropped to DIM_F.
fn pack_input(spec_l: &[Complex<f32>], spec_r: &[Complex<f32>]) -> Vec<f32> {
    let mut flat = vec![0.0f32; TENSOR_LEN];
    for (c, spec) in [spec_l, spec_r].into_iter().enumerate() {
        for f in 0..DIM_F {
            for t in 0..DIM_T {
                let z = spec[t * BINS + f];
                flat[((2 * c) * DIM_F + f) * DIM_T + t] = z.re;
                flat[((2 * c + 1) * DIM_F + f) * DIM_T + t] = z.im;
            }
        }
    }
    flat
}

/// Unpack the model's flat output back into two frame-major spectrograms.
/// Bins at and above DIM_F stay zero — the crop restored as silence, which
/// for a vocal stem is honest (nothing sung lives above ~17.6 kHz).
fn unpack_output(flat: &[f32]) -> (Vec<Complex<f32>>, Vec<Complex<f32>>) {
    let mut left = vec![Complex::default(); DIM_T * BINS];
    let mut right = vec![Complex::default(); DIM_T * BINS];
    for (c, spec) in [&mut left, &mut right].into_iter().enumerate() {
        for f in 0..DIM_F {
            for t in 0..DIM_T {
                spec[t * BINS + f] = Complex::new(
                    flat[((2 * c) * DIM_F + f) * DIM_T + t],
                    flat[((2 * c + 1) * DIM_F + f) * DIM_T + t],
                );
            }
        }
    }
    (left, right)
}

/// Chunk a stereo 44.1 kHz track through `run` (the [1,4,3072,256] →
/// same-shape spectrogram transform) and reassemble the returned stem.
/// UVR's non-overlap framing: TRIM zeros lead the track so chunk c's kept
/// middle lands exactly at c·GEN, edge effects fall in the trimmed halves,
/// and segments butt-join.
fn process_chunks(
    stft: &Stft,
    left: &[f32],
    right: &[f32],
    mut run: impl FnMut(&[f32]) -> Result<Vec<f32>, String>,
) -> Result<(Vec<f32>, Vec<f32>), String> {
    let len = left.len().min(right.len());
    if len < HOP {
        return Err("audio too short to separate".into());
    }
    let chunks = len.div_ceil(GEN);
    let padded_len = TRIM + chunks * GEN + TRIM;
    let mut mix_l = vec![0.0f32; padded_len];
    let mut mix_r = vec![0.0f32; padded_len];
    mix_l[TRIM..TRIM + len].copy_from_slice(&left[..len]);
    mix_r[TRIM..TRIM + len].copy_from_slice(&right[..len]);
    let mut stem_l = vec![0.0f32; chunks * GEN];
    let mut stem_r = vec![0.0f32; chunks * GEN];
    for c in 0..chunks {
        let s = c * GEN;
        let spec_l = stft.forward(&mix_l[s..s + CHUNK])?;
        let spec_r = stft.forward(&mix_r[s..s + CHUNK])?;
        let out = run(&pack_input(&spec_l, &spec_r))?;
        if out.len() != TENSOR_LEN {
            return Err(format!("unexpected chunk output length {}", out.len()));
        }
        let (mut vl, mut vr) = unpack_output(&out);
        let wl = stft.inverse(&mut vl)?;
        let wr = stft.inverse(&mut vr)?;
        stem_l[s..s + GEN].copy_from_slice(&wl[TRIM..TRIM + GEN]);
        stem_r[s..s + GEN].copy_from_slice(&wr[TRIM..TRIM + GEN]);
        // Politeness between chunks — this thread must never contend with
        // the audio path for sustained stretches (the passage loop's rule).
        std::thread::sleep(std::time::Duration::from_millis(3));
    }
    stem_l.truncate(len);
    stem_r.truncate(len);
    Ok((stem_l, stem_r))
}

/// Whole-file vocal stem as 16 kHz mono for the CTC aligner: decode →
/// chunked WinML separation → mid fold + anti-alias resample (player crate).
/// Fails soft by design — the caller falls back to the band-passed mixture.
pub fn separate_vocals_16k(
    app: &tauri::AppHandle,
    path: &std::path::Path,
) -> Result<Vec<f32>, String> {
    if !model_present(app) {
        return Err("separation model not downloaded".into());
    }
    let [left, right] = player::vocals::separation_input_44k1(path)?;
    ensure_engine(app)?;
    let guard = ENGINE.lock().unwrap();
    let engine = guard.as_ref().ok_or("engine unavailable")?;
    let stft = Stft::new();
    let started = Instant::now();
    let (stem_l, stem_r) = process_chunks(&stft, &left, &right, |flat| {
        eval_chunk(&engine.session, flat)
    })?;
    log::info!(
        "vocal_sep: {:.0}s of audio separated in {:.1}s on {}",
        stem_l.len() as f32 / 44_100.0,
        started.elapsed().as_secs_f32(),
        engine.device
    );
    Ok(player::vocals::stem_pcm_16k(&stem_l, &stem_r))
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The WOLA plumbing alone — forward STFT straight back through the
    /// inverse, no crop — must be per-sample exact: the numerator and the
    /// envelope share the same window algebra, so any real error here is a
    /// tiling or normalization bug, not numerics.
    #[test]
    fn stft_wola_roundtrip_is_exact() {
        let chunk: Vec<f32> = (0..CHUNK)
            .map(|i| (2.0 * std::f32::consts::PI * 440.0 * i as f32 / 44_100.0).sin() * 0.5)
            .collect();
        let stft = Stft::new();
        let mut spec = stft.forward(&chunk).unwrap();
        let back = stft.inverse(&mut spec).unwrap();
        let worst = chunk
            .iter()
            .zip(&back)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(worst < 1e-4, "un-cropped WOLA must be exact, worst {worst}");
    }

    /// An identity model (output spectrogram = input spectrogram) must hand
    /// the track back through the whole STFT → crop → pack → unpack →
    /// zero-pad → iSTFT → trim → butt-join pipeline: proves the frame
    /// tiling, the chunk joins, and the layout round-trip all at once.
    /// Tolerance: the reflect pad puts a derivative kink at each chunk
    /// edge whose broadband energy reaches above bin 3072, and the model
    /// crop discards it — torch.stft + UVR's own framing lose exactly the
    /// same (~1% here); the model trained on that world.
    #[test]
    fn identity_model_reconstructs_across_chunk_joins() {
        let len = GEN + GEN / 2; // two chunks, one join, one truncation
        let track: Vec<f32> = (0..len)
            .map(|i| (2.0 * std::f32::consts::PI * 440.0 * i as f32 / 44_100.0).sin() * 0.5)
            .collect();
        let stft = Stft::new();
        let (out_l, out_r) =
            process_chunks(&stft, &track, &track, |flat| Ok(flat.to_vec())).unwrap();
        assert_eq!(out_l.len(), len);
        assert_eq!(out_r.len(), len);
        let mut worst = 0.0f32;
        for (a, b) in track.iter().zip(&out_l) {
            worst = worst.max((a - b).abs());
        }
        assert!(
            worst < 2e-2,
            "identity round-trip must reconstruct the tone, worst err {worst}"
        );
        // The join at GEN must not be a seam beyond that same crop error.
        let seam = (GEN - 50..GEN + 50)
            .map(|i| (track[i] - out_l[i]).abs())
            .fold(0.0f32, f32::max);
        assert!(seam < 2e-2, "butt-join must be seamless, err {seam}");
    }

    /// The kept region of every chunk sits in full window-overlap coverage:
    /// the WOLA envelope there is flat and far from zero.
    #[test]
    fn envelope_is_stable_where_samples_are_kept() {
        let stft = Stft::new();
        let kept = &stft.envelope[2 * TRIM..2 * TRIM + GEN];
        let min = kept.iter().cloned().fold(f32::INFINITY, f32::min);
        let max = kept.iter().cloned().fold(0.0f32, f32::max);
        assert!(min > 2.0, "envelope must be well-conditioned, min {min}");
        assert!(
            max / min < 1.5,
            "envelope must be near-flat in the kept region, {min}..{max}"
        );
    }
}