// 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::collections::{HashSet, VecDeque};
use std::io::IoSliceMut;

use crate::api::box_parser::CodestreamInput;
use crate::api::codestream_parser::{
    CodestreamParser, ProcessMode, check_size_limit, validate_output_buffers,
};
use crate::api::{
    JxlColorProfile, JxlDecoderOptions, JxlOutputBuffer, JxlParallelRunner, JxlPixelFormat,
    SmallBuffer,
};
use crate::bit_reader::BitReader;
use crate::error::{Error, Result};
use crate::frame::{DecoderState, Frame, HfMetaSplitter, LfImageSplitter, Section};
use crate::headers::FileHeader;
use crate::headers::encodings::UnconditionalCoder;
use crate::headers::frame_header::{Encoding, FrameHeader, FrameType};
use crate::headers::toc::{IncrementalTocReader, Toc};
use crate::image::BufferRecycler;
use crate::util::NewWithCapacity;
use crate::util::sync::Arc;

struct SectionBuffer {
    len: usize,
    data: Vec<u8>,
    section: Section,
}

#[derive(Debug)]
struct SectionState {
    lf_global_done: bool,
    remaining_lf: usize,
    hf_global_done: bool,
    completed_passes: Vec<u8>,
    lf_global_flush_len: usize,
}

impl SectionState {
    fn new(num_lf_groups: usize, num_groups: usize) -> Self {
        Self {
            lf_global_done: false,
            remaining_lf: num_lf_groups,
            hf_global_done: false,
            completed_passes: vec![0; num_groups],
            lf_global_flush_len: 0,
        }
    }
}

pub struct FrameInfo {
    frame_header: Option<FrameHeader>,
    toc_parser: Option<IncrementalTocReader>,
    frame: Option<Box<Frame>>,
    // Keeps track of whether pixels have been modified.
    pixels_dirty: bool,
    buffer_recycler: Arc<BufferRecycler>,

    // Section information.
    sections: VecDeque<SectionBuffer>,
    section_size: u64,
    ready_section_data: usize,
    section_state: SectionState,

    // Or only section if in single section special case.
    lf_global_section: Option<SectionBuffer>,
    lf_sections: Vec<SectionBuffer>,
    hf_global_section: Option<SectionBuffer>,
    // indexed by group, then by pass.
    hf_sections: Vec<Vec<Option<SectionBuffer>>>,
    // group indices that *might* have new renderable data.
    candidate_hf_sections: HashSet<usize>,
    // The frame counters before the next frame (visible frames, and non-visible frames since the
    // last visible one), which seed its noise. They are counted here rather than in the decoder
    // state so that frames skipped by a seek count too.
    frame_counters: (usize, usize),
    // Which stored frames to keep in the next frame's decoder state (see `reset`).
    keep_slots: [bool; 8],
}

impl FrameInfo {
    pub fn new() -> Self {
        Self {
            frame_header: None,
            toc_parser: None,
            frame: None,
            sections: VecDeque::new(),
            section_size: 0,
            ready_section_data: 0,
            section_state: SectionState::new(0, 0),
            lf_global_section: None,
            lf_sections: vec![],
            hf_global_section: None,
            hf_sections: vec![],
            candidate_hf_sections: HashSet::new(),
            pixels_dirty: false,
            buffer_recycler: Arc::new(BufferRecycler::new()),
            frame_counters: (0, 0),
            keep_slots: [true; 8],
        }
    }

    /// Prepares for a next frame preceded by `frame_counters`, keeping only the stored frames in
    /// `keep_slots`.
    pub fn reset(&mut self, frame_counters: (usize, usize), keep_slots: [bool; 8]) {
        self.clear();
        if keep_slots == [false; 8] {
            self.frame = None;
            self.pixels_dirty = false;
        }
        self.frame_counters = frame_counters;
        self.keep_slots = keep_slots;
    }

    pub fn clear(&mut self) {
        self.frame_header = None;
        self.toc_parser = None;
        self.section_size = 0;
        self.ready_section_data = 0;

        // Clear sections
        self.sections.clear();
        self.section_state = SectionState::new(0, 0);
        self.lf_global_section = None;
        self.lf_sections.clear();
        self.hf_global_section = None;
        self.hf_sections.clear();
        self.candidate_hf_sections.clear();
    }

    pub fn current_frame_header(&self) -> Option<&FrameHeader> {
        self.frame_header
            .as_ref()
            .or(self.frame.as_ref().map(|x| x.header()))
    }

    #[cfg(test)]
    pub fn current_toc(&self) -> Option<&crate::headers::toc::Toc> {
        self.frame.as_ref().map(|x| x.toc())
    }

    pub fn parse_frame_header(
        &mut self,
        is_preview: bool,
        decode_options: &JxlDecoderOptions,
        file_header: &FileHeader,
        br: &mut BitReader,
        bits: &mut usize,
        level5_limits: bool,
    ) -> Result<()> {
        // For preview frames, use the preview dimensions instead of main image dimensions
        let nonserialized = if is_preview {
            file_header
                .preview_frame_header_nonserialized()
                .unwrap_or_else(|| file_header.frame_header_nonserialized())
        } else {
            file_header.frame_header_nonserialized()
        };

        let mut frame_header = FrameHeader::read_unconditional(&(), br, &nonserialized)?;
        frame_header.postprocess(&nonserialized);
        check_size_limit(
            decode_options.sample_limit,
            frame_header.size_upsampled(),
            frame_header.num_extra_channels as usize,
        )?;

        if level5_limits
            && frame_header.duration > 0
            && let Some(anim) = &file_header.image_metadata.animation
            && (frame_header.duration as u128) * (anim.tps_denominator as u128) * 120
                < (anim.tps_numerator as u128)
        {
            let fps = (anim.tps_numerator as u64)
                .div_ceil(((frame_header.duration as u64) * (anim.tps_denominator as u64)).max(1));
            return Err(Error::level_limit(
                level5_limits,
                "frame rate (fps)",
                fps,
                120,
            ));
        }

        self.frame_header = Some(frame_header);
        *bits = br.total_bits_read();
        Ok(())
    }

    pub fn parse_toc(&mut self, br: &mut BitReader, bits: &mut usize) -> Result<Toc> {
        if self.toc_parser.is_none() {
            let num_toc_entries = self.frame_header.as_ref().unwrap().num_toc_entries();
            self.toc_parser = Some(IncrementalTocReader::new(num_toc_entries as u32, br)?);
        }

        let toc_parser = self.toc_parser.as_mut().unwrap();
        *bits = br.total_bits_read();
        while !toc_parser.is_complete() {
            match toc_parser.read_step(br) {
                Ok(()) => *bits = br.total_bits_read(),
                Err(Error::OutOfBounds(c)) => {
                    return Err(Error::OutOfBounds(c));
                }
                Err(e) => return Err(e),
            }
        }
        br.jump_to_byte_boundary()?;
        *bits = br.total_bits_read();
        Ok(self.toc_parser.take().unwrap().finalize())
    }

    #[allow(clippy::too_many_arguments)]
    pub(super) fn make_frame(
        &mut self,
        cbuf: &mut SmallBuffer,
        toc: Toc,
        file_header: &FileHeader,
        decode_options: &JxlDecoderOptions,
        pixel_format: &JxlPixelFormat,
        output_profile: &JxlColorProfile,
        process_mode: ProcessMode,
        level5_limits: bool,
    ) -> Result<()> {
        self.section_size = toc.entries.iter().map(|x| *x as u64).sum();
        self.ready_section_data = 0;

        self.lf_global_section = None;
        self.lf_sections.clear();
        self.hf_global_section = None;
        self.candidate_hf_sections.clear();
        self.hf_sections.clear();

        self.frame_counters = if self.frame_header.as_ref().unwrap().is_visible() {
            (self.frame_counters.0 + 1, 0)
        } else {
            (self.frame_counters.0, self.frame_counters.1 + 1)
        };
        if !matches!(process_mode, ProcessMode::Skip(_)) {
            // If we have a previous frame, compute the decoder state from there.
            // Otherwise, compute a new decoder state.
            // We finalize the previous frame here to allow progressive rendering
            // to work properly if a flush is requested while we parse a frame
            // header.
            let mut decoder_state = self
                .frame
                .take()
                .map(|x| x.finalize())
                .transpose()?
                .unwrap_or_else(|| {
                    DecoderState::new(file_header.clone(), decode_options, level5_limits)
                });
            decoder_state.level5_limits = level5_limits;
            decoder_state.keep_slots(std::mem::replace(&mut self.keep_slots, [true; 8]));
            let mut frame = Frame::from_header_and_toc(
                self.frame_header.take().unwrap(),
                toc,
                decoder_state,
                self.frame_counters,
                self.buffer_recycler.clone(),
            )?;

            let num_groups = frame.header().num_groups();
            let num_passes = frame.header().passes.num_passes as usize;
            let mut hf_sections = Vec::new_with_capacity(num_groups)?;
            for _ in 0..num_groups {
                let mut row = Vec::new_with_capacity(num_passes)?;
                for _ in 0..num_passes {
                    row.push(None);
                }
                hf_sections.push(row);
            }
            self.hf_sections = hf_sections;

            let mut sections: Vec<_> = frame
                .toc()
                .entries
                .iter()
                .map(|x| SectionBuffer {
                    len: *x as usize,
                    data: vec![],
                    section: Section::LfGlobal, // will be fixed later
                })
                .collect();

            let order = if frame.toc().permuted {
                frame.toc().permutation.0.clone()
            } else {
                (0..sections.len() as u32).collect()
            };

            if sections.len() > 1 {
                let base_sections = [Section::LfGlobal, Section::HfGlobal];
                let lf_sections =
                    (0..frame.header().num_lf_groups()).map(|x| Section::Lf { group: x });
                let hf_sections = (0..frame.header().passes.num_passes).flat_map(|p| {
                    (0..frame.header().num_groups()).map(move |g| Section::Hf {
                        group: g,
                        pass: p as usize,
                    })
                });

                for section in base_sections
                    .into_iter()
                    .chain(lf_sections)
                    .chain(hf_sections)
                {
                    sections[order[frame.get_section_idx(section)] as usize].section = section;
                }
            }
            self.sections = sections.into_iter().collect();

            // Move data from the pre-section buffer into the sections.
            // Only allocate as much of each section as the pre-section buffer
            // can fill; the TOC-declared section length is untrusted and may
            // be much larger than the actual input.
            for buf in self.sections.iter_mut() {
                if cbuf.is_empty() {
                    break;
                }
                let target = buf.len.min(cbuf.len());
                let mut data = Vec::new();
                data.try_reserve_exact(target)?;
                data.resize(target, 0);
                buf.data = data;
                let n = cbuf.take(&mut [IoSliceMut::new(&mut buf.data)]);
                self.ready_section_data += n;
            }

            self.section_state =
                SectionState::new(frame.header().num_lf_groups(), frame.header().num_groups());
            frame.prepare_render_pipeline(pixel_format, output_profile)?;
            self.frame = Some(frame);
        } else {
            let num = (cbuf.len() as u64).min(self.section_size) as usize;
            cbuf.consume(num);
            self.ready_section_data += num;
        }

        Ok(())
    }

    pub fn dequeue_ready_sections(&mut self) {
        // Dequeue ready sections.
        while self
            .sections
            .front()
            .is_some_and(|s| s.len <= self.ready_section_data)
        {
            let s = self.sections.pop_front().unwrap();
            self.ready_section_data -= s.len;

            match s.section {
                Section::LfGlobal => {
                    self.lf_global_section = Some(s);
                }
                Section::HfGlobal => {
                    self.hf_global_section = Some(s);
                }
                Section::Lf { .. } => {
                    self.lf_sections.push(s);
                }
                Section::Hf { group, pass } => {
                    self.hf_sections[group][pass] = Some(s);
                    self.candidate_hf_sections.insert(group);
                }
            }
        }
    }

    pub fn skip_sections(
        &mut self,
        input: &mut CodestreamInput,
        buf: &mut SmallBuffer,
    ) -> Result<()> {
        let total_size = self.section_size;
        let need_skip =
            (total_size - self.ready_section_data as u64).min(usize::MAX as u64) as usize;
        let skipped = input.skip(need_skip)?;
        buf.mark_consumed(skipped as u64);
        self.ready_section_data += skipped;
        if (self.ready_section_data as u64) < total_size {
            Err(Error::OutOfBounds(
                (total_size - self.ready_section_data as u64).min(usize::MAX as u64) as usize,
            ))
        } else {
            self.sections.clear();
            Ok(())
        }
    }

    pub fn fill_sections(
        &mut self,
        input: &mut CodestreamInput,
        buf: &mut SmallBuffer,
    ) -> Result<()> {
        // TODO(veluca): consider re-using `section_buffers` between different iterations.
        let mut total = 0;
        loop {
            let available_codestream = input.available_bytes();
            let mut available_codestream = match (available_codestream, total) {
                (Ok(a), _) => a.max(1),
                (Err(Error::OutOfBounds(_)), t) if t > 0 => return Ok(()),
                (Err(e), _) => return Err(e),
            };
            let mut readable_section_data = self.ready_section_data + available_codestream;
            // Ensure enough section buffer space is available for reading the
            // available data. Section lengths come from the untrusted TOC, so
            // grow buffers only as far as the data we can actually read; a
            // truncated stream must not trigger a huge upfront allocation.
            for buf in self.sections.iter_mut() {
                if buf.data.len() < buf.len && buf.data.len() < readable_section_data {
                    let target = buf.len.min(readable_section_data);
                    if target == buf.len {
                        buf.data.try_reserve_exact(target - buf.data.len())?;
                    } else {
                        buf.data.try_reserve(target - buf.data.len())?;
                    }
                    buf.data.resize(target, 0);
                }
                readable_section_data = readable_section_data.saturating_sub(buf.data.len());
                if readable_section_data == 0 {
                    break;
                }
            }
            // Prepare buffers to read into.
            let mut section_buffers = vec![];
            let mut ready = self.ready_section_data;
            for buf in self.sections.iter_mut() {
                if buf.data.is_empty() && buf.len != 0 {
                    break;
                }
                let len = buf.data.len();
                if len > ready {
                    let readable = (available_codestream + ready).min(len);
                    section_buffers.push(IoSliceMut::new(&mut buf.data[ready..readable]));
                    available_codestream = available_codestream.saturating_sub(readable - ready);
                    if available_codestream == 0 {
                        break;
                    }
                }
                ready = ready.saturating_sub(len);
            }
            let num = input.read(&mut section_buffers[..])?;
            if num == 0 {
                break;
            }
            self.ready_section_data += num;
            total += num;
            buf.mark_consumed(num as u64);
            self.dequeue_ready_sections();
        }
        Ok(())
    }

    // Returns whether we modified the pixels.
    #[allow(clippy::too_many_arguments)]
    fn process_single_section(
        frame: &mut Frame,
        buf: &[u8],
        is_complete: bool,
        output_buffers: &mut Option<&mut [JxlOutputBuffer<'_>]>,
        output_profile: &JxlColorProfile,
        pixel_format: &JxlPixelFormat,
        do_flush: bool,
        parallel_runner: &mut dyn JxlParallelRunner,
    ) -> Result<bool> {
        let mut br = BitReader::new(buf);
        frame.decode_lf_global(&mut br, !is_complete)?;
        {
            let lf_splitter = frame.lf_image.as_mut().map(LfImageSplitter::new);
            let hf_meta_splitter = frame.hf_meta.as_mut().map(HfMetaSplitter::new);
            Frame::decode_lf_group(
                &frame.header,
                &frame.decoder_state,
                frame.lf_global.as_ref().unwrap(),
                0,
                &mut br,
                lf_splitter.as_ref(),
                hf_meta_splitter.as_ref(),
            )?;
        }
        frame.post_decode_lf_group(0);
        frame.decode_hf_global(&mut br)?;
        frame.finalize_lf(parallel_runner)?;
        frame.decode_and_render_hf_groups(
            output_buffers,
            pixel_format,
            vec![(0, vec![(0, br)])],
            do_flush,
            output_profile,
            parallel_runner,
        )
    }

    // Returns:
    // Ok(None) if the frame is complete
    // Ok(Some(len)) if we need `len` bytes to make progress
    // Err(_) if there was an error.
    pub fn process_sections(
        &mut self,
        mut output_buffers: Option<&mut [JxlOutputBuffer<'_>]>,
        output_profile: &JxlColorProfile,
        pixel_format: &JxlPixelFormat,
        parallel_runner: &mut dyn JxlParallelRunner,
    ) -> Result<Option<usize>> {
        let data_for_next_section = self
            .sections
            .front()
            .map(|x| x.len - self.ready_section_data);

        let frame = self.frame.as_mut().unwrap();
        let frame_header = frame.header();
        if frame_header.num_groups() == 1 && frame_header.passes.num_passes == 1 {
            // Single-group special case.
            let Some(buf) = self.lf_global_section.take() else {
                return Ok(data_for_next_section);
            };
            assert!(self.sections.is_empty());
            self.pixels_dirty |= Self::process_single_section(
                frame,
                &buf.data,
                true,
                &mut output_buffers,
                output_profile,
                pixel_format,
                false,
                parallel_runner,
            )?;
            return Ok(None);
        }

        if let Some(buf) = self.lf_global_section.take() {
            frame.decode_lf_global(&mut BitReader::new(&buf.data), false)?;
            self.section_state.lf_global_done = true;
        }

        if !self.section_state.lf_global_done {
            return Ok(data_for_next_section);
        }

        {
            let lf_splitter = frame.lf_image.as_mut().map(LfImageSplitter::new);
            let hf_meta_splitter = frame.hf_meta.as_mut().map(HfMetaSplitter::new);
            let header = &frame.header;
            let decoder_state = &frame.decoder_state;
            let lf_global = frame.lf_global.as_ref().unwrap();

            parallel_runner.run_ordered(
                self.lf_sections.len(),
                None,
                &|i: usize| -> Result<()> {
                    let lf_section = &self.lf_sections[i];
                    let Section::Lf { group } = &lf_section.section else {
                        unreachable!()
                    };
                    Frame::decode_lf_group(
                        header,
                        decoder_state,
                        lf_global,
                        *group,
                        &mut BitReader::new(&lf_section.data),
                        lf_splitter.as_ref(),
                        hf_meta_splitter.as_ref(),
                    )?;
                    Ok(())
                },
            )?;
        }

        for lf_section in self.lf_sections.drain(..) {
            let Section::Lf { group } = lf_section.section else {
                unreachable!()
            };
            frame.post_decode_lf_group(group);
            self.section_state.remaining_lf -= 1;
        }

        if self.section_state.remaining_lf != 0 {
            return Ok(data_for_next_section);
        }

        if let Some(hf_global) = self.hf_global_section.take() {
            frame.decode_hf_global(&mut BitReader::new(&hf_global.data))?;
            frame.finalize_lf(parallel_runner)?;
            self.section_state.hf_global_done = true;
        }

        if !self.section_state.hf_global_done {
            return Ok(data_for_next_section);
        }

        let mut group_readers = vec![];
        let mut processed_groups = vec![];

        let mut check_group = |g: usize| {
            let mut sections = vec![];
            for (pass, grp) in self.hf_sections[g]
                .iter()
                .enumerate()
                .skip(self.section_state.completed_passes[g] as usize)
            {
                let Some(s) = &grp else {
                    break;
                };
                self.section_state.completed_passes[g] += 1;
                sections.push((pass, BitReader::new(&s.data)));
            }
            if !sections.is_empty() {
                group_readers.push((g, sections));
                processed_groups.push(g);
            }
        };

        if self.candidate_hf_sections.len() * 4 < self.hf_sections.len() {
            for g in self.candidate_hf_sections.drain() {
                check_group(g)
            }
            // Processing sections in order is more efficient because it lets us flush
            // the pipeline faster.
            group_readers.sort_unstable_by_key(|x| x.0);
        } else {
            for g in 0..self.hf_sections.len() {
                if self.candidate_hf_sections.contains(&g) {
                    check_group(g);
                }
            }
            self.candidate_hf_sections.clear();
        }

        self.pixels_dirty |= frame.decode_and_render_hf_groups(
            &mut output_buffers,
            pixel_format,
            group_readers,
            false,
            output_profile,
            parallel_runner,
        )?;

        for g in processed_groups.into_iter() {
            for i in 0..self.section_state.completed_passes[g] {
                self.hf_sections[g][i as usize] = None;
            }
        }

        Ok(data_for_next_section)
    }

    pub fn do_flush(
        &mut self,
        output_buffers: &mut [JxlOutputBuffer<'_>],
        output_profile: &JxlColorProfile,
        pixel_format: &JxlPixelFormat,
        parallel_runner: &mut dyn JxlParallelRunner,
    ) -> Result<()> {
        let Some(frame) = self.frame.as_mut() else {
            return Ok(());
        };
        validate_output_buffers(output_buffers, Some(pixel_format))?;
        let frame_header = frame.header();

        let has_partial_lf = self
            .sections
            .front()
            .is_some_and(|s| s.section == Section::LfGlobal)
            && 2 * self.ready_section_data > 3 * self.section_state.lf_global_flush_len
            && frame_header.encoding == Encoding::Modular
            && matches!(
                frame_header.frame_type,
                FrameType::RegularFrame | FrameType::LFFrame
            );

        if has_partial_lf {
            self.section_state.lf_global_flush_len = self.ready_section_data;
            let section = &self.sections[0].data[..self.ready_section_data];
            // These are partial renders, so we ignore any errors.
            if frame_header.num_groups() == 1 && frame_header.passes.num_passes == 1 {
                if let Ok(dirty) = Self::process_single_section(
                    frame,
                    section,
                    false,
                    &mut Some(output_buffers),
                    output_profile,
                    pixel_format,
                    true,
                    parallel_runner,
                ) {
                    self.pixels_dirty |= dirty;
                    return Ok(());
                }
            } else {
                let _ = frame.decode_lf_global(&mut BitReader::new(section), true);
            }
        }

        if frame.can_do_early_rendering() || self.section_state.hf_global_done {
            self.pixels_dirty |= frame.decode_and_render_hf_groups(
                &mut Some(output_buffers),
                pixel_format,
                vec![],
                true,
                output_profile,
                parallel_runner,
            )?;
        }
        Ok(())
    }
}

impl CodestreamParser {
    pub fn get_and_clear_pixels_dirty(&mut self) -> bool {
        let r = self.frame_info.pixels_dirty;
        self.frame_info.pixels_dirty = false;
        r
    }
}