//! C ABI for the native neural sink (features `neural` and `ffi`).
//!
//! One-model entry points, one per media class, for a `WAI1` envelope:
//!
//!   wai_neural_decode_audio(envelope, len, model_path, *samples, *sr)
//!   wai_neural_decode_image(envelope, len, model_path, *rgb, *w, *h)
//!   wai_neural_decode_video(envelope, len, model_path,
//!                            *frames_concat, *n_frames, *w, *h, *fps_x_1000)
//!
//! and one entry point for a `WAI1` or `WAI2` envelope against a registry of the
//! parameter sets the sink holds:
//!
//!   wai_neural_decode_ex(envelope, len, registry_json, flags, *WaiDecoded)
//!
//! Every pin is verified before anything is decoded (SPEC §3.1), and a pinned
//! decode reads only the bytes it verified. Errors carry a status code and a
//! message in `wai_last_error`: a refused pin is `WAI_ERR_PRIOR` with the reason
//! first (`prior_not_held: …`), an unsound `replicate` under the strict flag is
//! `WAI_ERR_UNSOUND_REPLICATE`, and a `WAI2` envelope with nothing to pick is
//! `WAI_ERR_NO_RENDITION`. Output buffers are heap-allocated here; free them with
//! `wai_buffer_free`.
//!
//! Integer video, PCM and splat attributes (`neural::decodes_to_int`) have no
//! `wai_decoded_t` kind; the Rust API returns them (`int_payload::decode_int`). These
//! entry points do not advertise those capabilities, so such a payload is refused
//! before anything is decoded, and dispatch moves on to a fallback this ABI can
//! return.

use std::ffi::{c_char, c_int, CStr};
use std::panic;
use std::path::Path;
use std::slice;

use crate::container::{Wai, WaiMulti};
use crate::ffi::{set_err, WaiBuffer, WAI_ERR_CODEC, WAI_ERR_CONTAINER, WAI_ERR_INVALID_ARG,
                 WAI_ERR_NO_RENDITION, WAI_ERR_PANIC, WAI_ERR_PRIOR, WAI_ERR_UNSOUND_REPLICATE, WAI_OK};
use crate::neural::{decode_envelope, decode_envelope_strict, decode_multi, decode_multi_strict, decodes_to_int, DecodePath,
                    DecodeError, Decoded, ModelRegistry};

unsafe fn input_envelope<'a>(bytes: *const u8, len: usize) -> Option<&'a [u8]> {
    if bytes.is_null() && len > 0 { return None; }
    if bytes.is_null() { return Some(&[]); }
    Some(unsafe { slice::from_raw_parts(bytes, len) })
}

unsafe fn input_cstr<'a>(s: *const c_char) -> Option<&'a str> {
    if s.is_null() { return None; }
    unsafe { CStr::from_ptr(s) }.to_str().ok()
}

/// The status code and message for a decode error.
fn decode_status(e: DecodeError) -> (c_int, String) {
    match e {
        DecodeError::PriorMismatch(r) => (WAI_ERR_PRIOR, format!("{}: {r}", r.code())),
        DecodeError::UnsoundReplicate(why) => (WAI_ERR_UNSOUND_REPLICATE, why),
        DecodeError::NoRendition(_) => (WAI_ERR_NO_RENDITION, e.to_string()),
        DecodeError::Encrypted(why) => (WAI_ERR_CONTAINER, why),
        other => (WAI_ERR_CODEC, other.to_string()),
    }
}

fn run<F: FnOnce() -> Result<(), (c_int, String)> + panic::UnwindSafe>(f: F) -> c_int {
    match panic::catch_unwind(f) {
        Ok(Ok(())) => WAI_OK,
        Ok(Err((code, msg))) => {
            set_err(msg);
            code
        }
        Err(_) => {
            set_err("internal panic in the neural sink");
            WAI_ERR_PANIC
        }
    }
}

/// The refusal of a capability whose output this ABI has no kind for.
fn no_kind(caps: &[&str]) -> (c_int, String) {
    (WAI_ERR_CODEC, format!(
        "{} decodes to integer video, PCM or splat attributes, which have no wai_decoded_t kind: \
         not decoded here; decode it with the Rust API (int_payload::decode_int)", caps.join(", ")))
}

/// Parse a `WAI1` envelope and decode it with `model_path` registered for its
/// capability. A capability whose output has no `wai_decoded_t` kind is refused before
/// anything is decoded.
fn decode_one(envelope: *const u8, envelope_len: usize, model_path: *const c_char) -> Result<Decoded, (c_int, String)> {
    let bytes = unsafe { input_envelope(envelope, envelope_len) }
        .ok_or((WAI_ERR_INVALID_ARG, "envelope null with len>0".to_string()))?;
    let path = unsafe { input_cstr(model_path) }
        .ok_or((WAI_ERR_INVALID_ARG, "model_path null or not UTF-8".to_string()))?;
    let env = Wai::from_bytes(bytes).map_err(|e| (WAI_ERR_CONTAINER, format!("envelope parse: {e}")))?;
    if decodes_to_int(env.capability()) {
        return Err(no_kind(&[env.capability()]));
    }
    let reg = ModelRegistry::new().register(env.capability().to_string(), Path::new(path));
    decode_envelope(&env, &reg).map(|d| d.output).map_err(decode_status)
}

fn wrong_media(want: &str) -> (c_int, String) {
    (WAI_ERR_CONTAINER, format!("the envelope's capability does not decode to {want}"))
}

fn f32_bytes(samples: &[f32]) -> Vec<u8> {
    samples.iter().flat_map(|s| s.to_ne_bytes()).collect()
}

/// Decode a `wai.neural.<audio>` envelope (encodec32 / dac / mimi /
/// wavtokenizer / snac). Writes mono `f32` samples into `out_samples`; the
/// caller reads them as `float*` of length `out_samples.len / 4`.
///
/// # Safety
/// `envelope` valid for `envelope_len` bytes. `model_path` NUL-terminated
/// UTF-8 path to the decoder. `out_samples` writable `WaiBuffer*`.
/// `out_sample_rate` writable `uint32_t*` (or null to skip).
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wai_neural_decode_audio(
    envelope: *const u8, envelope_len: usize, model_path: *const c_char,
    out_samples: *mut WaiBuffer, out_sample_rate: *mut u32,
) -> c_int {
    if out_samples.is_null() { return WAI_ERR_INVALID_ARG; }
    run(|| {
        let Decoded::Audio(audio) = decode_one(envelope, envelope_len, model_path)? else {
            return Err(wrong_media("audio"));
        };
        unsafe { *out_samples = WaiBuffer::from_vec(f32_bytes(&audio.samples)); }
        if !out_sample_rate.is_null() { unsafe { *out_sample_rate = audio.sample_rate; } }
        Ok(())
    })
}

/// Decode a neural image envelope. Writes packed RGB (row-major, 3 bytes/pixel)
/// into `out_rgb`; writes width/height into `*out_width` / `*out_height`.
///
/// # Safety
/// See `wai_neural_decode_audio`.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wai_neural_decode_image(
    envelope: *const u8, envelope_len: usize, model_path: *const c_char,
    out_rgb: *mut WaiBuffer, out_width: *mut u32, out_height: *mut u32,
) -> c_int {
    if out_rgb.is_null() { return WAI_ERR_INVALID_ARG; }
    run(|| {
        let Decoded::Image(img) = decode_one(envelope, envelope_len, model_path)? else {
            return Err(wrong_media("an image"));
        };
        unsafe {
            if !out_width.is_null()  { *out_width  = img.width;  }
            if !out_height.is_null() { *out_height = img.height; }
            *out_rgb = WaiBuffer::from_vec(img.rgb);
        }
        Ok(())
    })
}

/// Decode a neural video envelope. Writes all frames concatenated (frame 0
/// RGB, then frame 1 RGB, …) into `out_frames_concat`; the caller splits at
/// `width*height*3` strides.
///
/// # Safety
/// See `wai_neural_decode_audio`.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wai_neural_decode_video(
    envelope: *const u8, envelope_len: usize, model_path: *const c_char,
    out_frames_concat: *mut WaiBuffer,
    out_n_frames: *mut u32, out_width: *mut u32, out_height: *mut u32,
    out_fps_x_1000: *mut u32,
) -> c_int {
    if out_frames_concat.is_null() { return WAI_ERR_INVALID_ARG; }
    run(|| {
        let Decoded::Video(v) = decode_one(envelope, envelope_len, model_path)? else {
            return Err(wrong_media("video"));
        };
        let concat: Vec<u8> = v.frames_rgb.concat();
        unsafe {
            if !out_n_frames.is_null()    { *out_n_frames    = v.frames_rgb.len() as u32; }
            if !out_width.is_null()       { *out_width       = v.width;  }
            if !out_height.is_null()      { *out_height      = v.height; }
            if !out_fps_x_1000.is_null()  { *out_fps_x_1000  = (v.fps * 1000.0).round() as u32; }
            *out_frames_concat = WaiBuffer::from_vec(concat);
        }
        Ok(())
    })
}

/// The media kind of a [`WaiDecoded`].
pub const WAI_DECODED_AUDIO: u32 = 1;
/// The media kind of a [`WaiDecoded`].
pub const WAI_DECODED_IMAGE: u32 = 2;
/// The media kind of a [`WaiDecoded`].
pub const WAI_DECODED_VIDEO: u32 = 3;

/// `wai_neural_decode_ex` flag: apply SPEC §7's replicate rule (strict).
pub const WAI_DECODE_STRICT: u32 = 1;

/// The decode path of a [`WaiDecoded`]: the capability has one decode path
/// (`neural::DecodePath::Sole`). A session claim's decoder entry states `variant` `null`.
pub const WAI_DECODE_PATH_SOLE: u32 = 1;
/// The decode path of a [`WaiDecoded`]: `wai.image.jpegai`'s exact decoder
/// (`neural::DecodePath::Exact`). A session claim's decoder entry states `variant` `exact`.
pub const WAI_DECODE_PATH_EXACT: u32 = 2;
/// The decode path of a [`WaiDecoded`]: `wai.image.jpegai`'s float decode
/// (`neural::DecodePath::Float`). A session claim's decoder entry states `variant` `float`.
pub const WAI_DECODE_PATH_FLOAT: u32 = 3;

/// What `wai_neural_decode_ex` decoded.
///
/// `kind` is `WAI_DECODED_AUDIO` (`data` holds `f32` samples, `sample_rate` set),
/// `WAI_DECODED_IMAGE` (`data` holds RGB8, `width` and `height` set) or
/// `WAI_DECODED_VIDEO` (`data` holds the frames' RGB8 concatenated, `width`,
/// `height`, `n_frames` and `fps_x_1000` set). `index` is the `WAI2` rendition
/// picked (0 for a `WAI1` envelope). `path` is the decode path that ran
/// (`WAI_DECODE_PATH_SOLE`, `WAI_DECODE_PATH_EXACT` or `WAI_DECODE_PATH_FLOAT`,
/// extensions/determinism-tiers.md §5), which a receipt of the decode states as its
/// decoder entry's `variant` (extensions/jwp-receipts.md §2.4.12). `fallback` is 1 when
/// a `WAI1` envelope decoded through its manifest's `fallback` capability and 0 when
/// through its `capability`; for a `WAI2` envelope `index` names the rendition, and with
/// it the capability, and `fallback` is 0. Free `data` with `wai_buffer_free`.
/// `wai_neural_decode_ex` writes the whole struct on every call that passes a
/// non-null `out`: on failure `kind`, `path` and `fallback` are 0 and `data` is empty,
/// so `data` is always safe to free.
#[repr(C)]
pub struct WaiDecoded {
    pub kind: u32,
    pub index: u32,
    pub data: WaiBuffer,
    pub width: u32,
    pub height: u32,
    pub sample_rate: u32,
    pub n_frames: u32,
    pub fps_x_1000: u32,
    pub path: u32,
    pub fallback: u32,
}

impl WaiDecoded {
    /// Nothing decoded: `kind` 0, `data` empty.
    fn empty() -> Self {
        WaiDecoded {
            kind: 0, index: 0, data: WaiBuffer::empty(),
            width: 0, height: 0, sample_rate: 0, n_frames: 0, fps_x_1000: 0, path: 0, fallback: 0,
        }
    }
}

/// Decode a `WAI1` or `WAI2` envelope against the parameter sets in
/// `registry_json`: a NUL-terminated JSON object mapping each capability to an
/// array of paths. The first path is the registered artifact; the others are
/// further sets a pin may resolve against (a set and its successor). Every pin is
/// verified first (SPEC §3.1); a `WAI2` rendition is picked by SPEC §7.1. With
/// `flags & WAI_DECODE_STRICT`, SPEC §7's replicate rule applies too. A capability
/// whose output has no `wai_decoded_t` kind (integer video, PCM, splat attributes) is
/// not advertised, whatever `registry_json` says: dispatch moves on past it before
/// decoding, and when nothing else decodes the error names it.
///
/// `*out` is written on every call with a non-null `out`, before anything else is
/// read: on failure it holds `kind` 0 and an empty `data`, so the caller may free
/// `out->data` whatever the status.
///
/// # Safety
/// `envelope` valid for `envelope_len` bytes; `registry_json` NUL-terminated;
/// `out` a writable `WaiDecoded*`, which need not be initialised.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn wai_neural_decode_ex(
    envelope: *const u8, envelope_len: usize,
    registry_json: *const c_char, flags: u32,
    out: *mut WaiDecoded,
) -> c_int {
    if out.is_null() {
        set_err("decode_ex: out null");
        return WAI_ERR_INVALID_ARG;
    }
    // `write`, not assignment: the caller's struct may be uninitialised, and its
    // old contents are not ours to drop.
    unsafe { out.write(WaiDecoded::empty()) };
    run(|| {
        let bytes = unsafe { input_envelope(envelope, envelope_len) }
            .ok_or((WAI_ERR_INVALID_ARG, "envelope null with len>0".to_string()))?;
        let rj = unsafe { input_cstr(registry_json) }
            .ok_or((WAI_ERR_INVALID_ARG, "registry_json null or not UTF-8".to_string()))?;
        let map: std::collections::BTreeMap<String, Vec<String>> = serde_json::from_str(rj)
            .map_err(|e| (WAI_ERR_INVALID_ARG, format!("registry_json: {e}")))?;
        let mut reg = ModelRegistry::new();
        let mut kindless: Vec<String> = Vec::new();
        for (cap, paths) in map {
            if decodes_to_int(&cap) {
                kindless.push(cap);
                continue;
            }
            let mut it = paths.into_iter();
            if let Some(first) = it.next() {
                reg = reg.register(cap.clone(), first);
            }
            for p in it {
                reg = reg.hold(cap.clone(), p);
            }
        }
        let strict = flags & WAI_DECODE_STRICT != 0;
        let version = crate::container::detect_version(bytes)
            .map_err(|e| (WAI_ERR_CONTAINER, e.to_string()))?;
        // A refusal at an envelope that names a capability left unadvertised above says so.
        let explain = |named: Vec<&str>, e: DecodeError| {
            let left: Vec<&str> = named.into_iter().filter(|c| kindless.iter().any(|k| k == c)).collect();
            let (code, msg) = decode_status(e);
            if left.is_empty() { (code, msg) } else { (code, format!("{msg}; {}", no_kind(&left).1)) }
        };
        let (index, fallback, d) = if version == 1 {
            let env = Wai::from_bytes(bytes).map_err(|e| (WAI_ERR_CONTAINER, format!("envelope parse: {e}")))?;
            let d = if strict { decode_envelope_strict(&env, &reg) } else { decode_envelope(&env, &reg) };
            let named = [Some(env.capability()), env.fallback()].into_iter().flatten().collect();
            let d = d.map_err(|e| explain(named, e))?;
            (0, d.capability != env.capability(), d)
        } else {
            let m = WaiMulti::from_bytes(bytes).map_err(|e| (WAI_ERR_CONTAINER, format!("envelope parse: {e}")))?;
            let r = if strict { decode_multi_strict(&m, &reg) } else { decode_multi(&m, &reg) };
            let named = m.manifest.renditions.iter().map(|r| r.capability.as_str()).collect();
            let (i, d) = r.map_err(|e| explain(named, e))?;
            (i, false, d)
        };
        let path = match d.path {
            DecodePath::Exact => WAI_DECODE_PATH_EXACT,
            DecodePath::Float => WAI_DECODE_PATH_FLOAT,
            DecodePath::Sole => WAI_DECODE_PATH_SOLE,
        };
        let mut o = WaiDecoded { index: index as u32, path, fallback: u32::from(fallback), ..WaiDecoded::empty() };
        match d.output {
            Decoded::Audio(a) => {
                o.kind = WAI_DECODED_AUDIO;
                o.sample_rate = a.sample_rate;
                o.data = WaiBuffer::from_vec(f32_bytes(&a.samples));
            }
            Decoded::Image(i) => {
                o.kind = WAI_DECODED_IMAGE;
                o.width = i.width;
                o.height = i.height;
                o.data = WaiBuffer::from_vec(i.rgb);
            }
            Decoded::Video(v) => {
                o.kind = WAI_DECODED_VIDEO;
                o.width = v.width;
                o.height = v.height;
                o.n_frames = v.frames_rgb.len() as u32;
                o.fps_x_1000 = (v.fps * 1000.0).round() as u32;
                o.data = WaiBuffer::from_vec(v.frames_rgb.concat());
            }
            // Not reached: such a capability is not advertised above. Refused rather
            // than converted, so no caller could read one as another kind.
            Decoded::Int(_) => {
                return Err((WAI_ERR_CODEC, "this integer output (video, PCM or splat attributes) has no wai_decoded_t kind; decode it with the Rust API".to_string()));
            }
        }
        unsafe { out.write(o) };
        Ok(())
    })
}
