//! Tier B word alignment: CTC forced alignment with a small acoustic model
//! through WinML — the same built-in runtime the hero detector uses, so no
//! new dependency and no child processes (the preparation-worker charter).
//!
//! The model (MMS-300m forced-aligner, ONNX f32 export — Meta's multilingual
//! wav2vec2 fine-tuned specifically FOR forced alignment, the checkpoint
//! torchaudio's alignment API standardized on) emits per-20ms CHARACTER
//! probabilities; a Viterbi pass aligns the KNOWN lyric text against them —
//! no transcription, no guessing, just "when does each known word occur".
//! Its vocab is lowercase romanized characters with NO word separator and
//! blank = `<blank>`; the loader reads all of that from the repo's own
//! vocab.json rather than hardcoding (the vocab-never-hardcoded rule).
//! It is downloaded ONLY from the user's explicit click in Settings
//! (`lyric_model_download`) — the charter's consent rule — verified by
//! length, and loaded lazily. When it is absent or fails, the worker's
//! envelope estimate stands; nothing ever depends on this module working.

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Mutex;

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://huggingface.co/onnx-community/mms-300m-1130-forced-aligner-ONNX/resolve/main/onnx/model.onnx";
const VOCAB_URL: &str =
    "https://huggingface.co/onnx-community/mms-300m-1130-forced-aligner-ONNX/resolve/main/vocab.json";
/// Exact f32 model size, shown in Settings and sanity-checked after download.
const MODEL_BYTES: u64 = 1_262_529_881;
/// wav2vec2-family conv stride: 320 samples at 16 kHz (MMS keeps it).
const FRAME_MS: i64 = 20;

struct Engine {
    session: LearningModelSession,
    /// char → token id, from the repo's own vocab.json (never hardcoded).
    chars: HashMap<char, usize>,
    vocab_len: usize,
    blank: usize,
    /// Word-separator token — wav2vec2-960h had `|`, MMS has none.
    sep: Option<usize>,
}

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

pub(crate) fn models_dir(app: &tauri::AppHandle) -> Result<PathBuf, String> {
    let dir = crate::paths::app_data(app)
        .map_err(|e| e.to_string())?
        .join("lyric-models");
    std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
    Ok(dir)
}

fn model_file(app: &tauri::AppHandle) -> Result<PathBuf, String> {
    Ok(models_dir(app)?.join("mms-300m-fa.onnx"))
}

fn vocab_file(app: &tauri::AppHandle) -> Result<PathBuf, String> {
    Ok(models_dir(app)?.join("mms-300m-fa.vocab.json"))
}

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

#[tauri::command]
pub fn lyric_model_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,
        // The tier the machine would actually prepare under right now —
        // ctc-v5 once the separator model is ALSO on disk.
        "aligner": crate::lyric_prep::active_aligner(&app),
        "downloadBytes": MODEL_BYTES,
    })
}

/// Download the model + vocab. Runs ONLY from the user's explicit Settings
/// click — the consent-before-download charter rule. Streams to a temp file,
/// renames atomically, emits `lyrics:model-progress` along the way.
#[tauri::command]
pub async fn lyric_model_download(app: tauri::AppHandle) -> Result<(), String> {
    let client = reqwest::Client::builder()
        .user_agent("SoulStream/0.1 (lyric alignment)")
        .build()
        .map_err(|e| e.to_string())?;
    // Vocab first — tiny, and the engine cannot load without it.
    let vocab_dest = vocab_file(&app)?;
    let vocab = client
        .get(VOCAB_URL)
        .send()
        .await
        .map_err(|e| e.to_string())?
        .error_for_status()
        .map_err(|e| e.to_string())?
        .bytes()
        .await
        .map_err(|e| e.to_string())?;
    serde_json::from_slice::<HashMap<String, usize>>(&vocab)
        .map_err(|_| "vocab.json did not parse".to_string())?;
    std::fs::write(&vocab_dest, &vocab).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(0);
    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 > 4 * 1024 * 1024 {
            last_emit = received;
            let _ = app.emit(
                "lyrics:model-progress",
                serde_json::json!({ "received": received, "total": total }),
            );
        }
    }
    drop(file);
    if total > 0 && received != total {
        let _ = std::fs::remove_file(&tmp);
        return Err(format!("download truncated: {received} of {total} bytes"));
    }
    std::fs::rename(&tmp, &dest).map_err(|e| e.to_string())?;
    // The superseded wav2vec2 tier's files are dead weight once this lands.
    if let Ok(dir) = models_dir(&app) {
        for legacy in ["wav2vec2-base-960h.onnx", "wav2vec2-base-960h.vocab.json"] {
            let _ = std::fs::remove_file(dir.join(legacy));
        }
    }
    let _ = app.emit(
        "lyrics: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("model not downloaded".into());
    }
    let raw = std::fs::read(vocab_file(app)?).map_err(|e| e.to_string())?;
    let vocab: HashMap<String, usize> =
        serde_json::from_slice(&raw).map_err(|_| "vocab.json did not parse".to_string())?;
    let vocab_len = vocab.values().copied().max().unwrap_or(0) + 1;
    // CTC blank: MMS names it `<blank>` (its pad_token_id points there);
    // the wav2vec2-960h family used `<pad>` itself. Prefer the explicit one.
    let blank = *vocab
        .get("<blank>")
        .or_else(|| vocab.get("<pad>"))
        .ok_or("vocab has no blank token")?;
    let mut chars = HashMap::new();
    for (token, id) in &vocab {
        let mut it = token.chars();
        if let (Some(c), None) = (it.next(), it.next()) {
            chars.insert(c, *id);
        }
    }
    let sep = chars.get(&'|').copied();
    // Match the hero detector's WinRT posture: 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 alignment model: {e}"))?;
    let device = LearningModelDevice::Create(LearningModelDeviceKind::Cpu)
        .map_err(|e| format!("could not create WinML device: {e}"))?;
    let session = LearningModelSession::CreateFromModelOnDevice(&model, &device)
        .map_err(|e| format!("could not create alignment session: {e}"))?;
    *guard = Some(Engine {
        session,
        chars,
        vocab_len,
        blank,
        sep,
    });
    Ok(())
}

/// Bare-ASCII fold for the Latin letters MMS's romanized vocab expects:
/// Western lyrics keep their diacritics, the model's chars do not. Anything
/// outside this stays unmapped and the line falls back to the estimate.
fn fold_latin(c: char) -> char {
    match c {
        'à'..='å' | 'À'..='Å' => 'a',
        'è'..='ë' | 'È'..='Ë' => 'e',
        'ì'..='ï' | 'Ì'..='Ï' => 'i',
        'ò'..='ö' | 'Ò'..='Ö' | 'ø' | 'Ø' => 'o',
        'ù'..='ü' | 'Ù'..='Ü' => 'u',
        'ý' | 'ÿ' | 'Ý' => 'y',
        'ñ' | 'Ñ' => 'n',
        'ç' | 'Ç' => 'c',
        // Typographic apostrophes down to the vocab's plain one.
        '’' | 'ʼ' | '`' => '\'',
        _ => c,
    }
}

pub struct AlignedWord {
    pub start_ms: i64,
    pub end_ms: i64,
}

pub struct AlignedLine {
    pub words: Vec<AlignedWord>,
    pub confidence: f32,
}

/// Force-align a PASSAGE — one or more consecutive lines' words — against a
/// 16 kHz mono window. Returns per-line word times RELATIVE to the window
/// plus a per-line confidence, or Err when the model cannot carry the
/// passage (the caller retries per line, then falls back to the estimate).
///
/// Aligning lines TOGETHER is what lets word boundaries pull free of a
/// slightly-late source cue: inside a passage the line cues are only hints,
/// not walls — the research bundle's "whole-song alignment using line
/// starts as hints", bounded to passage windows. `floors_ms[l]` is line
/// l's earliest permitted emission (window-relative): without that floor,
/// a repeating chorus lets the path lock whole lines onto the PREVIOUS
/// acoustic instance of the same words (the charter's "competing repeated
/// sections") — hint means "at most this much early", never "anywhere".
/// `ceils_ms[l]` is the floor's mirror, the latest permitted emission:
/// without it a line may smear arbitrarily LATE across its successors'
/// audio when the space between lines still sounds vocal-ish (a separated
/// stem keeps reverb/delay tails), sustaining one token for seconds at
/// HIGH margin confidence — free decoding hears the same tail, so the
/// score cannot object. Hint means "at most this much late", too.
pub fn align_passage(
    app: &tauri::AppHandle,
    pcm: &[f32],
    lines: &[Vec<String>],
    floors_ms: &[i64],
    ceils_ms: &[i64],
) -> Result<Vec<AlignedLine>, String> {
    let word_total: usize = lines.iter().map(|l| l.len()).sum();
    if pcm.len() < 1600
        || word_total == 0
        || floors_ms.len() != lines.len()
        || ceils_ms.len() != lines.len()
    {
        return Err("window too short".into());
    }
    ensure_engine(app)?;
    let guard = ENGINE.lock().unwrap();
    let engine = guard.as_ref().ok_or("engine unavailable")?;

    // Zero-mean / unit-variance, the model's expected normalization.
    let mean = pcm.iter().sum::<f32>() / pcm.len() as f32;
    let var = pcm.iter().map(|s| (s - mean) * (s - mean)).sum::<f32>() / pcm.len() as f32;
    let scale = 1.0 / (var.sqrt() + 1e-7);
    let normed: Vec<f32> = pcm.iter().map(|s| (s - mean) * scale).collect();

    let input = TensorFloat::CreateFromShapeArrayAndDataArray(&[1, normed.len() as i64], &normed)
        .map_err(|e| format!("could not create input tensor: {e}"))?;
    let binding = LearningModelBinding::CreateFromSession(&engine.session)
        .map_err(|e| format!("could not create binding: {e}"))?;
    binding
        .Bind(&HSTRING::from("input_values"), &input)
        .map_err(|e| format!("could not bind input: {e}"))?;
    let result = engine
        .session
        .Evaluate(&binding, &HSTRING::from("lyric-align"))
        .map_err(|e| format!("inference failed: {e}"))?;
    if !result.Succeeded().unwrap_or(false) {
        return Err("inference did not succeed".into());
    }
    let logits: TensorFloat = result
        .Outputs()
        .and_then(|o| o.Lookup(&HSTRING::from("logits")))
        .map_err(|e| format!("logits missing: {e}"))?
        .cast()
        .map_err(|e| format!("logits wrong type: {e}"))?;
    let values: Vec<f32> = logits
        .GetAsVectorView()
        .map_err(|e| format!("could not read logits: {e}"))?
        .into_iter()
        .collect();
    let v = engine.vocab_len;
    if v == 0 || values.len() % v != 0 {
        return Err(format!("unexpected logits length {}", values.len()));
    }
    let t_frames = values.len() / v;

    // log-softmax per frame.
    let mut logp = vec![0.0f32; values.len()];
    for t in 0..t_frames {
        let row = &values[t * v..(t + 1) * v];
        let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let sum: f32 = row.iter().map(|x| (x - max).exp()).sum();
        let log_z = max + sum.ln();
        for (i, x) in row.iter().enumerate() {
            logp[t * v + i] = x - log_z;
        }
    }

    // Token sequence: each token knows its WORD (passage-global index), each
    // word knows its line; a separator joins words only when the vocab has
    // one (wav2vec2-960h did) — MMS aligns bare chars. Mapping is
    // vocab-driven and case-agnostic, with a Latin fold so accented lyrics
    // reach the romanized vocab. A word that maps to nothing (digits,
    // symbols, other scripts) sinks the whole passage to the per-line
    // ladder rather than being silently skipped.
    let map_char = |c: char| -> Option<usize> {
        engine
            .chars
            .get(&c)
            .or_else(|| c.to_lowercase().next().and_then(|l| engine.chars.get(&l)))
            .or_else(|| c.to_uppercase().next().and_then(|u| engine.chars.get(&u)))
            .or_else(|| {
                let f = fold_latin(c);
                engine
                    .chars
                    .get(&f)
                    .or_else(|| f.to_uppercase().next().and_then(|u| engine.chars.get(&u)))
            })
            .copied()
    };
    let mut tokens: Vec<(usize, usize)> = Vec::new();
    // Per-token emission bounds, frames — the line's cue-derived floor and
    // its successor-derived ceiling.
    let mut floor_frame: Vec<usize> = Vec::new();
    let mut ceil_frame: Vec<usize> = Vec::new();
    let mut next_wi = 0usize;
    for (li, line) in lines.iter().enumerate() {
        let floor = (floors_ms[li].max(0) / FRAME_MS) as usize;
        let ceil = (ceils_ms[li].max(0) / FRAME_MS) as usize;
        for word in line {
            let wi = next_wi;
            next_wi += 1;
            if wi > 0 {
                if let Some(sep) = engine.sep {
                    tokens.push((sep, wi));
                    floor_frame.push(floor);
                    ceil_frame.push(ceil);
                }
            }
            let before = tokens.len();
            for c in word.chars() {
                if let Some(id) = map_char(c) {
                    tokens.push((id, wi));
                    floor_frame.push(floor);
                    ceil_frame.push(ceil);
                }
            }
            if tokens.len() == before {
                return Err("word has no mappable characters".into());
            }
        }
    }
    let n = tokens.len();
    if t_frames < n {
        return Err("window shorter than token sequence".into());
    }

    // Real CTC trellis: at each frame stay on blank, SUSTAIN the current
    // token by re-emitting it (the self-loop), or advance by emitting the
    // next. The self-loop turns single emission frames into token SPANS —
    // without it every word "ended" one frame after it started and sustained
    // notes were invisible. Known approximation, kept deliberately: identical
    // CONSECUTIVE tokens (double letters, "LL") properly require a mandatory
    // blank between them; collapsing that costs ~one frame (20ms).
    const NEG: f32 = f32::NEG_INFINITY;
    let mut trellis = vec![NEG; (t_frames + 1) * (n + 1)];
    // Backpointers (0 = blank, 1 = self-loop, 2 = advance) make the
    // backtrack exact instead of re-deriving float comparisons.
    let mut back = vec![0u8; (t_frames + 1) * (n + 1)];
    let idx = |t: usize, j: usize| t * (n + 1) + j;
    trellis[idx(0, 0)] = 0.0;
    for t in 0..t_frames {
        let blank_p = logp[t * v + engine.blank];
        for j in 0..=n.min(t + 1) {
            let mut best = trellis[idx(t, j)] + blank_p;
            let mut way = 0u8;
            // A token may not be emitted before its line's floor (repeated
            // choruses would lock onto an earlier acoustic instance) nor at
            // or past its ceiling (stem reverb tails would smear a sustain
            // across the successors' audio). A line that cannot fit between
            // its bounds leaves no path — the caller's ladder handles it.
            if j > 0 && t >= floor_frame[j - 1] && t < ceil_frame[j - 1] {
                let tok_p = logp[t * v + tokens[j - 1].0];
                let sustain = trellis[idx(t, j)] + tok_p;
                if sustain > best {
                    best = sustain;
                    way = 1;
                }
                let advance = trellis[idx(t, j - 1)] + tok_p;
                if advance > best {
                    best = advance;
                    way = 2;
                }
            }
            trellis[idx(t + 1, j)] = best;
            back[idx(t + 1, j)] = way;
        }
    }
    if trellis[idx(t_frames, n)] == NEG {
        return Err("no alignment path".into());
    }
    // Backtrack the pointers into per-token spans: an advance opens the
    // token's span, self-loops extend it. Walking backwards, the first
    // emission met is the token's LAST frame. Every backtrack step consumes
    // exactly one frame emitting blank or a token — recording that symbol's
    // logp gives the path's per-frame score, which per-LINE confidence
    // needs below.
    let mut first_emit = vec![usize::MAX; n];
    let mut last_emit = vec![usize::MAX; n];
    let mut path_logp = vec![0.0f32; t_frames];
    let mut j = n;
    let mut t = t_frames;
    while t > 0 {
        match back[idx(t, j)] {
            2 => {
                first_emit[j - 1] = t - 1;
                if last_emit[j - 1] == usize::MAX {
                    last_emit[j - 1] = t - 1;
                }
                path_logp[t - 1] = logp[(t - 1) * v + tokens[j - 1].0];
                j -= 1;
            }
            1 => {
                if last_emit[j - 1] == usize::MAX {
                    last_emit[j - 1] = t - 1;
                }
                path_logp[t - 1] = logp[(t - 1) * v + tokens[j - 1].0];
            }
            _ => {
                path_logp[t - 1] = logp[(t - 1) * v + engine.blank];
            }
        }
        t -= 1;
    }
    if j != 0 {
        return Err("backtrack failed".into());
    }

    // Per-word spans, then grouped into lines.
    let mut word_span: Vec<Option<(usize, usize)>> = vec![None; word_total];
    for (k, (id, w)) in tokens.iter().enumerate() {
        if Some(*id) == engine.sep || first_emit[k] == usize::MAX {
            continue;
        }
        word_span[*w] = Some(match word_span[*w] {
            None => (first_emit[k], last_emit[k]),
            Some((f, _)) => (f, last_emit[k]),
        });
    }

    // Confidence = margin to FREE decoding, per LINE: exp of the mean
    // per-frame gap between the forced path and the best unconstrained path
    // over the line's own sung span — "how close is the known lyric to what
    // the model actually hears", the calibration CTC-segmentation systems
    // use. Near 1.0 when text matches audio, sinking as it diverges. (v1
    // sampled exp(logprob) at each token's single emission frame, where CTC
    // piles mass on blank — every line pinned to the clamp floor.)
    let free: Vec<f32> = (0..t_frames)
        .map(|t| logp[t * v..(t + 1) * v].iter().cloned().fold(NEG, f32::max))
        .collect();
    let mut out = Vec::with_capacity(lines.len());
    let mut wi = 0usize;
    for line in lines {
        let mut words_out = Vec::with_capacity(line.len());
        let mut span_first = usize::MAX;
        let mut span_last = 0usize;
        for _ in line {
            match word_span[wi] {
                Some((f, l)) => {
                    span_first = span_first.min(f);
                    span_last = span_last.max(l);
                    words_out.push(AlignedWord {
                        start_ms: f as i64 * FRAME_MS,
                        end_ms: (l as i64 + 1) * FRAME_MS,
                    });
                }
                None => return Err("word lost in backtrack".into()),
            }
            wi += 1;
        }
        let confidence = if span_first == usize::MAX {
            0.05
        } else {
            let gap: f32 = (span_first..=span_last.min(t_frames - 1))
                .map(|t| path_logp[t] - free[t])
                .sum();
            let frames = (span_last.min(t_frames - 1) - span_first + 1) as f32;
            (gap / frames).exp().clamp(0.05, 0.99)
        };
        out.push(AlignedLine {
            words: words_out,
            confidence,
        });
    }
    Ok(out)
}