// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#[cfg(feature = "brotli")]
use std::borrow::Cow;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::io::{IoSliceMut, Read};

use crate::api::{JxlBitstreamInput, JxlSignature, ProfileLevel, SmallBuffer, check_signature};
use crate::error::{Error, Result};
#[cfg(feature = "brotli")]
use crate::util::NewWithCapacity;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ParseState {
    SignatureNeeded,
    // We need at least this many bytes in the buffer to determine what to do
    // with the next box.
    BoxNeeded(usize),
    // The next bytes are raw codestream. None = no limit.
    Codestream(Option<u64>),
    // Skip the next bytes.
    Skip(Option<u64>),
    // The next bytes are auxiliary box contents. None = no limit.
    Aux(Option<u64>),
    // The next bytes should be buffered in a jxlp box.
    OOOJxlp(u32, Option<u64>),
    // After the last codestream box, no more container bytes: no further
    // codestream in file.
    Complete,
}

// Relevant box types.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CodestreamBoxType {
    None,
    Jxlc,
    // (index, is_last)
    Jxlp(u32, bool),
}

impl CodestreamBoxType {
    fn is_last(&self) -> bool {
        matches!(self, Self::Jxlc | Self::Jxlp(_, true))
    }
}

struct OOOJxlpBox {
    data: Vec<u8>,
    consumed: usize,
    is_last: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct JxlAuxBoxType(pub [u8; 4]);

impl JxlAuxBoxType {
    pub const EXIF: Self = JxlAuxBoxType(*b"Exif");
    pub const XMP: Self = JxlAuxBoxType(*b"xml ");
    pub const JUMBF: Self = JxlAuxBoxType(*b"jumb");
    pub const JPEG_RECONSTRUCTION_DATA: Self = JxlAuxBoxType(*b"jbrd");
    pub const HDR_GAIN_MAP: Self = JxlAuxBoxType(*b"jhgm");
}

pub struct JxlAuxBox {
    ty: JxlAuxBoxType,
    data: Vec<u8>,
    brotli: bool,
}

impl std::fmt::Debug for JxlAuxBox {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let len = self.data.len();
        f.debug_struct("JxlAuxBox")
            .field("ty", &self.ty)
            .field(
                "data",
                &format_args!("({len} byte{})", if len == 1 { "" } else { "s" }),
            )
            .field("brotli", &self.brotli)
            .finish()
    }
}

impl JxlAuxBox {
    pub fn box_type(&self) -> JxlAuxBoxType {
        self.ty
    }

    /// Returns the raw box data, potentially compressed.
    pub fn raw_data(&self) -> &[u8] {
        &self.data
    }

    /// Returns if the raw data is Brotli-compressed.
    pub fn is_compressed(&self) -> bool {
        self.brotli
    }

    #[cfg(feature = "brotli")]
    pub fn data(&self, mut trailing_data: &[u8]) -> Result<Cow<'_, [u8]>> {
        use std::io::prelude::*;

        let output = if self.is_compressed() {
            let mut data = &self.data[..];
            let r = (&mut data).chain(&mut trailing_data);
            let mut output = Vec::new();
            let mut brotli = brotli_decompressor::Decompressor::new(r, 1024);
            let mut buf = [0u8; 1024];
            loop {
                let n = brotli.read(&mut buf).map_err(crate::error::Error::Brotli)?;
                if n == 0 {
                    break;
                }
                output.try_reserve(n)?;
                output.extend_from_slice(&buf[..n]);
            }
            Cow::Owned(output)
        } else if trailing_data.is_empty() {
            Cow::Borrowed(&self.data[..])
        } else {
            let mut output = Vec::new_with_capacity(self.data.len() + trailing_data.len())?;
            output.extend_from_slice(&self.data);
            output.extend_from_slice(trailing_data);
            Cow::Owned(output)
        };
        Ok(output)
    }
}

#[derive(Debug, Clone, Copy)]
pub(crate) struct BoxParserCheckpoint {
    box_type: CodestreamBoxType,
    pub(crate) file_position: u64,
    codestream_left: Option<u64>,
    is_valid_checkpoint: bool,
    pub(super) consumed_codestream: u64,
    next_aux_box_idx: usize,
}

pub(super) struct BoxParser {
    local_buffer: SmallBuffer,
    version: Option<u32>,
    ooo_jxlp_buffer: HashMap<u32, OOOJxlpBox>,
    state: ParseState,
    latest_codestream_box: CodestreamBoxType,
    // Map from codestream position at the start of the box
    // to box info
    codestream_pos_to_box: BTreeMap<u64, BoxParserCheckpoint>,
    allow_checkpoint: bool,
    container_level: Option<ProfileLevel>,
    aux: AuxBoxState,
    input_closed: bool,
}

#[derive(Default)]
struct AuxBoxState {
    boxes_to_extract: HashSet<JxlAuxBoxType>,
    boxes: HashMap<JxlAuxBoxType, Vec<JxlAuxBox>>,
    box_buffer: Option<JxlAuxBox>,
    next_box_idx: usize,
    seen_box_count: usize,
}

impl BoxParser {
    pub(super) fn with_aux_boxes(box_types: impl IntoIterator<Item = JxlAuxBoxType>) -> Self {
        BoxParser {
            local_buffer: SmallBuffer::new(128),
            state: ParseState::SignatureNeeded,
            latest_codestream_box: CodestreamBoxType::None,
            ooo_jxlp_buffer: HashMap::new(),
            version: None,
            container_level: None,
            codestream_pos_to_box: BTreeMap::new(),
            allow_checkpoint: true,
            aux: AuxBoxState {
                boxes_to_extract: box_types.into_iter().collect(),
                ..Default::default()
            },
            input_closed: false,
        }
    }

    fn start_aux_box(&mut self, box_type: JxlAuxBoxType, brotli: bool, content_len: Option<u64>) {
        let box_idx = self.aux.next_box_idx;
        self.aux.next_box_idx += 1;
        let is_new = box_idx >= self.aux.seen_box_count;
        self.aux.seen_box_count = self.aux.seen_box_count.max(self.aux.next_box_idx);

        if is_new && self.aux.boxes_to_extract.contains(&box_type) {
            self.aux.box_buffer = Some(JxlAuxBox {
                ty: box_type,
                brotli,
                data: Vec::new(),
            });
            self.state = ParseState::Aux(content_len);
        } else {
            self.state = ParseState::Skip(content_len);
        }
    }

    pub(super) fn reset(&mut self) {
        self.local_buffer.reset();
        self.version = None;
        self.ooo_jxlp_buffer.clear();
        self.state = ParseState::SignatureNeeded;
        self.latest_codestream_box = CodestreamBoxType::None;
        self.codestream_pos_to_box.clear();
        self.allow_checkpoint = true;
        self.container_level = None;
        self.aux.boxes.clear();
        self.aux.box_buffer = None;
        self.aux.next_box_idx = 0;
        self.aux.seen_box_count = 0;
        self.input_closed = false;
    }

    pub(super) fn reset_to_checkpoint(&mut self, box_checkpoint: BoxParserCheckpoint) {
        self.local_buffer.reset();
        self.ooo_jxlp_buffer.clear();
        self.latest_codestream_box = box_checkpoint.box_type;
        self.state = ParseState::Codestream(box_checkpoint.codestream_left);
        // Do not allow creating new checkpoints after a reset.
        self.allow_checkpoint = false;

        if self.aux.box_buffer.take().is_some() {
            self.aux.seen_box_count -= 1;
        }
        self.aux.next_box_idx = box_checkpoint.next_aux_box_idx;
        self.input_closed = false;
    }

    pub(super) fn close_input(&mut self) {
        self.input_closed = true;
    }

    pub(super) fn state_checkpoint(
        &self,
        codestream_pos: u64,
    ) -> Result<Option<BoxParserCheckpoint>> {
        if !self.allow_checkpoint {
            Ok(None)
        } else {
            // last element with key <= codestream_pos.
            let (&start, &(mut b)) = self
                .codestream_pos_to_box
                .range(..=codestream_pos)
                .last()
                .unwrap();
            if !b.is_valid_checkpoint {
                // Disallow files that have frame starts in boxes that are not
                // valid checkpoints.
                return Err(Error::InvalidBox);
            }
            let overflow = codestream_pos - start;
            if let Some(bytes) = &mut b.codestream_left {
                *bytes -= overflow;
            }
            b.file_position += overflow;
            b.consumed_codestream = codestream_pos;
            Ok(Some(b))
        }
    }

    pub(super) fn total_bytes_consumed(&self, codestream_bytes_consumed: u64) -> u64 {
        let (start, b) = self.codestream_pos_to_box.last_key_value().unwrap();
        b.file_position + codestream_bytes_consumed.saturating_sub(*start)
    }

    pub(super) fn aux_boxes(&self, box_type: JxlAuxBoxType) -> &[JxlAuxBox] {
        self.aux
            .boxes
            .get(&box_type)
            .map(|v| &**v)
            .unwrap_or_default()
    }

    pub(super) fn container_level(&self) -> Option<ProfileLevel> {
        self.container_level
    }

    pub(super) fn trailing_box(&self) -> Option<&JxlAuxBox> {
        if self.state != ParseState::Aux(None) {
            return None;
        }

        self.aux.box_buffer.as_ref()
    }

    pub(super) fn is_codestream_complete(&self) -> bool {
        if !self.latest_codestream_box.is_last() {
            return false;
        }

        match self.state {
            ParseState::Codestream(Some(0)) => true,
            ParseState::Codestream(_) => false,
            _ => true,
        }
    }

    fn add_checkpoint(&mut self) {
        if !self.allow_checkpoint {
            return;
        }
        // For a codestream box to be a valid checkpoint, it needs to be decoded
        // in-order (i.e. all other logically-preceding jxlp boxes have been seen already)
        // and we must not have seen any future jxlp boxes (which would be stored in the
        // ooo jxlp buffer).
        let (codestream_left, is_valid_checkpoint) = match self.state {
            ParseState::Codestream(c) => (c, self.ooo_jxlp_buffer.is_empty()),
            ParseState::OOOJxlp(_, c) => (c, false),
            _ => (None, false),
        };
        let codestream_position = self
            .codestream_pos_to_box
            .last_key_value()
            .map(|(p, c)| *p + c.codestream_left.unwrap())
            .unwrap_or(0);
        self.codestream_pos_to_box.insert(
            codestream_position,
            BoxParserCheckpoint {
                file_position: self.local_buffer.consumed(),
                box_type: self.latest_codestream_box,
                codestream_left,
                is_valid_checkpoint,
                consumed_codestream: 0,
                next_aux_box_idx: self.aux.next_box_idx,
            },
        );
    }

    fn injected_jxlp(&mut self) -> Option<&mut OOOJxlpBox> {
        // Skip 0-sized jxlp boxes.
        while matches!(self.state, ParseState::BoxNeeded(_) | ParseState::Complete)
            && let CodestreamBoxType::Jxlp(j, _) = self.latest_codestream_box
            && self
                .ooo_jxlp_buffer
                .get(&j)
                .is_some_and(|b| b.consumed == b.data.len())
        {
            self.check_ooo_jxlp_done();
        }
        if matches!(self.state, ParseState::BoxNeeded(_) | ParseState::Complete)
            && let CodestreamBoxType::Jxlp(j, _) = self.latest_codestream_box
            && let Some(b) = self.ooo_jxlp_buffer.get_mut(&j)
        {
            Some(b)
        } else {
            None
        }
    }

    fn check_ooo_jxlp_done(&mut self) {
        let CodestreamBoxType::Jxlp(j, is_last) = self.latest_codestream_box else {
            unreachable!()
        };
        let Some(b) = self.ooo_jxlp_buffer.get_mut(&j) else {
            unreachable!()
        };
        if b.consumed == b.data.len() {
            self.ooo_jxlp_buffer.remove(&j);
            if is_last {
                self.state = ParseState::Complete;
            } else if let Some(b) = self.ooo_jxlp_buffer.get(&(j + 1)) {
                self.latest_codestream_box = CodestreamBoxType::Jxlp(j + 1, b.is_last);
            }
        }
    }

    fn mark_codestream_read(&mut self, n: usize) {
        let ParseState::Codestream(c) = &mut self.state else {
            unreachable!()
        };
        if c.is_none() {
            return;
        }
        let new_size = c.unwrap() - n as u64;
        if new_size == 0 {
            self.state = if self.latest_codestream_box.is_last() {
                ParseState::Complete
            } else {
                if let CodestreamBoxType::Jxlp(j, _) = self.latest_codestream_box
                    && let Some(b) = self.ooo_jxlp_buffer.get(&(j + 1))
                {
                    self.latest_codestream_box = CodestreamBoxType::Jxlp(j + 1, b.is_last);
                    self.add_checkpoint();
                }
                ParseState::BoxNeeded(8)
            }
        } else {
            *c = Some(new_size);
        }
    }

    fn available_bytes_inner(&mut self, input: &mut dyn JxlBitstreamInput) -> Result<usize> {
        Ok(self.local_buffer.len() + input.available_bytes()?)
    }

    fn read_inner(
        &mut self,
        input: &mut dyn JxlBitstreamInput,
        mut buf: &mut [IoSliceMut],
    ) -> Result<usize> {
        let mut n = 0;
        if !self.local_buffer.is_empty() {
            let mut ll = &self.local_buffer[..];
            n = ll.read_vectored(buf)?;
            IoSliceMut::advance_slices(&mut buf, n);
            self.local_buffer.consume(n);
        }
        let read = input.read(buf)?;
        self.local_buffer.mark_consumed(read as u64);
        Ok(n + read)
    }

    fn skip_inner(
        &mut self,
        input: &mut dyn JxlBitstreamInput,
        mut amount: usize,
    ) -> Result<usize> {
        let mut n = 0;
        if !self.local_buffer.is_empty() {
            let ll = self.local_buffer.len();
            n = ll.min(amount);
            amount -= n;
            self.local_buffer.consume(n);
        }
        let skip = input.skip(amount)?;
        self.local_buffer.mark_consumed(skip as u64);
        Ok(n + skip)
    }

    fn read_until_at_least(
        &mut self,
        input: &mut dyn JxlBitstreamInput,
        min_size: usize,
    ) -> Result<()> {
        loop {
            let n = self.local_buffer.refill(|b| Ok(input.read(b)?))?;
            if n == 0 || self.local_buffer.len() >= min_size {
                break;
            }
        }
        if self.local_buffer.len() < min_size {
            return Err(Error::OutOfBounds(min_size - self.local_buffer.len()));
        }
        Ok(())
    }

    fn advance_to_codestream(&mut self, input: &mut dyn JxlBitstreamInput) -> Result<()> {
        loop {
            match self.state {
                ParseState::Codestream(Some(0)) => self.state = ParseState::BoxNeeded(8),
                ParseState::Complete | ParseState::Codestream(_) => return Ok(()),
                ParseState::SignatureNeeded => {
                    let codestream_signature_len = JxlSignature::Codestream.signature_len();
                    self.read_until_at_least(input, codestream_signature_len)?;
                    match check_signature(&self.local_buffer) {
                        JxlSignature::NeedMoreInput { size_hint } => {
                            return Err(Error::OutOfBounds(size_hint));
                        }
                        JxlSignature::None => {
                            return Err(Error::InvalidSignature);
                        }
                        JxlSignature::Codestream => {
                            self.state = ParseState::Codestream(None);
                            self.latest_codestream_box = CodestreamBoxType::Jxlc;
                            self.add_checkpoint();
                            return Ok(());
                        }
                        JxlSignature::Container => {
                            let l = JxlSignature::Container.signature_len();
                            self.local_buffer.consume(l);
                            self.state = ParseState::BoxNeeded(8);
                        }
                    }
                }
                ParseState::Skip(count) => {
                    if count == Some(0) {
                        self.state = ParseState::BoxNeeded(8);
                        continue;
                    }
                    let to_skip = count.unwrap_or(u64::MAX).min(usize::MAX as u64) as usize;
                    let n = self.skip_inner(input, to_skip)? as u64;
                    if n == 0 {
                        return Err(Error::OutOfBounds(to_skip));
                    }
                    self.state = ParseState::Skip(count.map(|x| x - n));
                }
                ParseState::Aux(None) => {
                    let local_buffer_len = self.local_buffer.len();
                    let buf = self.aux.box_buffer.as_mut().unwrap();
                    if local_buffer_len > 0 {
                        buf.data.extend_from_slice(&self.local_buffer);
                        self.local_buffer.consume(local_buffer_len);
                    }
                    return Ok(());
                }
                ParseState::Aux(Some(count)) => {
                    if count == 0 {
                        let buf = self.aux.box_buffer.take().unwrap();
                        self.aux.boxes.entry(buf.ty).or_default().push(buf);
                        self.state = if self.latest_codestream_box.is_last() {
                            ParseState::Complete
                        } else {
                            ParseState::BoxNeeded(8)
                        };
                        continue;
                    }

                    let total = self.handle_aux_box(input, count)?;
                    if total == 0 {
                        return Err(Error::OutOfBounds(count.min(usize::MAX as u64) as usize));
                    }
                    self.state = ParseState::Aux(Some(count - total));
                }
                ParseState::OOOJxlp(id, count) => {
                    if count == Some(0) {
                        self.state = ParseState::BoxNeeded(8);
                        continue;
                    }
                    // Temporarily take the buffer out of the jxlp state to make the borrow checker happy.
                    let mut buf = self.ooo_jxlp_buffer.remove(&id).unwrap();
                    let mut total = 0;
                    loop {
                        let space = buf
                            .data
                            .len()
                            .max(1024)
                            .min(self.available_bytes_inner(input)?)
                            as u64;
                        let space = count.map(|x| (x - total).min(space)).unwrap_or(space) as usize;
                        if space == 0 {
                            break;
                        }
                        buf.data.try_reserve(space)?;
                        let cur = buf.data.len();
                        buf.data.resize(cur + space, 0);
                        let n =
                            self.read_inner(input, &mut [IoSliceMut::new(&mut buf.data[cur..])])?;
                        buf.data.truncate(cur + n);
                        if n == 0 {
                            break;
                        }
                        total += n as u64;
                    }
                    self.ooo_jxlp_buffer.insert(id, buf);
                    if total == 0 {
                        return Err(Error::OutOfBounds(
                            count.map(|x| x.min(usize::MAX as u64)).unwrap_or(1) as usize,
                        ));
                    }
                    self.state = ParseState::OOOJxlp(id, count.map(|x| x - total));
                }
                ParseState::BoxNeeded(min_size) => self.parse_box(input, min_size)?,
            }
        }
    }

    fn consume_trailing_data(&mut self, input: &mut dyn JxlBitstreamInput) -> Result<()> {
        if self.aux.boxes_to_extract.is_empty() {
            return Ok(());
        }
        loop {
            match self.state {
                ParseState::Codestream(None) => {
                    return Ok(());
                }
                ParseState::Complete => {
                    if self.available_bytes_inner(input)? == 0 {
                        if self.input_closed {
                            return Ok(());
                        }
                        return Err(Error::OutOfBounds(8));
                    }
                    self.state = ParseState::BoxNeeded(8);
                }
                ParseState::Skip(None) => return Ok(()),
                ParseState::Codestream(count) | ParseState::Skip(count) => {
                    if count == Some(0) {
                        self.state = ParseState::Complete;
                        continue;
                    }
                    let to_skip = count.unwrap_or(u64::MAX).min(usize::MAX as u64) as usize;
                    let n = self.skip_inner(input, to_skip)? as u64;
                    if n == 0 {
                        if self.input_closed {
                            return Err(Error::InvalidBox);
                        }
                        return Err(Error::OutOfBounds(to_skip));
                    }
                    self.state = ParseState::Skip(count.map(|x| x - n));
                }
                ParseState::Aux(None) => {
                    let local_buffer_len = self.local_buffer.len();
                    let buf = self.aux.box_buffer.as_mut().unwrap();
                    if local_buffer_len > 0 {
                        buf.data.extend_from_slice(&self.local_buffer);
                        self.local_buffer.consume(local_buffer_len);
                    }
                    return Ok(());
                }
                ParseState::Aux(Some(count)) => {
                    if count == 0 {
                        let buf = self.aux.box_buffer.take().unwrap();
                        self.aux.boxes.entry(buf.ty).or_default().push(buf);
                        self.state = ParseState::Complete;
                        continue;
                    }

                    let total = self.handle_aux_box(input, count)?;
                    if total == 0 {
                        if self.input_closed {
                            return Err(Error::InvalidBox);
                        }
                        return Err(Error::OutOfBounds(count.min(usize::MAX as u64) as usize));
                    }
                    self.state = ParseState::Aux(Some(count - total));
                }
                ParseState::BoxNeeded(min_size) => self.parse_box(input, min_size)?,
                _ => {
                    return Err(Error::InvalidBox);
                }
            }
        }
    }

    fn handle_aux_box(&mut self, input: &mut dyn JxlBitstreamInput, count: u64) -> Result<u64> {
        let mut buf = self.aux.box_buffer.take().unwrap();
        let mut total = 0;
        loop {
            let space = buf
                .data
                .len()
                .max(1024)
                .min(self.available_bytes_inner(input)?) as u64;
            let space = (count - total).min(space) as usize;
            if space == 0 {
                break;
            }
            buf.data.try_reserve(space)?;
            let cur = buf.data.len();
            buf.data.resize(cur + space, 0);
            let n = self.read_inner(input, &mut [IoSliceMut::new(&mut buf.data[cur..])])?;
            buf.data.truncate(cur + n);
            if n == 0 {
                break;
            }
            total += n as u64;
        }
        self.aux.box_buffer = Some(buf);
        Ok(total)
    }

    fn parse_box(&mut self, input: &mut dyn JxlBitstreamInput, required_size: usize) -> Result<()> {
        match self.read_until_at_least(input, required_size) {
            Err(Error::OutOfBounds(_)) if self.input_closed => return Err(Error::InvalidBox),
            res => res?,
        }

        let min_len = match &self.local_buffer[..] {
            [0, 0, 0, 1, ..] => 16,
            _ => 8,
        };

        let ty: [_; 4] = self.local_buffer[4..8].try_into().unwrap();
        if &ty != b"ftyp" && self.version.is_none() {
            // ftyp should be the first box.
            return Err(Error::InvalidBox);
        }

        let extra_len = match &ty {
            b"jxlp" => 4,
            b"ftyp" => 8,
            b"brob" => 4,
            b"jxll" => 1,
            _ => 0,
        };

        if self.local_buffer.len() < extra_len + min_len {
            if self.input_closed {
                return Err(Error::InvalidBox);
            }
            self.state = ParseState::BoxNeeded(extra_len + min_len);
            return Ok(());
        }

        let box_len = match &self.local_buffer[..] {
            [0, 0, 0, 1, ..] => u64::from_be_bytes(self.local_buffer[8..16].try_into().unwrap()),
            _ => u32::from_be_bytes(self.local_buffer[0..4].try_into().unwrap()) as u64,
        };

        self.local_buffer.consume(min_len);

        let content_len = if box_len == 0 {
            None
        } else if let Some(d) = box_len.checked_sub((extra_len + min_len) as u64) {
            Some(d)
        } else {
            return Err(Error::InvalidBox);
        };

        match &ty {
            b"ftyp" => {
                if self.version.is_some() {
                    return Err(Error::InvalidBox);
                }
                if &self.local_buffer[0..4] != b"jxl " {
                    return Err(Error::InvalidBox);
                }
                let ver = u32::from_be_bytes(self.local_buffer[4..8].try_into().unwrap());
                if ver > 1 {
                    return Err(Error::InvalidBox);
                }
                self.version = Some(ver);
                self.local_buffer.consume(8);
                self.state = ParseState::Skip(content_len);
            }
            b"jxlc" => {
                if matches!(self.latest_codestream_box, CodestreamBoxType::Jxlp(..)) {
                    return Err(Error::InvalidBox);
                }
                self.latest_codestream_box = CodestreamBoxType::Jxlc;
                self.state = ParseState::Codestream(content_len);
                self.add_checkpoint();
            }
            b"jxlp" => {
                let index = u32::from_be_bytes(self.local_buffer[..4].try_into().unwrap());
                self.local_buffer.consume(4);
                let last = index & 0x80000000 != 0;
                let idx = index & 0x7fffffff;
                let wanted_idx = match self.latest_codestream_box {
                    CodestreamBoxType::Jxlc | CodestreamBoxType::Jxlp(_, true) => {
                        return Err(Error::InvalidBox);
                    }
                    CodestreamBoxType::None => 0,
                    CodestreamBoxType::Jxlp(i, _) => i + 1,
                };
                if idx < wanted_idx {
                    return Err(Error::InvalidBox);
                }
                if idx > wanted_idx {
                    // Out-of-order jxlp
                    if self.version != Some(1) {
                        return Err(Error::InvalidBox);
                    }
                    if self.ooo_jxlp_buffer.contains_key(&idx) {
                        return Err(Error::InvalidBox);
                    }
                    if content_len.is_none() {
                        return Err(Error::InvalidBox);
                    }
                    self.state = ParseState::OOOJxlp(idx, content_len);
                    self.ooo_jxlp_buffer.insert(
                        idx,
                        OOOJxlpBox {
                            data: vec![],
                            consumed: 0,
                            is_last: last,
                        },
                    );
                    // We create a known-invalid checkpoint here to ensure
                    // that the checks work properly, and to count codestream
                    // length correctly.
                    self.add_checkpoint();
                    return Ok(());
                }
                // In-order jxlp
                self.latest_codestream_box = CodestreamBoxType::Jxlp(idx, last);
                self.state = ParseState::Codestream(content_len);
                self.add_checkpoint();
            }
            b"brob" => {
                let inner_type = &self.local_buffer[..4];
                if matches!(inner_type, b"brob" | [b'j', b'x', b'l', _] | b"jbrd") {
                    return Err(Error::InvalidBox);
                }
                let inner_type = JxlAuxBoxType(inner_type.try_into().unwrap());
                self.local_buffer.consume(4);

                self.start_aux_box(inner_type, true, content_len);
            }
            b"jxll" => {
                if self.version.is_none() || self.latest_codestream_box != CodestreamBoxType::None {
                    return Err(Error::InvalidBox);
                }
                if self.container_level.is_some() {
                    return Err(Error::InvalidBox);
                }
                if content_len != Some(0) {
                    return Err(Error::InvalidBox);
                }
                let level = match self.local_buffer[0] {
                    5 => ProfileLevel::Main5,
                    10 => ProfileLevel::Main10,
                    _ => return Err(Error::InvalidBox),
                };
                self.local_buffer.consume(1);
                self.container_level = Some(level);
                self.state = ParseState::BoxNeeded(8);
            }
            code => {
                let ty = JxlAuxBoxType(*code);
                self.start_aux_box(ty, false, content_len);
            }
        }
        Ok(())
    }
}

pub(super) struct CodestreamInput<'a> {
    box_parser: &'a mut BoxParser,
    input: &'a mut dyn JxlBitstreamInput,
}

impl<'a> CodestreamInput<'a> {
    pub(super) fn new(box_parser: &'a mut BoxParser, input: &'a mut dyn JxlBitstreamInput) -> Self {
        Self { box_parser, input }
    }

    pub(super) fn box_parser(&self) -> &BoxParser {
        self.box_parser
    }

    pub(super) fn consume_trailing_data(&mut self) -> Result<()> {
        self.box_parser.consume_trailing_data(self.input)
    }

    // The methods below have the same semantics as the methods in JxlBitstreamInput.

    pub(super) fn available_bytes(&mut self) -> Result<usize> {
        if let Some(b) = self.box_parser.injected_jxlp() {
            return Ok(b.data.len() - b.consumed);
        }

        self.box_parser.advance_to_codestream(self.input)?;
        let ParseState::Codestream(r) = self.box_parser.state else {
            return Ok(0);
        };
        let available = self.box_parser.available_bytes_inner(self.input)?;
        Ok(r.map(|x| x.min(available as u64) as usize)
            .unwrap_or(available))
    }

    pub(super) fn read(&mut self, buf: &mut [IoSliceMut]) -> Result<usize> {
        if let Some(b) = self.box_parser.injected_jxlp() {
            let mut r = &b.data[b.consumed..];
            let n = r.read_vectored(buf)?;
            b.consumed += n;
            self.box_parser.check_ooo_jxlp_done();
            return Ok(n);
        }

        self.box_parser.advance_to_codestream(self.input)?;
        let ParseState::Codestream(r) = self.box_parser.state else {
            return Ok(0);
        };
        let limit = r.unwrap_or(u64::MAX);
        // If the input buffers have more space than the remaining codestream in the current
        // box, split the read to ensure we don't over-read.
        let n = if buf.iter().map(|x| x.len() as u64).sum::<u64>() <= limit {
            self.box_parser.read_inner(self.input, buf)?
        } else {
            let mut cumsum = 0;
            let mut upper = 0;
            while cumsum + buf[upper].len() as u64 <= limit {
                cumsum += buf[upper].len() as u64;
                upper += 1;
            }
            let n = self.box_parser.read_inner(self.input, &mut buf[..upper])?;
            let left = (limit - n as u64) as usize;
            n + self
                .box_parser
                .read_inner(self.input, &mut [IoSliceMut::new(&mut buf[upper][..left])])?
        };
        self.box_parser.mark_codestream_read(n);
        Ok(n)
    }

    pub(super) fn skip(&mut self, amount: usize) -> Result<usize> {
        if let Some(b) = self.box_parser.injected_jxlp() {
            let n = (b.data.len() - b.consumed).min(amount);
            b.consumed += n;
            self.box_parser.check_ooo_jxlp_done();
            return Ok(n);
        }

        self.box_parser.advance_to_codestream(self.input)?;
        let ParseState::Codestream(r) = self.box_parser.state else {
            return Ok(0);
        };
        let n = self.box_parser.skip_inner(
            self.input,
            r.unwrap_or(u64::MAX).min(amount as u64) as usize,
        )?;
        self.box_parser.mark_codestream_read(n);
        Ok(n)
    }
}

#[cfg(test)]
mod tests {
    use std::io::IoSliceMut;

    use super::{BoxParser, CodestreamInput};

    /// Regression: a zero-length skippable box must not leave the parser stuck at
    /// `SkippableBox(0)` when more container input is available.
    #[test]
    fn zero_length_skippable_box_does_not_hang() {
        let data = include_bytes!("../../tests/testdata/zero_length_skippable_box.jxl");
        let mut parser = BoxParser::with_aux_boxes(None);
        let mut input = data.as_slice();

        {
            let mut input = CodestreamInput::new(&mut parser, &mut input);
            let mut buf = vec![0; 1024];
            while matches!(input.read(&mut [IoSliceMut::new(&mut buf)]), Ok(x) if x > 0) {}
        }

        // All input should be consumed.
        assert!(input.is_empty());
    }
}