//! MJPEG decoding through libjpeg-turbo's TurboJPEG API.
//!
//! The ATEM Mini Pro emits baseline 4:2:2 JPEG. Rather than paying for a CPU
//! YUV->RGB conversion and a 4-bytes-per-pixel upload, we pull the raw planes
//! out of the decoder, interleave them into packed YUY2 (one sequential pass,
//! 2 bytes per pixel) and let the GPU's sampler do the colour conversion.
//!
//! 4:2:0 sources skip even the interleave: their planes go straight into the
//! frame buffer in I420 order. Anything else falls back to BGRX, which is
//! always correct but does the conversion on the CPU.

use crate::frame::{BufferPool, Colorspace, Frame, PixelFormat};
use std::ffi::{c_char, c_int, c_uchar, c_void};
use std::fmt;

// --- TurboJPEG FFI ----------------------------------------------------------
// Values transcribed from turbojpeg.h (libjpeg-turbo 3.x). The API is an
// opaque handle plus plain ints, so there is no struct layout to get wrong.

type TjHandle = *mut c_void;

const TJINIT_DECOMPRESS: c_int = 1;

const TJPARAM_STOPONWARNING: c_int = 0;
const TJPARAM_SUBSAMP: c_int = 4;
const TJPARAM_JPEGWIDTH: c_int = 5;
const TJPARAM_JPEGHEIGHT: c_int = 6;
const TJPARAM_PRECISION: c_int = 7;
const TJPARAM_COLORSPACE: c_int = 8;
const TJPARAM_MAXPIXELS: c_int = 24;

const TJSAMP_444: c_int = 0;
const TJSAMP_422: c_int = 1;
const TJSAMP_420: c_int = 2;
const TJSAMP_GRAY: c_int = 3;

const TJPF_BGRX: c_int = 3;

const TJERR_WARNING: c_int = 0;

const TJCS_RGB: c_int = 0;
const TJCS_YCBCR: c_int = 1;
const TJCS_GRAY: c_int = 2;
const TJCS_CMYK: c_int = 3;
const TJCS_YCCK: c_int = 4;

#[link(name = "turbojpeg")]
extern "C" {
    fn tj3Init(init_type: c_int) -> TjHandle;
    fn tj3Destroy(handle: TjHandle);
    fn tj3GetErrorStr(handle: TjHandle) -> *const c_char;
    fn tj3GetErrorCode(handle: TjHandle) -> c_int;
    fn tj3Set(handle: TjHandle, param: c_int, value: c_int) -> c_int;
    fn tj3Get(handle: TjHandle, param: c_int) -> c_int;
    fn tj3DecompressHeader(handle: TjHandle, jpeg_buf: *const c_uchar, jpeg_size: usize) -> c_int;
    fn tj3Decompress8(
        handle: TjHandle,
        jpeg_buf: *const c_uchar,
        jpeg_size: usize,
        dst_buf: *mut c_uchar,
        pitch: c_int,
        pixel_format: c_int,
    ) -> c_int;
    fn tj3DecompressToYUVPlanes8(
        handle: TjHandle,
        jpeg_buf: *const c_uchar,
        jpeg_size: usize,
        dst_planes: *mut *mut c_uchar,
        strides: *mut c_int,
    ) -> c_int;
}

/// Refuse absurd dimensions before allocating. 8K worth of pixels is far beyond
/// anything this appliance handles and keeps a corrupt header from asking for a
/// multi-gigabyte buffer.
const MAX_PIXELS: c_int = 7680 * 4320;

#[derive(Debug)]
pub enum DecodeError {
    Init(String),
    Header(String),
    Decode(String),
    /// Header parsed but described something we will not render.
    Unsupported(String),
}

impl fmt::Display for DecodeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DecodeError::Init(m) => write!(f, "turbojpeg init failed: {m}"),
            DecodeError::Header(m) => write!(f, "bad JPEG header: {m}"),
            DecodeError::Decode(m) => write!(f, "JPEG decode failed: {m}"),
            DecodeError::Unsupported(m) => write!(f, "unsupported JPEG: {m}"),
        }
    }
}

impl std::error::Error for DecodeError {}

/// Description of the incoming JPEG, read from its header.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct JpegInfo {
    pub width: u32,
    pub height: u32,
    pub subsamp: i32,
    pub colorspace: i32,
}

impl JpegInfo {
    /// Which SDL-uploadable layout this JPEG decodes into most cheaply.
    pub fn target_format(&self) -> PixelFormat {
        match self.subsamp {
            TJSAMP_422 => PixelFormat::Yuy2,
            TJSAMP_420 => PixelFormat::I420,
            // 4:4:4, 4:4:0, 4:1:1, grayscale and anything else: let turbojpeg
            // produce packed BGRX rather than growing a converter per layout.
            _ => PixelFormat::Bgrx,
        }
    }

    pub fn colorspace_name(&self) -> &'static str {
        match self.colorspace {
            TJCS_RGB => "RGB",
            TJCS_YCBCR => "YCbCr",
            TJCS_GRAY => "grayscale",
            TJCS_CMYK => "CMYK",
            TJCS_YCCK => "YCCK",
            _ => "unknown",
        }
    }

    pub fn subsamp_name(&self) -> &'static str {
        match self.subsamp {
            TJSAMP_444 => "4:4:4",
            TJSAMP_422 => "4:2:2",
            TJSAMP_420 => "4:2:0",
            TJSAMP_GRAY => "grayscale",
            4 => "4:4:0",
            5 => "4:1:1",
            6 => "4:4:1",
            _ => "unknown",
        }
    }
}

/// How far back from the end to look for the EOI marker.
///
/// The ATEM leaves 2-4 bytes of padding there; other devices append a handful
/// of trailing zeros. Anything further than this and the frame is not simply
/// padded, it is damaged.
const EOI_SEARCH_WINDOW: usize = 64;

/// Cheap structural check for a whole JPEG: starts with SOI, ends with EOI.
///
/// This is what separates "padded" from "truncated". Entropy-coded data can
/// never contain a bare `FF D9` (an `FF` in the entropy stream is always
/// followed by a stuffed `00` or a restart marker), so finding EOI near the end
/// is a reliable signal that the frame arrived whole.
pub fn looks_complete(jpeg: &[u8]) -> bool {
    if jpeg.len() < 4 || jpeg[0] != 0xFF || jpeg[1] != 0xD8 {
        return false;
    }
    let start = jpeg.len().saturating_sub(EOI_SEARCH_WINDOW).max(2);
    jpeg[start..].windows(2).any(|w| w == [0xFF, 0xD9])
}

pub struct MjpegDecoder {
    handle: TjHandle,
    /// Scratch planar buffer for the 4:2:2 path, reused across frames.
    planar: Vec<u8>,
    /// Frames that decoded fully but tripped a libjpeg warning.
    warnings: u64,
}

// SAFETY: a tjhandle is not shared; the decoder is owned by exactly one thread
// and every call goes through `&mut self`.
unsafe impl Send for MjpegDecoder {}

impl MjpegDecoder {
    pub fn new() -> Result<Self, DecodeError> {
        let mut d = MjpegDecoder { handle: std::ptr::null_mut(), planar: Vec::new(), warnings: 0 };
        d.reinit()?;
        Ok(d)
    }

    /// Builds a fresh decompressor, discarding any existing one.
    ///
    /// libjpeg keeps error and warning state on the handle. A truncated frame
    /// leaves "Premature end of JPEG file" latched, and every subsequent
    /// *valid* frame then fails -- one bad USB frame would kill the feed for
    /// good. Rebuilding the handle after a failure is the only reliable way to
    /// clear that, and it costs nothing on the happy path.
    fn reinit(&mut self) -> Result<(), DecodeError> {
        if !self.handle.is_null() {
            // SAFETY: handle came from tj3Init and is destroyed once.
            unsafe { tj3Destroy(self.handle) };
            self.handle = std::ptr::null_mut();
        }
        // SAFETY: plain C call with a constant argument.
        let handle = unsafe { tj3Init(TJINIT_DECOMPRESS) };
        if handle.is_null() {
            return Err(DecodeError::Init("tj3Init returned NULL".into()));
        }
        self.handle = handle;
        // Warnings are NOT treated as fatal. The ATEM pads each frame with a
        // couple of bytes before the EOI marker, which makes libjpeg warn
        // ("extraneous bytes before marker 0xd9") on an otherwise perfect
        // image -- rejecting those would drop about 97% of frames. Truncated
        // frames are caught structurally by `looks_complete` instead.
        self.set(TJPARAM_STOPONWARNING, 0)?;
        self.set(TJPARAM_MAXPIXELS, MAX_PIXELS)?;
        Ok(())
    }

    /// Clears latched libjpeg state after a failed frame. Best effort: if the
    /// rebuild itself fails the handle is left null and later calls report it
    /// rather than dereferencing it.
    fn recover(&mut self) {
        if let Err(e) = self.reinit() {
            crate::warn_!("could not reinitialise the JPEG decoder: {e}");
        }
    }

    fn set(&self, param: c_int, value: c_int) -> Result<(), DecodeError> {
        // SAFETY: `self.handle` is a live decompressor handle.
        if unsafe { tj3Set(self.handle, param, value) } < 0 {
            return Err(DecodeError::Init(self.last_error()));
        }
        Ok(())
    }

    /// Interprets a TurboJPEG return code.
    ///
    /// TurboJPEG returns -1 for a *warning* as well as for a hard failure
    /// (`if (this->jerr.warning) retval = -1;` in turbojpeg.c), and only
    /// `tj3GetErrorCode` tells them apart. On a warning libjpeg still finished
    /// and the output buffer is complete, so the frame is kept.
    ///
    /// This is the ATEM's everyday case: it pads each frame with a few bytes
    /// before the EOI marker, which warns on every single frame. Treating that
    /// as failure discards the entire stream.
    fn check(&mut self, rc: c_int, what: &str) -> Result<(), DecodeError> {
        if rc >= 0 {
            return Ok(());
        }
        // SAFETY: handle is live whenever a tj3 call has just returned.
        let code = unsafe { tj3GetErrorCode(self.handle) };
        if code == TJERR_WARNING {
            self.warnings = self.warnings.saturating_add(1);
            return Ok(());
        }
        let msg = self.last_error();
        self.recover();
        Err(DecodeError::Decode(format!("{what}: {msg}")))
    }

    /// Number of frames decoded with a recoverable warning since the last call.
    pub fn take_warning_count(&mut self) -> u64 {
        std::mem::take(&mut self.warnings)
    }

    fn last_error(&self) -> String {
        // SAFETY: tj3GetErrorStr always returns a valid NUL-terminated string
        // for a live handle.
        unsafe {
            let p = tj3GetErrorStr(self.handle);
            if p.is_null() {
                "unknown turbojpeg error".to_string()
            } else {
                std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned()
            }
        }
    }

    /// Reads dimensions and subsampling without decoding pixel data.
    pub fn read_header(&mut self, jpeg: &[u8]) -> Result<JpegInfo, DecodeError> {
        if self.handle.is_null() {
            self.reinit()?;
        }
        if jpeg.len() < 4 {
            return Err(DecodeError::Header(format!("frame too short ({} bytes)", jpeg.len())));
        }
        if !looks_complete(jpeg) {
            return Err(DecodeError::Header(format!(
                "incomplete frame ({} bytes, no end-of-image marker)",
                jpeg.len()
            )));
        }
        // SAFETY: pointer and length describe the caller's live slice.
        let rc = unsafe { tj3DecompressHeader(self.handle, jpeg.as_ptr(), jpeg.len()) };
        self.check(rc, "reading the header")
            .map_err(|e| DecodeError::Header(e.to_string()))?;
        // SAFETY: handle is live; these params are always readable post-header.
        let (w, h, subsamp, colorspace, precision) = unsafe {
            (
                tj3Get(self.handle, TJPARAM_JPEGWIDTH),
                tj3Get(self.handle, TJPARAM_JPEGHEIGHT),
                tj3Get(self.handle, TJPARAM_SUBSAMP),
                tj3Get(self.handle, TJPARAM_COLORSPACE),
                tj3Get(self.handle, TJPARAM_PRECISION),
            )
        };
        if w <= 0 || h <= 0 {
            self.recover();
            return Err(DecodeError::Header(format!("nonsensical dimensions {w}x{h}")));
        }
        if precision != 8 {
            self.recover();
            return Err(DecodeError::Unsupported(format!("{precision}-bit samples")));
        }
        Ok(JpegInfo {
            width: w as u32,
            height: h as u32,
            subsamp,
            colorspace,
        })
    }

    /// Decodes one JPEG into a [`Frame`] backed by an allocation from `pool`.
    ///
    /// On failure the borrowed buffer is returned to the pool, so a stream of
    /// corrupt frames neither leaks nor thrashes the allocator.
    pub fn decode(
        &mut self,
        jpeg: &[u8],
        pool: &dyn BufferPool,
        seq: u64,
    ) -> Result<Frame, DecodeError> {
        let info = self.read_header(jpeg)?;
        let format = info.target_format();
        let layout = format.layout(info.width, info.height);
        let mut buf = pool.take_buffer(layout.total);

        let result = match format {
            PixelFormat::Yuy2 => self.decode_422_to_yuy2(jpeg, &info, &mut buf, layout.pitches[0]),
            PixelFormat::I420 => self.decode_420_to_i420(jpeg, &info, &mut buf, &layout),
            PixelFormat::Bgrx => self.decode_to_bgrx(jpeg, &mut buf, layout.pitches[0]),
            PixelFormat::Uyvy => unreachable!("MJPEG never targets UYVY"),
        };

        if let Err(e) = result {
            // `check` has already rebuilt the handle on a fatal error.
            pool.recycle(buf);
            return Err(e);
        }

        Ok(Frame {
            width: info.width,
            height: info.height,
            format,
            // JFIF is BT.601 primaries at full range regardless of subsampling.
            // Grayscale still rides the BGRX path, where the matrix is unused.
            colorspace: Colorspace::Jpeg,
            layout,
            data: buf,
            seq,
        })
    }

    /// 4:2:2: decode to planes, then interleave to Y0 U Y1 V.
    fn decode_422_to_yuy2(
        &mut self,
        jpeg: &[u8],
        info: &JpegInfo,
        out: &mut [u8],
        out_pitch: usize,
    ) -> Result<(), DecodeError> {
        let w = info.width as usize;
        let h = info.height as usize;
        let cw = w.div_ceil(2);

        let y_size = w * h;
        let c_size = cw * h;
        let need = y_size + 2 * c_size;
        if self.planar.len() != need {
            self.planar.resize(need, 0);
        }

        let rc = {
            let (y_plane, rest) = self.planar.split_at_mut(y_size);
            let (u_plane, v_plane) = rest.split_at_mut(c_size);
            let mut planes: [*mut c_uchar; 3] =
                [y_plane.as_mut_ptr(), u_plane.as_mut_ptr(), v_plane.as_mut_ptr()];
            let mut strides: [c_int; 3] = [w as c_int, cw as c_int, cw as c_int];

            // SAFETY: the three plane pointers address disjoint regions of
            // `self.planar`, each large enough for `stride * h` bytes as
            // required for 4:2:2 output.
            unsafe {
                tj3DecompressToYUVPlanes8(
                    self.handle,
                    jpeg.as_ptr(),
                    jpeg.len(),
                    planes.as_mut_ptr(),
                    strides.as_mut_ptr(),
                )
            }
        };
        self.check(rc, "decoding 4:2:2 planes")?;

        let (y_plane, rest) = self.planar.split_at(y_size);
        let (u_plane, v_plane) = rest.split_at(c_size);
        interleave_422_to_yuy2(y_plane, u_plane, v_plane, w, h, cw, out, out_pitch);
        Ok(())
    }

    /// 4:2:0: the frame buffer already has the I420 plane order SDL wants, so
    /// point the decoder straight at it.
    fn decode_420_to_i420(
        &mut self,
        jpeg: &[u8],
        info: &JpegInfo,
        out: &mut [u8],
        layout: &crate::frame::PlaneLayout,
    ) -> Result<(), DecodeError> {
        debug_assert_eq!(out.len(), layout.total);
        let base = out.as_mut_ptr();
        // SAFETY: the three offsets come from `PixelFormat::I420.layout`, which
        // partitions exactly `layout.total` bytes into non-overlapping planes,
        // and `out` is that long.
        let mut planes: [*mut c_uchar; 3] = unsafe {
            [
                base.add(layout.offsets[0]),
                base.add(layout.offsets[1]),
                base.add(layout.offsets[2]),
            ]
        };
        let mut strides: [c_int; 3] = [
            layout.pitches[0] as c_int,
            layout.pitches[1] as c_int,
            layout.pitches[2] as c_int,
        ];
        let _ = info;
        // SAFETY: as above; planes and strides describe a valid I420 buffer.
        let rc = unsafe {
            tj3DecompressToYUVPlanes8(
                self.handle,
                jpeg.as_ptr(),
                jpeg.len(),
                planes.as_mut_ptr(),
                strides.as_mut_ptr(),
            )
        };
        self.check(rc, "decoding 4:2:0 planes")
    }

    fn decode_to_bgrx(&mut self, jpeg: &[u8], out: &mut [u8], pitch: usize) -> Result<(), DecodeError> {
        // SAFETY: `out` is `pitch * height` bytes, matching what turbojpeg
        // writes for TJPF_BGRX at the header's dimensions.
        let rc = unsafe {
            tj3Decompress8(
                self.handle,
                jpeg.as_ptr(),
                jpeg.len(),
                out.as_mut_ptr(),
                pitch as c_int,
                TJPF_BGRX,
            )
        };
        self.check(rc, "decoding to BGRX")
    }
}

impl Drop for MjpegDecoder {
    fn drop(&mut self) {
        if !self.handle.is_null() {
            // SAFETY: handle was produced by tj3Init and is dropped once.
            unsafe { tj3Destroy(self.handle) };
        }
    }
}

/// Interleaves planar 4:2:2 into packed YUY2 (Y0 U Y1 V).
///
/// Split out from the FFI so it can be tested directly.
pub fn interleave_422_to_yuy2(
    y: &[u8],
    u: &[u8],
    v: &[u8],
    width: usize,
    height: usize,
    chroma_width: usize,
    out: &mut [u8],
    out_pitch: usize,
) {
    let pairs = width / 2;
    for row in 0..height {
        let y_row = &y[row * width..row * width + width];
        let u_row = &u[row * chroma_width..row * chroma_width + chroma_width];
        let v_row = &v[row * chroma_width..row * chroma_width + chroma_width];
        let o_row = &mut out[row * out_pitch..row * out_pitch + pairs * 4];

        for (i, quad) in o_row.chunks_exact_mut(4).enumerate() {
            quad[0] = y_row[2 * i];
            quad[1] = u_row[i];
            quad[2] = y_row[2 * i + 1];
            quad[3] = v_row[i];
        }
    }
}

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

    #[test]
    fn interleave_produces_yuy2_byte_order() {
        // 4x2 image, chroma width 2.
        let y = vec![10, 11, 12, 13, /**/ 20, 21, 22, 23];
        let u = vec![100, 101, /**/ 110, 111];
        let v = vec![200, 201, /**/ 210, 211];
        let mut out = vec![0u8; 4 * 2 * 2];

        interleave_422_to_yuy2(&y, &u, &v, 4, 2, 2, &mut out, 8);

        assert_eq!(&out[0..8], &[10, 100, 11, 200, 12, 101, 13, 201]);
        assert_eq!(&out[8..16], &[20, 110, 21, 210, 22, 111, 23, 211]);
    }

    #[test]
    fn target_format_follows_subsampling() {
        let base = JpegInfo { width: 1920, height: 1080, subsamp: TJSAMP_422, colorspace: TJCS_YCBCR };
        assert_eq!(base.target_format(), PixelFormat::Yuy2);
        assert_eq!(JpegInfo { subsamp: TJSAMP_420, ..base }.target_format(), PixelFormat::I420);
        assert_eq!(JpegInfo { subsamp: TJSAMP_444, ..base }.target_format(), PixelFormat::Bgrx);
        assert_eq!(
            JpegInfo { subsamp: TJSAMP_GRAY, colorspace: TJCS_GRAY, ..base }.target_format(),
            PixelFormat::Bgrx
        );
    }

    #[test]
    fn subsamp_and_colorspace_names_are_human_readable() {
        let i = JpegInfo { width: 2, height: 2, subsamp: TJSAMP_422, colorspace: TJCS_YCBCR };
        assert_eq!(i.subsamp_name(), "4:2:2");
        assert_eq!(i.colorspace_name(), "YCbCr");
        assert_eq!(JpegInfo { subsamp: TJSAMP_GRAY, colorspace: TJCS_GRAY, ..i }.colorspace_name(), "grayscale");
    }

    #[test]
    fn looks_complete_accepts_a_padded_frame_and_rejects_a_truncated_one() {
        // Minimal shapes: the predicate only inspects the ends.
        let whole = [0xFFu8, 0xD8, 0x11, 0x22, 0xFF, 0xD9];
        assert!(looks_complete(&whole));

        // The ATEM's quirk: two stray bytes before EOI.
        let padded = [0xFFu8, 0xD8, 0x11, 0x22, 0x00, 0x00, 0xFF, 0xD9];
        assert!(looks_complete(&padded), "padding before EOI is benign");

        // Trailing zeros after EOI, as some devices emit.
        let trailing = [0xFFu8, 0xD8, 0x11, 0xFF, 0xD9, 0, 0, 0, 0];
        assert!(looks_complete(&trailing));

        assert!(!looks_complete(&[0xFF, 0xD8, 0x11, 0x22]), "no EOI at all");
        assert!(!looks_complete(&[0x00, 0x00, 0xFF, 0xD9]), "no SOI");
        assert!(!looks_complete(&[]));
        assert!(!looks_complete(&[0xFF]));
    }

    #[test]
    fn looks_complete_rejects_an_eoi_that_is_too_far_from_the_end() {
        let mut v = vec![0xFFu8, 0xD8, 0xFF, 0xD9];
        v.extend(std::iter::repeat(0u8).take(EOI_SEARCH_WINDOW * 2));
        assert!(!looks_complete(&v), "EOI buried far from the end means damage, not padding");
    }

    #[test]
    fn decoder_initialises() {
        MjpegDecoder::new().expect("turbojpeg should initialise");
    }

    #[test]
    fn empty_input_is_rejected_without_panicking() {
        let mut d = MjpegDecoder::new().unwrap();
        assert!(matches!(d.read_header(&[]), Err(DecodeError::Header(_))));
        assert!(d.decode(&[], &NullPool, 0).is_err());
    }

    #[test]
    fn garbage_input_is_rejected_without_panicking() {
        let mut d = MjpegDecoder::new().unwrap();
        let junk: Vec<u8> = (0..4096u32).map(|i| (i * 7 % 251) as u8).collect();
        assert!(d.decode(&junk, &NullPool, 0).is_err());
    }

    /// A real JPEG header followed by nothing must fail cleanly rather than
    /// reading past the buffer.
    #[test]
    fn truncated_jpeg_is_rejected_without_panicking() {
        let mut d = MjpegDecoder::new().unwrap();
        let mut buf = vec![0xFFu8, 0xD8, 0xFF, 0xDB];
        buf.extend_from_slice(&[0u8; 64]);
        assert!(d.decode(&buf, &NullPool, 0).is_err());
    }
}