// 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.

use std::io::IoSliceMut;

use crate::container::frame_index::FrameIndexBox;
use crate::error::{Error, Result};

use crate::api::{
    JxlBitstreamInput, JxlSignatureType, check_signature_internal, inner::process::SmallBuffer,
};

#[derive(Clone, Debug)]
pub enum ParseState {
    SignatureNeeded,
    BoxNeeded,
    CodestreamBox(u64),
    SkippableBox(u64),
    /// Buffering a jxli box: (remaining bytes, accumulated content).
    BufferingFrameIndex(u64, Vec<u8>),
}

#[derive(Debug)]
enum CodestreamBoxType {
    None,
    Jxlc,
    Jxlp(u32),
    LastJxlp,
}

pub(super) struct BoxParser {
    pub(super) box_buffer: SmallBuffer,
    pub(super) state: ParseState,
    box_type: CodestreamBoxType,
    /// Parsed frame index box, if present in the file.
    pub(super) frame_index: Option<FrameIndexBox>,
    /// Total file bytes consumed from the underlying input.
    pub(super) total_file_consumed: u64,
    skip_jxlp_checks: bool,
}

impl BoxParser {
    pub(super) fn new() -> Self {
        BoxParser {
            box_buffer: SmallBuffer::new(128),
            state: ParseState::SignatureNeeded,
            box_type: CodestreamBoxType::None,
            frame_index: None,
            total_file_consumed: 0,
            skip_jxlp_checks: false,
        }
    }

    // Reads input until the next byte of codestream is available.
    // This function might over-read bytes. Thus, the contents of self.box_buffer should always be
    // read after this function call.
    // Returns the number of codestream bytes that will be available to be read after this call,
    // including any bytes in self.box_buffer.
    // Might return `u64::MAX`, indicating that the rest of the file is codestream.
    pub(super) fn get_more_codestream(&mut self, input: &'color:red'>mut dyn JxlBitstreamInput) -> Result<u64> {
        loop {
            match self.state.clone() {
                ParseState::SignatureNeeded => {
                    let read = self.box_buffer.refill(|b| input.read(b), None)?;
                    self.total_file_consumed += read as u64;
                    match check_signature_internal(&self.box_buffer)? {
                        None => return Err(Error::InvalidSignature),
                        Some(JxlSignatureType::Codestream) => {
                            self.state = ParseState::CodestreamBox(u64::MAX);
                            return Ok(u64::MAX);
                        }
                        Some(JxlSignatureType::Container) => {
                            self.box_buffer
                                .consume(JxlSignatureType::Container.signature().len());
                            self.state = ParseState::BoxNeeded;
                        }
                    }
                }
                ParseState::CodestreamBox(b) => {
                    return Ok(b);
                }
                ParseState::SkippableBox(mut s) => {
                    let num = s.min(usize::MAX as u64) as usize;
                    let skipped = if !self.box_buffer.is_empty() {
                        self.box_buffer.consume(num)
                    } else {
                        let skipped = input.skip(num)?;
                        self.total_file_consumed += skipped as u64;
                        skipped
                    };
                    if skipped == 0 {
                        return Err(Error::OutOfBounds(num));
                    }
                    s -= skipped as u64;
                    if s == 0 {
                        self.state = ParseState::BoxNeeded;
                    } else {
                        self.state = ParseState::SkippableBox(s);
                    }
                }
                ParseState::BufferingFrameIndex(mut remaining, mut buf) => {
                    let num = remaining.min(usize::MAX as u64) as usize;
                    if !self.box_buffer.is_empty() {
                        let take = num.min(self.box_buffer.len());
                        buf.extend_from_slice(&self.box_buffer[..take]);
                        self.box_buffer.consume(take);
                        remaining -= take as u64;
                    } else {
                        let old_len = buf.len();
                        buf.resize(old_len + num, 0);
                        let read = input.read(&mut [IoSliceMut::new(&mut buf[old_len..])])?;
                        self.total_file_consumed += read as u64;
                        if read == 0 {
                            return Err(Error::OutOfBounds(num));
                        }
                        buf.truncate(old_len + read);
                        remaining -= read as u64;
                    }
                    if remaining == 0 {
                        // Parse the buffered frame index box.
                        self.frame_index = Some(FrameIndexBox::parse(&buf)?);
                        self.state = ParseState::BoxNeeded;
                    } else {
                        self.state = ParseState::BufferingFrameIndex(remaining, buf);
                    }
                }
                ParseState::BoxNeeded => {
                    let read = self.box_buffer.refill(|b| input.read(b), None)?;
                    self.total_file_consumed += read as u64;
                    let min_len = match &self.box_buffer[..] {
                        [0, 0, 0, 1, ..] => 16,
                        _ => 8,
                    };
                    if self.box_buffer.len() <= min_len {
                        return Err(Error::OutOfBounds(min_len - self.box_buffer.len()));
                    }
                    let ty: [_; 4] = self.box_buffer[4..8].try_into().unwrap();
                    let extra_len = if &ty == b"jxlp" { 4 } else { 0 };
                    if self.box_buffer.len() <= min_len + extra_len {
                        return Err(Error::OutOfBounds(
                            min_len + extra_len - self.box_buffer.len(),
                        ));
                    }
                    let box_len = match &self.box_buffer[..] {
                        [0, 0, 0, 1, ..] => {
                            u64::from_be_bytes(self.box_buffer[8..16].try_into().unwrap())
                        }
                        _ => u32::from_be_bytes(self.box_buffer[0..4].try_into().unwrap()) as u64,
                    };
                    // Per JXL spec: jxlc box with length 0 has special meaning "extends to EOF"
                    let content_len = if box_len == 0 && (&ty == b"jxlp" || &ty == b"jxlc") {
                        u64::MAX
                    } else {
                        if box_len <= (min_len + extra_len) as u64 {
                            return Err(Error::InvalidBox);
                        }
                        box_len - min_len as u64 - extra_len as u64
                    };
                    match &ty {
                        b"jxlc" => {
                            if matches!(
                                self.box_type,
                                CodestreamBoxType::Jxlp(..) | CodestreamBoxType::LastJxlp
                            ) {
                                return Err(Error::InvalidBox);
                            }
                            self.box_type = CodestreamBoxType::Jxlc;
                            self.state = ParseState::CodestreamBox(content_len);
                        }
                        b"jxlp" => {
                            let index = u32::from_be_bytes(
                                self.box_buffer[min_len..min_len + 4].try_into().unwrap(),
                            );
                            let last = index & 0x80000000 != 0;
                            let idx = index & 0x7fffffff;
                            if !self.skip_jxlp_checks {
                                let wanted_idx = match self.box_type {
                                    CodestreamBoxType::Jxlc | CodestreamBoxType::LastJxlp => {
                                        return Err(Error::InvalidBox);
                                    }
                                    CodestreamBoxType::None => 0,
                                    CodestreamBoxType::Jxlp(i) => i + 1,
                                };
                                if idx != wanted_idx {
                                    return Err(Error::InvalidBox);
                                }
                            }
                            self.box_type = if last {
                                CodestreamBoxType::LastJxlp
                            } else {
                                CodestreamBoxType::Jxlp(idx)
                            };
                            self.state = ParseState::CodestreamBox(content_len);
                        }
                        b"jxli" => {
                            if content_len == u64::MAX {
                                return Err(Error::InvalidBox);
                            }
                            // Reasonable size limit for a frame index box (16 MB).
                            if content_len > 16 * 1024 * 1024 {
                                self.state = ParseState::SkippableBox(content_len);
                            } else {
                                self.state = ParseState::BufferingFrameIndex(
                                    content_len,
                                    Vec::with_capacity(content_len as usize),
                                );
                            }
                        }
                        _ => {
                            self.state = ParseState::SkippableBox(content_len);
                        }
                    }
                    self.box_buffer.consume(min_len + extra_len);
                }
            }
        }
    }

    /// Accounts file bytes consumed directly by codestream parser reads/skips.
    pub(super) fn mark_file_consumed(&mut self, amount: usize) {
        self.total_file_consumed += amount as u64;
    }

    /// Resets the box parser for seeking to a specific codestream position.
    ///
    /// Sets the parser to `CodestreamBox(remaining)` state with cleared
    /// buffers.  The caller must provide raw input starting from the file
    /// position that corresponds to the target codestream offset.
    ///
    /// `remaining` is the number of codestream bytes left in the current
    /// box from the target file position.  For bare-codestream files this
    /// is `u64::MAX`.
    pub(super) fn reset_for_codestream_seek(&mut self, remaining: u64) {
        self.box_buffer = SmallBuffer::new(128);
        self.state = ParseState::CodestreamBox(remaining);
        self.skip_jxlp_checks = true;
        // Keep frame_index unchanged.
    }

    pub(super) fn consume_codestream(&mut self, amount: u64) {
        if let ParseState::CodestreamBox(cb) = &mut self.state {
            *cb = cb.checked_sub(amount).unwrap();
            if *cb == 0 {
                self.state = ParseState::BoxNeeded;
            }
        } else if amount != 0 {
            unreachable!()
        }
    }
}