//! Android hardware H.264 decode via `MediaCodec` (the NDK `AMediaCodec` API).
//!
//! This drives a byte-buffer-mode `AMediaCodec` H.264 decoder: the host's
//! Annex-B access units go in as input buffers, and the decoder hands back
//! YUV 4:2:0 output buffers (NV12 semi-planar or I420 planar depending on the
//! device) which we convert to the same packed-RGB8 [`Frame`] the openh264 path
//! produces, so the viewer is unchanged. On essentially all Android devices this
//! binds to the SoC's hardware video decoder (OMX / Codec2).
//!
//! The module is only compiled on Android (see the `#[cfg(target_os = "android")]`
//! gate in `decode.rs`). It talks to the stable NDK Media C ABI in `libmediandk`
//! directly — the same ABI the `ndk`/`ndk-sys` crates wrap — so the FFI surface
//! is self-contained and version-independent. If init fails at runtime,
//! [`super::best_decoder`] falls back to the portable openh264 decoder.
//!
//! ## Flow
//! * The codec must be configured with the stream's SPS/PPS as codec-specific
//!   data (`csd-0`/`csd-1`) before it can start, so it is created lazily on the
//!   first access unit carrying both (a keyframe). AUs seen before that are
//!   dropped (`Ok(None)`).
//! * Each `decode` call queues the AU as one input buffer and dequeues at most
//!   one decoded output buffer; the host stream is low-latency with no
//!   B-frames, so the steady state is one frame in, one frame out.
//! * The decoder reports the true picture rectangle through the crop keys on the
//!   output format, so the emitted `Frame` is the real display size (e.g.
//!   1920x1080) even when the YUV planes are stride/slice-height padded.

use super::{Frame, VideoDecoder};
use crate::error::{NdError, Result};

use core::ffi::{c_char, c_void};

use self::sys::*;

/// `MediaCodec` hardware H.264 decoder.
pub struct MediaCodecDecoder {
    /// The codec; `null` until the first keyframe primes and starts it.
    codec: *mut AMediaCodec,
    sps: Vec<u8>,
    pps: Vec<u8>,
    pts_us: i64,
}

impl MediaCodecDecoder {
    /// Probe that an H.264 decoder can be created on this device, then discard
    /// it (the real codec is configured lazily once the SPS/PPS arrive). Returns
    /// an error — so the caller falls back to openh264 — if no AVC decoder
    /// exists at all.
    pub fn new() -> Result<Self> {
        unsafe {
            let probe = AMediaCodec_createDecoderByType(MIME.as_ptr() as *const c_char);
            if probe.is_null() {
                return Err(other("no AMediaCodec AVC decoder available"));
            }
            AMediaCodec_delete(probe);
        }
        Ok(Self {
            codec: core::ptr::null_mut(),
            sps: Vec::new(),
            pps: Vec::new(),
            pts_us: 0,
        })
    }

    /// Create, configure and start the codec from the cached SPS/PPS.
    fn prime(&mut self) -> Result<()> {
        self.teardown();
        unsafe {
            let codec = AMediaCodec_createDecoderByType(MIME.as_ptr() as *const c_char);
            if codec.is_null() {
                return Err(other("AMediaCodec_createDecoderByType returned null"));
            }

            let fmt = AMediaFormat_new();
            if fmt.is_null() {
                AMediaCodec_delete(codec);
                return Err(other("AMediaFormat_new returned null"));
            }
            AMediaFormat_setString(fmt, KEY_MIME.as_ptr() as _, MIME.as_ptr() as _);
            // A nominal size is required; the decoder learns the real one from
            // the codec-specific data and reports it back via the output format.
            AMediaFormat_setInt32(fmt, KEY_WIDTH.as_ptr() as _, 1920);
            AMediaFormat_setInt32(fmt, KEY_HEIGHT.as_ptr() as _, 1080);

            // csd-0 = SPS, csd-1 = PPS, each as a start-code-prefixed NAL.
            let csd0 = with_start_code(&self.sps);
            let csd1 = with_start_code(&self.pps);
            AMediaFormat_setBuffer(
                fmt,
                KEY_CSD0.as_ptr() as _,
                csd0.as_ptr() as *const c_void,
                csd0.len(),
            );
            AMediaFormat_setBuffer(
                fmt,
                KEY_CSD1.as_ptr() as _,
                csd1.as_ptr() as *const c_void,
                csd1.len(),
            );

            let st = AMediaCodec_configure(
                codec,
                fmt,
                core::ptr::null_mut(), // no output Surface: byte-buffer mode
                core::ptr::null_mut(), // no crypto
                0,
            );
            AMediaFormat_delete(fmt);
            if st != AMEDIA_OK {
                AMediaCodec_delete(codec);
                return Err(other(&format!("AMediaCodec_configure failed ({st})")));
            }
            let st = AMediaCodec_start(codec);
            if st != AMEDIA_OK {
                AMediaCodec_delete(codec);
                return Err(other(&format!("AMediaCodec_start failed ({st})")));
            }
            self.codec = codec;
        }
        Ok(())
    }

    /// Queue one Annex-B access unit as an input buffer.
    unsafe fn queue_input(&mut self, au: &[u8]) {
        let idx = AMediaCodec_dequeueInputBuffer(self.codec, DEQUEUE_TIMEOUT_US);
        if idx < 0 {
            return;
        }
        let mut cap: usize = 0;
        let buf = AMediaCodec_getInputBuffer(self.codec, idx as usize, &mut cap);
        if buf.is_null() || au.len() > cap {
            // Still submit an empty buffer so the index is not leaked.
            let _ = AMediaCodec_queueInputBuffer(self.codec, idx as usize, 0, 0, self.pts_us, 0);
            return;
        }
        core::ptr::copy_nonoverlapping(au.as_ptr(), buf, au.len());
        let _ = AMediaCodec_queueInputBuffer(self.codec, idx as usize, 0, au.len(), self.pts_us, 0);
        self.pts_us += PTS_STEP_US;
    }

    /// Dequeue at most one decoded output buffer and convert it to RGB.
    unsafe fn pull_output(&mut self) -> Result<Option<Frame>> {
        let mut info = AMediaCodecBufferInfo {
            offset: 0,
            size: 0,
            presentationTimeUs: 0,
            flags: 0,
        };
        let idx = AMediaCodec_dequeueOutputBuffer(self.codec, &mut info, DEQUEUE_TIMEOUT_US);
        if idx < 0 {
            // INFO_TRY_AGAIN_LATER / OUTPUT_FORMAT_CHANGED / BUFFERS_CHANGED: all
            // mean "no frame this time"; the next call picks the format up.
            return Ok(None);
        }

        let mut cap: usize = 0;
        let buf = AMediaCodec_getOutputBuffer(self.codec, idx as usize, &mut cap);
        let frame = if !buf.is_null() && info.size > 0 {
            let fmt = AMediaCodec_getOutputFormat(self.codec);
            let geom = read_geometry(fmt);
            if !fmt.is_null() {
                AMediaFormat_delete(fmt);
            }
            let end = (info.offset as usize)
                .saturating_add(info.size as usize)
                .min(cap);
            let start = (info.offset as usize).min(end);
            let data = core::slice::from_raw_parts(buf.add(start), end - start);
            yuv420_to_rgb(data, &geom)
        } else {
            None
        };
        let _ = AMediaCodec_releaseOutputBuffer(self.codec, idx as usize, false);
        Ok(frame)
    }

    fn teardown(&mut self) {
        unsafe {
            if !self.codec.is_null() {
                let _ = AMediaCodec_stop(self.codec);
                AMediaCodec_delete(self.codec);
                self.codec = core::ptr::null_mut();
            }
        }
    }
}

impl VideoDecoder for MediaCodecDecoder {
    fn decode(&mut self, au: &[u8]) -> Result<Option<Frame>> {
        if au.is_empty() {
            return Ok(None);
        }

        // Cache parameter sets from keyframes; (re)prime when they change.
        let mut new_sps: Option<Vec<u8>> = None;
        let mut new_pps: Option<Vec<u8>> = None;
        for_each_nal(au, |nal| match nal.first().map(|b| b & 0x1F) {
            Some(7) => new_sps = Some(nal.to_vec()),
            Some(8) => new_pps = Some(nal.to_vec()),
            _ => {}
        });
        if let (Some(sps), Some(pps)) = (new_sps, new_pps) {
            if sps != self.sps || pps != self.pps || self.codec.is_null() {
                self.sps = sps;
                self.pps = pps;
                self.prime()?;
            }
        }
        if self.codec.is_null() {
            return Ok(None); // still waiting for the first keyframe
        }

        unsafe {
            // Feed the whole Annex-B AU (start codes included — MediaCodec wants
            // the elementary stream byte format in byte-buffer mode).
            self.queue_input(au);
            self.pull_output()
        }
    }
}

impl Drop for MediaCodecDecoder {
    fn drop(&mut self) {
        self.teardown();
    }
}

// SAFETY: the decoder is created on, and used only from, the single background
// decode thread that owns it. `AMediaCodec` handles are not shared between
// threads here (no `Sync`); moving ownership to that thread and calling it there
// is sound.
unsafe impl Send for MediaCodecDecoder {}

/// ~29.97 fps in microseconds; only needs to be monotonic.
const PTS_STEP_US: i64 = 33_367;
/// 10 ms dequeue timeout: long enough to pace one frame, short enough to keep
/// the decode loop responsive.
const DEQUEUE_TIMEOUT_US: i64 = 10_000;

// NDK format keys and MIME, as NUL-terminated byte strings (the `AMEDIAFORMAT_*`
// macros expand to exactly these literals).
const MIME: &[u8] = b"video/avc\0";
const KEY_MIME: &[u8] = b"mime\0";
const KEY_WIDTH: &[u8] = b"width\0";
const KEY_HEIGHT: &[u8] = b"height\0";
const KEY_STRIDE: &[u8] = b"stride\0";
const KEY_SLICE_HEIGHT: &[u8] = b"slice-height\0";
const KEY_COLOR_FORMAT: &[u8] = b"color-format\0";
const KEY_CROP_LEFT: &[u8] = b"crop-left\0";
const KEY_CROP_TOP: &[u8] = b"crop-top\0";
const KEY_CROP_RIGHT: &[u8] = b"crop-right\0";
const KEY_CROP_BOTTOM: &[u8] = b"crop-bottom\0";
const KEY_CSD0: &[u8] = b"csd-0\0";
const KEY_CSD1: &[u8] = b"csd-1\0";

// Android `MediaCodecInfo.CodecCapabilities` colour formats we handle.
const COLOR_FORMAT_YUV420_PLANAR: i32 = 19; // I420 (Y, U, V planes)
const COLOR_FORMAT_YUV420_SEMIPLANAR: i32 = 21; // NV12 (Y, interleaved UV)

/// Decoded output-buffer geometry pulled from the codec's output `AMediaFormat`.
struct Geometry {
    color_format: i32,
    stride: usize,
    slice_height: usize,
    /// Visible picture rectangle (display size + top-left offset).
    crop_x: usize,
    crop_y: usize,
    disp_w: usize,
    disp_h: usize,
}

/// Read the stride/slice-height/colour-format and crop rectangle from the output
/// format, filling in sane fallbacks when a key is absent.
unsafe fn read_geometry(fmt: *mut AMediaFormat) -> Geometry {
    let mut width = 0i32;
    let mut height = 0i32;
    let mut stride = 0i32;
    let mut slice_h = 0i32;
    let mut color = COLOR_FORMAT_YUV420_SEMIPLANAR;
    let (mut cl, mut ct, mut cr, mut cb) = (0i32, 0i32, -1i32, -1i32);
    if !fmt.is_null() {
        get_i32(fmt, KEY_WIDTH, &mut width);
        get_i32(fmt, KEY_HEIGHT, &mut height);
        get_i32(fmt, KEY_STRIDE, &mut stride);
        get_i32(fmt, KEY_SLICE_HEIGHT, &mut slice_h);
        get_i32(fmt, KEY_COLOR_FORMAT, &mut color);
        get_i32(fmt, KEY_CROP_LEFT, &mut cl);
        get_i32(fmt, KEY_CROP_TOP, &mut ct);
        get_i32(fmt, KEY_CROP_RIGHT, &mut cr);
        get_i32(fmt, KEY_CROP_BOTTOM, &mut cb);
    }
    if stride <= 0 {
        stride = width;
    }
    if slice_h <= 0 {
        slice_h = height;
    }
    // Crop rect is inclusive; fall back to the full coded size when absent.
    let (disp_w, disp_h, crop_x, crop_y) = if cr >= cl && cb >= ct && cr >= 0 && cb >= 0 {
        (
            (cr - cl + 1) as usize,
            (cb - ct + 1) as usize,
            cl.max(0) as usize,
            ct.max(0) as usize,
        )
    } else {
        (width.max(0) as usize, height.max(0) as usize, 0, 0)
    };
    Geometry {
        color_format: color,
        stride: stride.max(0) as usize,
        slice_height: slice_h.max(0) as usize,
        crop_x,
        crop_y,
        disp_w,
        disp_h,
    }
}

/// Convert an I420 / NV12 output buffer to packed RGB8, cropping to the visible
/// rectangle described by `g` (BT.601 limited range, matching the other paths).
fn yuv420_to_rgb(data: &[u8], g: &Geometry) -> Option<Frame> {
    let (w, h, stride, plane_h) = (g.disp_w, g.disp_h, g.stride, g.slice_height);
    if w == 0 || h == 0 || stride == 0 || plane_h == 0 {
        return None;
    }
    let y_plane = stride * plane_h;
    let semiplanar = g.color_format != COLOR_FORMAT_YUV420_PLANAR;
    let mut rgb = vec![0u8; w * h * 3];
    for row in 0..h {
        let sy = g.crop_y + row;
        let y_off = sy * stride;
        let out_off = row * w * 3;
        let cy = sy / 2;
        for col in 0..w {
            let sx = g.crop_x + col;
            let y = *data.get(y_off + sx).unwrap_or(&16) as i32;
            let (u, v) = if semiplanar {
                // NV12: interleaved UV plane after Y, half-height, full stride.
                // The chroma pair for luma column `sx` sits at the even offset
                // `sx & !1` within the UV row (same stride as the Y plane).
                let uv_off = y_plane + cy * stride;
                let base = uv_off + (sx & !1);
                (
                    *data.get(base).unwrap_or(&128) as i32,
                    *data.get(base + 1).unwrap_or(&128) as i32,
                )
            } else {
                // I420: U plane then V plane, each half-width / half-height.
                let c_stride = stride / 2;
                let c_plane = (stride / 2) * (plane_h / 2);
                let cx = sx / 2;
                let u_off = y_plane + cy * c_stride + cx;
                let v_off = y_plane + c_plane + cy * c_stride + cx;
                (
                    *data.get(u_off).unwrap_or(&128) as i32,
                    *data.get(v_off).unwrap_or(&128) as i32,
                )
            };
            let c = y - 16;
            let d = u - 128;
            let e = v - 128;
            let r = ((298 * c + 409 * e + 128) >> 8).clamp(0, 255) as u8;
            let gg = ((298 * c - 100 * d - 208 * e + 128) >> 8).clamp(0, 255) as u8;
            let b = ((298 * c + 516 * d + 128) >> 8).clamp(0, 255) as u8;
            let o = out_off + col * 3;
            rgb[o] = r;
            rgb[o + 1] = gg;
            rgb[o + 2] = b;
        }
    }
    Some(Frame { w, h, rgb })
}

/// Prepend a 4-byte Annex-B start code to a raw NAL for use as csd.
fn with_start_code(nal: &[u8]) -> Vec<u8> {
    let mut v = Vec::with_capacity(nal.len() + 4);
    v.extend_from_slice(&[0, 0, 0, 1]);
    v.extend_from_slice(nal);
    v
}

/// `AMediaFormat_getInt32` wrapper that leaves `out` untouched when the key is
/// absent (so the caller's fallback stands).
unsafe fn get_i32(fmt: *mut AMediaFormat, key: &[u8], out: *mut i32) -> bool {
    AMediaFormat_getInt32(fmt, key.as_ptr() as *const c_char, out)
}

/// Iterate the Annex-B NAL units in `au`, start codes excluded and trailing zero
/// bytes trimmed.
fn for_each_nal<F: FnMut(&[u8])>(au: &[u8], mut f: F) {
    let n = au.len();
    let mut starts: Vec<usize> = Vec::new();
    let mut i = 0usize;
    while i + 2 < n {
        if au[i] == 0 && au[i + 1] == 0 && au[i + 2] == 1 {
            starts.push(i + 3);
            i += 3;
        } else {
            i += 1;
        }
    }
    for k in 0..starts.len() {
        let begin = starts[k];
        let end = if k + 1 < starts.len() {
            starts[k + 1] - 3
        } else {
            n
        };
        if begin >= end {
            continue;
        }
        let mut nal = &au[begin..end];
        while let [rest @ .., 0] = nal {
            nal = rest;
        }
        if !nal.is_empty() {
            f(nal);
        }
    }
}

fn other(msg: &str) -> NdError {
    NdError::Other(anyhow::anyhow!(msg.to_string()))
}

/// Self-contained FFI to the NDK Media C ABI (`libmediandk`). These signatures
/// match `<media/NdkMediaCodec.h>` / `<media/NdkMediaFormat.h>` and are stable
/// since Android API 21 (`AMediaFormat_getRect`/crop keys since API 21 as well
/// for the common decoders). The `ndk`/`ndk-sys` crates wrap the same library.
#[allow(non_snake_case, non_camel_case_types)]
mod sys {
    use core::ffi::{c_char, c_void};

    pub type AMediaCodec = c_void;
    pub type AMediaFormat = c_void;
    pub type media_status_t = i32;
    pub const AMEDIA_OK: media_status_t = 0;

    #[repr(C)]
    pub struct AMediaCodecBufferInfo {
        pub offset: i32,
        pub size: i32,
        pub presentationTimeUs: i64,
        pub flags: u32,
    }

    #[link(name = "mediandk")]
    extern "C" {
        pub fn AMediaCodec_createDecoderByType(mime_type: *const c_char) -> *mut AMediaCodec;
        pub fn AMediaCodec_configure(
            codec: *mut AMediaCodec,
            format: *mut AMediaFormat,
            surface: *mut c_void,
            crypto: *mut c_void,
            flags: u32,
        ) -> media_status_t;
        pub fn AMediaCodec_start(codec: *mut AMediaCodec) -> media_status_t;
        pub fn AMediaCodec_stop(codec: *mut AMediaCodec) -> media_status_t;
        pub fn AMediaCodec_delete(codec: *mut AMediaCodec) -> media_status_t;
        pub fn AMediaCodec_dequeueInputBuffer(codec: *mut AMediaCodec, timeout_us: i64) -> isize;
        pub fn AMediaCodec_getInputBuffer(
            codec: *mut AMediaCodec,
            idx: usize,
            out_size: *mut usize,
        ) -> *mut u8;
        pub fn AMediaCodec_queueInputBuffer(
            codec: *mut AMediaCodec,
            idx: usize,
            offset: i64,
            size: usize,
            time_us: i64,
            flags: u32,
        ) -> media_status_t;
        pub fn AMediaCodec_dequeueOutputBuffer(
            codec: *mut AMediaCodec,
            info: *mut AMediaCodecBufferInfo,
            timeout_us: i64,
        ) -> isize;
        pub fn AMediaCodec_getOutputBuffer(
            codec: *mut AMediaCodec,
            idx: usize,
            out_size: *mut usize,
        ) -> *mut u8;
        pub fn AMediaCodec_getOutputFormat(codec: *mut AMediaCodec) -> *mut AMediaFormat;
        pub fn AMediaCodec_releaseOutputBuffer(
            codec: *mut AMediaCodec,
            idx: usize,
            render: bool,
        ) -> media_status_t;

        pub fn AMediaFormat_new() -> *mut AMediaFormat;
        pub fn AMediaFormat_delete(fmt: *mut AMediaFormat) -> media_status_t;
        pub fn AMediaFormat_setString(
            fmt: *mut AMediaFormat,
            name: *const c_char,
            value: *const c_char,
        );
        pub fn AMediaFormat_setInt32(fmt: *mut AMediaFormat, name: *const c_char, value: i32);
        pub fn AMediaFormat_setBuffer(
            fmt: *mut AMediaFormat,
            name: *const c_char,
            data: *const c_void,
            size: usize,
        );
        pub fn AMediaFormat_getInt32(
            fmt: *mut AMediaFormat,
            name: *const c_char,
            out: *mut i32,
        ) -> bool;
    }
}