//! JPEG AI (ISO/IEC 6048-1) picture + tool header parsing.
//!
//! Headers are plain MSB-first fixed-width bit fields — no entropy coding.
//! Field widths are either explicit in the syntax or `ceil(log2(max+1))`.
//! The PIH and TON payloads tile contiguously (392 + 504 bits on the
//! conformance stream, no padding between fields).
//!
//! This parses the syntax subset exercised by the `decoder_profile_id = 2`
//! single-region 4:4:4 stream the conformance fixtures use; gated optional
//! blocks that are disabled there (region partitioning, masks, UDI/CICP
//! payloads) are read as their enable flags only and rejected if set, so a
//! stream needing them fails loudly instead of mis-parsing.
//!
//! IMPORTANT: PIH and TON are SEPARATE bit cursors (the reference keeps a
//! per-substream reader): the picture header carries profile/dims/model
//! blocks (without lsbs flags); the tool header starts with the two
//! lsbs_enable flags and then the filter blocks.

/// MSB-first bit reader over a header payload.
pub struct HeaderReader<'a> {
    buf: &'a [u8],
    pos: usize,
}

impl<'a> HeaderReader<'a> {
    pub fn new(buf: &'a [u8]) -> Self {
        Self { buf, pos: 0 }
    }

    /// Read `n` bits MSB-first; `None` past end.
    pub fn bits(&mut self, n: u32) -> Option<u64> {
        let mut v = 0u64;
        for _ in 0..n {
            let byte = *self.buf.get(self.pos >> 3)?;
            v = (v << 1) | ((byte >> (7 - (self.pos & 7))) & 1) as u64;
            self.pos += 1;
        }
        Some(v)
    }

    /// Width for a `max_symbol_value`-coded field.
    pub fn maxsym_bits(max: u64) -> u32 {
        64 - (max).leading_zeros()
    }

    pub fn ue_max(&mut self, max: u64) -> Option<u64> {
        self.bits(Self::maxsym_bits(max))
    }
}

/// Per-colour-model (Y / UV) header block.
#[derive(Debug, Default, Clone)]
pub struct ModelBlock {
    pub num_chs: usize,
    /// Signalled in the TOOL header (first two bits of TON), not in PIH.
    pub lsbs_enable: bool,
    pub use_cube: bool,
    pub rvs_enable: bool,
    pub grfs_enable: bool,
    /// Per-channel gating bitmap for the RVS/CWG table choice.
    pub grfs_channel_flag: Vec<u8>,
    pub synthesis_tile_enable: bool,
}

/// EFElinear signalled parameters (one of the two filter blocks).
#[derive(Debug, Default, Clone)]
pub struct EfeLinearHeader {
    pub mean_u_q100: u64,
    pub mean_v_q100: u64,
    pub best_cand_idx: [u64; 2],
    pub f_l: [u64; 2],
    pub min_symbol: u64,
    pub max_symbol: u64,
    pub weights_u: Vec<u64>,
    pub weights_u2: Vec<u64>,
    pub weights_v: Vec<u64>,
    pub weights_v2: Vec<u64>,
}

/// Everything the decode route needs from PIH + TON.
#[derive(Debug, Default, Clone)]
pub struct JpegAiHeader {
    pub stream_profile_idc: u64,
    pub decoder_profile_id: u64,
    pub synthesis_transform_ids: Vec<u64>,
    pub level_idc: u64,
    pub img_width: u64,
    pub img_height: u64,
    pub bit_depth_idc: u64,
    pub s_ver: u64,
    pub s_hor: u64,
    pub c_ver: u64,
    pub c_hor: u64,
    pub colour_transform_idx: u64,
    pub model_id: u64,
    /// Luma `beta_displacement_log` (de-biased by 2048).
    pub beta_displacement_log: i64,
    /// Chroma `beta_displacement_log`: its own 12-bit field when
    /// `independent_beta_uv` is set, otherwise the luma value.
    pub beta_displacement_log_uv: i64,
    pub models: Vec<ModelBlock>,
    pub efe_linear_enabled: bool,
    pub efe_linear: EfeLinearHeader,
    pub efe_linear2: EfeLinearHeader,
    pub icci_enabled: bool,
    pub icci_use: [bool; 3],
    pub icci_short_list: bool,
    pub icci_signalled_idx: [u64; 2],
    pub efe_nonlinear_enabled: bool,
    pub efn_u_enabled: bool,
    pub efn_v_enabled: bool,
    pub efn_min_luma_q100: u64,
    pub efn_max_luma_q100: u64,
    pub efn_weights: Vec<u64>,
    pub lef_enabled: bool,
    pub lef_ch_idx: u64,
}

#[derive(Debug)]
pub enum HeaderError {
    Truncated(&'static str),
    Unsupported(&'static str),
}

fn efe_linear_block(
    r: &mut HeaderReader,
    first: bool,
    out: &mut EfeLinearHeader,
) -> Result<(), HeaderError> {
    let t = |o: Option<u64>| o.ok_or(HeaderError::Truncated("EFElinear"));
    if first {
        out.mean_u_q100 = t(r.ue_max(32767))?;
        out.mean_v_q100 = t(r.ue_max(32767))?;
    } else {
        out.mean_u_q100 = u64::MAX; // second block reuses means; fields not present
        out.mean_v_q100 = u64::MAX;
    }
    out.best_cand_idx = [t(r.ue_max(15))?, t(r.ue_max(15))?];
    out.f_l = [t(r.ue_max(9))?, t(r.ue_max(9))?];
    out.min_symbol = t(r.ue_max(65535))?;
    out.max_symbol = t(r.ue_max(65535))?;
    // weight counts: fL taps² per candidate region; this profile signals one
    // weight per (U, U2, V, V2) at fL=1. Larger fL ⇒ fL² weights each.
    let n_u = (out.f_l[0] * out.f_l[0]) as usize;
    let n_v = (out.f_l[1] * out.f_l[1]) as usize;
    let range = out.max_symbol.max(1);
    let mut grab = |n: usize| -> Result<Vec<u64>, HeaderError> {
        (0..n).map(|_| t(r.ue_max(range))).collect()
    };
    out.weights_u = grab(n_u)?;
    out.weights_u2 = grab(1)?;
    out.weights_v = grab(n_v)?;
    out.weights_v2 = grab(1)?;
    Ok(())
}

/// Parse the PIH and TON payloads. Each substream has its OWN bit cursor.
pub fn parse_headers(pih: &[u8], ton: &[u8]) -> Result<JpegAiHeader, HeaderError> {
    let r = &mut HeaderReader::new(pih);
    let t = |o: Option<u64>| o.ok_or(HeaderError::Truncated("header"));

    let mut h = JpegAiHeader {
        stream_profile_idc: t(r.bits(4))?,
        decoder_profile_id: t(r.bits(4))?,
        ..Default::default()
    };
    let n_syn = t(r.bits(4))? + 1;
    for _ in 0..n_syn {
        h.synthesis_transform_ids.push(t(r.bits(4))?);
    }
    h.level_idc = t(r.bits(8))?;
    h.img_width = t(r.ue_max(65535))? + 64;
    h.img_height = t(r.ue_max(65535))? + 64;
    let _diff_disp_w = t(r.bits(6))?;
    let _diff_disp_h = t(r.bits(6))?;
    h.bit_depth_idc = t(r.ue_max(4))?;
    h.s_ver = t(r.ue_max(1))? + 1;
    h.s_hor = t(r.ue_max(1))? + 1;
    // c_ver / c_hor are coded only when s_ver / s_hor is 1; otherwise they are 2.
    h.c_ver = if h.s_ver == 1 { t(r.ue_max(1))? + 1 } else { 2 };
    h.c_hor = if h.s_hor == 1 { t(r.ue_max(1))? + 1 } else { 2 };
    h.colour_transform_idx = t(r.ue_max(2))?;
    if h.colour_transform_idx == 2 {
        return Err(HeaderError::Unsupported("signalled colour matrix"));
    }
    h.model_id = t(r.ue_max(15))?;
    let multi_threading_z = t(r.bits(1))?;
    if multi_threading_z != 0 {
        return Err(HeaderError::Unsupported("multi-threaded z"));
    }
    h.beta_displacement_log = t(r.bits(12))? as i64 - 2048;
    let region_partitioning = t(r.ue_max(1))?;
    if region_partitioning != 0 {
        return Err(HeaderError::Unsupported("region partitioning"));
    }
    // quantizer β O-header: a separate chroma β follows when signalled.
    h.beta_displacement_log_uv = if t(r.ue_max(1))? == 1 {
        t(r.bits(12))? as i64 - 2048
    } else {
        h.beta_displacement_log
    };

    for _ in 0..2 {
        let mt = t(r.bits(1))?;
        if mt != 0 {
            return Err(HeaderError::Unsupported("multi-threaded residual"));
        }
        let mut m = ModelBlock {
            num_chs: t(r.bits(8))? as usize,
            use_cube: t(r.ue_max(1))? == 1,
            rvs_enable: t(r.ue_max(1))? == 1,
            grfs_enable: t(r.ue_max(1))? == 1,
            ..Default::default()
        };
        if m.grfs_enable {
            m.grfs_channel_flag =
                (0..m.num_chs).map(|_| r.bits(1).map(|v| v as u8)).collect::<Option<_>>()
                    .ok_or(HeaderError::Truncated("grfs bitmap"))?;
        }
        m.synthesis_tile_enable = t(r.ue_max(1))? == 1;
        h.models.push(m);
    }
    let _gain_3d = t(r.ue_max(1))?;

    // ---- tool header (fresh bit cursor over TON) ----
    let r = &mut HeaderReader::new(ton);
    h.models[0].lsbs_enable = t(r.ue_max(1))? == 1;
    h.models[1].lsbs_enable = t(r.ue_max(1))? == 1;
    h.efe_linear_enabled = t(r.ue_max(1))? == 1;
    if h.efe_linear_enabled {
        efe_linear_block(r, true, &mut h.efe_linear)?;
        efe_linear_block(r, false, &mut h.efe_linear2)?;
    }
    h.icci_enabled = t(r.ue_max(1))? == 1;
    if h.icci_enabled {
        let tiles_enabled = t(r.ue_max(1))?;
        if tiles_enabled != 0 {
            return Err(HeaderError::Unsupported("icci tiling"));
        }
        for i in 0..3 {
            h.icci_use[i] = t(r.ue_max(1))? == 1;
        }
        h.icci_short_list = t(r.ue_max(1))? == 1;
        h.icci_signalled_idx = [t(r.ue_max(2))?, t(r.ue_max(2))?];
    }
    h.efe_nonlinear_enabled = t(r.ue_max(1))? == 1;
    if h.efe_nonlinear_enabled {
        let _min_symbol = t(r.ue_max(65535))?;
        let _max_symbol = t(r.ue_max(65535))?;
        let mask1 = t(r.ue_max(1))?;
        let mask2 = t(r.ue_max(1))?;
        if mask1 != 0 || mask2 != 0 {
            return Err(HeaderError::Unsupported("EFEnonlinear masks"));
        }
        h.efn_u_enabled = t(r.ue_max(1))? == 1;
        h.efn_v_enabled = t(r.ue_max(1))? == 1;
        if h.efn_u_enabled || h.efn_v_enabled {
            let _w = t(r.ue_max(65535))?;
            let _hh = t(r.ue_max(65535))?;
            let n_tiles = t(r.ue_max(65535))?;
            if n_tiles != 1 {
                return Err(HeaderError::Unsupported("EFEnonlinear multi-tile"));
            }
            h.efn_min_luma_q100 = t(r.ue_max(65535))?;
            h.efn_max_luma_q100 = t(r.ue_max(65535))?;
            let cand_num = t(r.ue_max(65535))? as usize;
            h.efn_weights = (0..cand_num)
                .map(|_| t(r.ue_max(65535)))
                .collect::<Result<_, _>>()?;
        }
    }
    h.lef_enabled = t(r.ue_max(1))? == 1;
    if h.lef_enabled {
        h.lef_ch_idx = t(r.ue_max(255))?;
    }
    Ok(h)
}
