//! SNAC — Multi-Scale Neural Audio Codec (`wai.neural.snac`).
//!
//! Unlike the flat-RVQ codecs (`audio.rs`: one `[1,1,q,t]` code grid), SNAC is
//! HIERARCHICAL: N codebooks each at a different temporal resolution. A 24/32/44
//! kHz model emits N variable-length integer token streams (e.g. lengths
//! `[12, 24, 48, 96]` — each level 2× the previous rate). `model.decode(codes)`
//! upsamples each level by its stride, sums to `z_q`, and runs the decoder net;
//! we export that whole path as one ONNX taking the N code tensors → waveform.
//!
//! Wire format (little-endian):
//!   <III>  sample_rate, n_samples, n_levels
//!   n_levels × u32   length T_i of each level (coarse → fine)
//!   then, per level i: T_i × u16 codes (row-major)
//!
//! ONNX decoder contract (from the exporter):
//!   inputs:  codes_0 … codes_{N-1}   int64 [1, T_i]
//!   output:  audio_values            float [1, 1, samples]

use std::borrow::Cow;

use ndarray::{Array2, IxDyn};
use ort::session::SessionInputValue;
use ort::value::Value;

use super::audio::DecodedAudio;
use super::{params::ParamSource, DecodeError};

const HEADER: usize = 12;

pub(crate) fn decode(payload: &[u8], params: &ParamSource) -> Result<DecodedAudio, DecodeError> {
    if payload.len() < HEADER {
        return Err(DecodeError::InvalidPayload(format!(
            "snac payload {} B < {HEADER} B header", payload.len())));
    }
    let sample_rate = u32::from_le_bytes(payload[0..4].try_into().unwrap());
    let n_samples = u32::from_le_bytes(payload[4..8].try_into().unwrap()) as usize;
    let n_levels = u32::from_le_bytes(payload[8..12].try_into().unwrap()) as usize;
    if n_levels == 0 || n_levels > 8 {
        return Err(DecodeError::InvalidPayload(format!("snac n_levels={n_levels} (want 1..=8)")));
    }

    let mut off = HEADER;
    let len_table = HEADER + 4 * n_levels;
    if payload.len() < len_table {
        return Err(DecodeError::InvalidPayload(format!(
            "snac payload {} B < {len_table} B (header + {n_levels} lengths)", payload.len())));
    }
    let mut lengths = Vec::with_capacity(n_levels);
    for _ in 0..n_levels {
        lengths.push(u32::from_le_bytes(payload[off..off + 4].try_into().unwrap()) as usize);
        off += 4;
    }
    let codes_bytes: usize = lengths.iter().map(|t| 2 * t).sum();
    if payload.len() != len_table + codes_bytes {
        return Err(DecodeError::InvalidPayload(format!(
            "snac payload {} != {len_table} + codes {codes_bytes}", payload.len())));
    }

    // Parse each level's codes → [1, T_i] i64, build the named ONNX inputs.
    let mut sess = params.session_model("decoder.onnx")?;
    let mut inputs: Vec<(Cow<'static, str>, SessionInputValue<'static>)> =
        Vec::with_capacity(n_levels);
    for (i, &t) in lengths.iter().enumerate() {
        let mut arr = Array2::<i64>::zeros((1, t));
        let slice = arr.as_slice_mut().unwrap();
        for j in 0..t {
            slice[j] = u16::from_le_bytes([payload[off + 2 * j], payload[off + 2 * j + 1]]) as i64;
        }
        off += 2 * t;
        let v = Value::from_array(arr)
            .map_err(|e| DecodeError::Ort(format!("snac codes_{i}: {e}")))?;
        inputs.push((Cow::Owned(format!("codes_{i}")), v.into()));
    }

    let outputs = sess.run(inputs)
        .map_err(|e| DecodeError::Ort(format!("snac session.run: {e}")))?;
    let out = outputs.get("audio_values")
        .ok_or_else(|| DecodeError::Ort("snac output 'audio_values' missing".into()))?;
    let view = out.try_extract_array::<f32>()
        .map_err(|e| DecodeError::Ort(format!("extract audio_values: {e}")))?;
    let view: ndarray::ArrayViewD<f32> = view.view().into_dyn()
        .into_dimensionality::<IxDyn>()
        .map_err(|e| DecodeError::Ort(format!("reshape audio_values: {e}")))?;
    let mut samples: Vec<f32> = view.iter().copied().collect();
    if samples.len() > n_samples {
        samples.truncate(n_samples);
    }
    Ok(DecodedAudio { samples, sample_rate })
}
