// 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 jxl_macros::UnconditionalCoder;

use crate::bit_reader::BitReader;
use crate::entropy_coding::ans::*;
use crate::entropy_coding::context_map::*;
use crate::entropy_coding::huffman::*;
use crate::entropy_coding::hybrid_uint::*;
use crate::error::{Error, Result};
use crate::headers::encodings::*;
use crate::util::tracing_wrappers::*;

pub fn decode_varint16(br: &mut BitReader) -> Result<u16> {
    if br.read(1)? != 0 {
        let nbits = br.read(4)? as usize;
        if nbits == 0 {
            Ok(1)
        } else {
            Ok((1 << nbits) + br.read(nbits)? as u16)
        }
    } else {
        Ok(0)
    }
}

pub fn unpack_signed(unsigned: u32) -> i32 {
    ((unsigned >> 1) ^ ((!unsigned) & 1).wrapping_sub(1)) as i32
}

#[derive(UnconditionalCoder, Debug)]
struct Lz77Params {
    pub enabled: bool,
    #[condition(enabled)]
    #[coder(u2S(224, 512, 4096, Bits(15) + 8))]
    pub min_symbol: Option<u32>,
    #[condition(enabled)]
    #[coder(u2S(3, 4, Bits(2) + 5, Bits(8) + 9))]
    pub min_length: Option<u32>,
}

#[derive(Debug)]
enum Codes {
    Huffman(HuffmanCodes),
    Ans(AnsCodes),
}

impl Codes {
    fn single_symbol(&self, ctx: usize) -> Option<u32> {
        match self {
            Self::Huffman(hc) => hc.single_symbol(ctx),
            Self::Ans(ans) => ans.single_symbol(ctx),
        }
    }
}

#[derive(Debug)]
pub struct Histograms {
    lz77_params: Lz77Params,
    lz77_length_uint: Option<HybridUint>,
    context_map: Vec<u8>,
    // TODO(veluca): figure out why this is unused.
    #[allow(dead_code)]
    log_alpha_size: usize,
    uint_configs: Vec<HybridUint>,
    codes: Codes,
}

#[derive(Debug)]
pub struct Lz77State {
    min_symbol: u32,
    min_length: u32,
    dist_multiplier: u32,
    window: Vec<u32>,
    num_to_copy: u32,
    copy_pos: u32,
    num_decoded: u32,
}

impl Lz77State {
    const LOG_WINDOW_SIZE: u32 = 20;
    const WINDOW_MASK: u32 = (1 << Self::LOG_WINDOW_SIZE) - 1;

    #[rustfmt::skip]
    const SPECIAL_DISTANCES: [(i8, u8); 120] = [
        ( 0, 1), ( 1, 0), ( 1, 1), (-1, 1), ( 0, 2), ( 2, 0), ( 1, 2), (-1, 2), ( 2, 1), (-2, 1),
        ( 2, 2), (-2, 2), ( 0, 3), ( 3, 0), ( 1, 3), (-1, 3), ( 3, 1), (-3, 1), ( 2, 3), (-2, 3),
        ( 3, 2), (-3, 2), ( 0, 4), ( 4, 0), ( 1, 4), (-1, 4), ( 4, 1), (-4, 1), ( 3, 3), (-3, 3),
        ( 2, 4), (-2, 4), ( 4, 2), (-4, 2), ( 0, 5), ( 3, 4), (-3, 4), ( 4, 3), (-4, 3), ( 5, 0),
        ( 1, 5), (-1, 5), ( 5, 1), (-5, 1), ( 2, 5), (-2, 5), ( 5, 2), (-5, 2), ( 4, 4), (-4, 4),
        ( 3, 5), (-3, 5), ( 5, 3), (-5, 3), ( 0, 6), ( 6, 0), ( 1, 6), (-1, 6), ( 6, 1), (-6, 1),
        ( 2, 6), (-2, 6), ( 6, 2), (-6, 2), ( 4, 5), (-4, 5), ( 5, 4), (-5, 4), ( 3, 6), (-3, 6),
        ( 6, 3), (-6, 3), ( 0, 7), ( 7, 0), ( 1, 7), (-1, 7), ( 5, 5), (-5, 5), ( 7, 1), (-7, 1),
        ( 4, 6), (-4, 6), ( 6, 4), (-6, 4), ( 2, 7), (-2, 7), ( 7, 2), (-7, 2), ( 3, 7), (-3, 7),
        ( 7, 3), (-7, 3), ( 5, 6), (-5, 6), ( 6, 5), (-6, 5), ( 8, 0), ( 4, 7), (-4, 7), ( 7, 4),
        (-7, 4), ( 8, 1), ( 8, 2), ( 6, 6), (-6, 6), ( 8, 3), ( 5, 7), (-5, 7), ( 7, 5), (-7, 5),
        ( 8, 4), ( 6, 7), (-6, 7), ( 7, 6), (-7, 6), ( 8, 5), ( 7, 7), (-7, 7), ( 8, 6), ( 8, 7),
    ];

    #[inline]
    fn apply_copy(&mut self, distance_sym: u32, num_to_copy: u32) {
        let distance_sub_1 = if self.dist_multiplier == 0 {
            distance_sym
        } else if let Some(distance) = distance_sym.checked_sub(120) {
            distance
        } else {
            let (offset, dist) = Lz77State::SPECIAL_DISTANCES[distance_sym as usize];
            let dist = (self.dist_multiplier * dist as u32).checked_add_signed(offset as i32 - 1);
            dist.unwrap_or(0)
        };

        let distance = (((1 << 20) - 1).min(distance_sub_1) + 1).min(self.num_decoded);
        self.copy_pos = self.num_decoded - distance;
        self.num_to_copy = num_to_copy;
    }

    #[inline]
    fn push_decoded_symbol(&mut self, token: u32) {
        let offset = (self.num_decoded & Self::WINDOW_MASK) as usize;
        if let Some(slot) = self.window.get_mut(offset) {
            *slot = token;
        } else {
            debug_assert_eq!(self.window.len(), offset);
            self.window.push(token);
        }
        self.num_decoded += 1;
    }

    #[inline]
    fn pull_symbol(&mut self) -> Option<u32> {
        if let Some(next_num_to_copy) = self.num_to_copy.checked_sub(1) {
            let sym = self.window[(self.copy_pos & Self::WINDOW_MASK) as usize];
            self.copy_pos += 1;
            self.num_to_copy = next_num_to_copy;
            Some(sym)
        } else {
            None
        }
    }
}

#[derive(Debug)]
struct RleState {
    min_symbol: u32,
    min_length: u32,
    last_sym: Option<u32>,
    repeat_count: u32,
}

impl RleState {
    #[inline]
    fn push_token(
        &mut self,
        token: u32,
        histograms: &Histograms,
        br: &mut BitReader,
        cluster: usize,
    ) {
        if let Some(token) = token.checked_sub(self.min_symbol) {
            let lz_length_conf = histograms.lz77_length_uint.as_ref().unwrap();
            let count = lz_length_conf.read(token, br);
            // If this calculation overflows, the bitstream is invalid (it would be rejected
            // on the LZ77 path), but we don't report an error.
            self.repeat_count = count.wrapping_add(self.min_length);
        } else {
            let sym = histograms.uint_configs[cluster].read(token, br);
            self.last_sym = Some(sym);
            self.repeat_count = 1;
        }
    }

    #[inline]
    fn pull_symbol(&mut self) -> Option<u32> {
        if self.repeat_count > 0 {
            self.repeat_count -= 1;
            self.last_sym
        } else {
            None
        }
    }
}

#[derive(Debug)]
enum SymbolReaderState {
    None,
    Lz77(Lz77State),
    Rle(RleState),
}

#[derive(Debug, Clone, Default)]
struct ErrorState {
    lz77_repeat: bool,
    arithmetic_overflow: bool,
}

impl ErrorState {
    fn new() -> Self {
        Self::default()
    }

    fn check_for_error(&self) -> Result<()> {
        if self.lz77_repeat {
            Err(Error::UnexpectedLz77Repeat)
        } else if self.arithmetic_overflow {
            Err(Error::ArithmeticOverflow)
        } else {
            Ok(())
        }
    }
}

#[derive(Debug)]
pub struct SymbolReader {
    state: SymbolReaderState,
    ans_reader: AnsReader,
    errors: ErrorState,
}

impl SymbolReader {
    pub fn new(
        histograms: &Histograms,
        br: &mut BitReader,
        image_width: Option<usize>,
    ) -> Result<Self> {
        let ans_reader = if matches!(histograms.codes, Codes::Ans(_)) {
            AnsReader::init(br)?
        } else {
            AnsReader::new_unused()
        };

        let Lz77Params {
            enabled: lz77_enabled,
            min_symbol,
            min_length,
        } = histograms.lz77_params;

        let state = if lz77_enabled {
            let min_symbol = min_symbol.unwrap();
            let min_length = min_length.unwrap();
            let dist_multiplier = image_width.unwrap_or(0) as u32;

            let lz_dist_cluster = *histograms.context_map.last().unwrap() as usize;
            let lz_conf = &histograms.uint_configs[lz_dist_cluster];
            let is_rle = histograms.codes.single_symbol(lz_dist_cluster) == Some(1)
                && lz_conf.is_split_exponent_zero();

            if is_rle {
                SymbolReaderState::Rle(RleState {
                    min_symbol,
                    min_length,
                    last_sym: None,
                    repeat_count: 0,
                })
            } else {
                SymbolReaderState::Lz77(Lz77State {
                    min_symbol,
                    min_length,
                    dist_multiplier,
                    window: Vec::new(),
                    num_to_copy: 0,
                    copy_pos: 0,
                    num_decoded: 0,
                })
            }
        } else {
            SymbolReaderState::None
        };

        Ok(Self {
            state,
            ans_reader,
            errors: ErrorState::new(),
        })
    }
}

impl SymbolReader {
    #[inline]
    pub fn read_unsigned(
        &mut self,
        histograms: &Histograms,
        br: &mut BitReader,
        context: usize,
    ) -> u32 {
        let cluster = histograms.map_context_to_cluster(context);
        self.read_unsigned_clustered(histograms, br, cluster)
    }

    #[inline(always)]
    pub fn read_signed(
        &mut self,
        histograms: &Histograms,
        br: &mut BitReader,
        context: usize,
    ) -> i32 {
        let unsigned = self.read_unsigned(histograms, br, context);
        unpack_signed(unsigned)
    }

    #[inline]
    pub fn read_unsigned_clustered(
        &mut self,
        histograms: &Histograms,
        br: &mut BitReader,
        cluster: usize,
    ) -> u32 {
        match &mut self.state {
            SymbolReaderState::None => {
                let token = match &histograms.codes {
                    Codes::Huffman(hc) => hc.read(br, cluster),
                    Codes::Ans(ans) => self.ans_reader.read(ans, br, cluster),
                };
                histograms.uint_configs[cluster].read(token, br)
            }

            SymbolReaderState::Lz77(lz77_state) => {
                if let Some(sym) = lz77_state.pull_symbol() {
                    lz77_state.push_decoded_symbol(sym);
                    return sym;
                }
                let token = match &histograms.codes {
                    Codes::Huffman(hc) => hc.read(br, cluster),
                    Codes::Ans(ans) => self.ans_reader.read(ans, br, cluster),
                };
                let Some(lz77_token) = token.checked_sub(lz77_state.min_symbol) else {
                    let sym = histograms.uint_configs[cluster].read(token, br);
                    lz77_state.push_decoded_symbol(sym);
                    return sym;
                };
                if lz77_state.num_decoded == 0 {
                    self.errors.lz77_repeat = true;
                    return 0;
                }

                let num_to_copy = histograms
                    .lz77_length_uint
                    .as_ref()
                    .unwrap()
                    .read(lz77_token, br);
                let Some(num_to_copy) = num_to_copy.checked_add(lz77_state.min_length) else {
                    warn!(
                        num_to_copy,
                        lz77_state.min_length, "LZ77 num_to_copy overflow"
                    );
                    self.errors.arithmetic_overflow = true;
                    return 0;
                };

                let lz_dist_cluster = *histograms.context_map.last().unwrap() as usize;
                let distance_sym = match &histograms.codes {
                    Codes::Huffman(hc) => hc.read(br, lz_dist_cluster),
                    Codes::Ans(ans) => self.ans_reader.read(ans, br, lz_dist_cluster),
                };
                let distance_sym = histograms.uint_configs[lz_dist_cluster].read(distance_sym, br);
                lz77_state.apply_copy(distance_sym, num_to_copy);

                let sym = lz77_state.pull_symbol().unwrap();
                lz77_state.push_decoded_symbol(sym);
                sym
            }

            SymbolReaderState::Rle(rle_state) => {
                if let Some(sym) = rle_state.pull_symbol() {
                    return sym;
                }

                let token = match &histograms.codes {
                    Codes::Huffman(hc) => hc.read(br, cluster),
                    Codes::Ans(ans) => self.ans_reader.read(ans, br, cluster),
                };
                rle_state.push_token(token, histograms, br, cluster);
                if let Some(sym) = rle_state.pull_symbol() {
                    sym
                } else {
                    self.errors.lz77_repeat = true;
                    0
                }
            }
        }
    }

    #[inline(always)]
    pub fn read_signed_clustered(
        &mut self,
        histograms: &Histograms,
        br: &mut BitReader,
        cluster: usize,
    ) -> i32 {
        let unsigned = self.read_unsigned_clustered(histograms, br, cluster);
        unpack_signed(unsigned)
    }

    pub fn check_final_state(self, histograms: &Histograms, br: &mut BitReader) -> Result<()> {
        self.errors.check_for_error()?;
        br.check_for_error()?;
        match &histograms.codes {
            Codes::Huffman(_) => Ok(()),
            Codes::Ans(_) => self.ans_reader.check_final_state(),
        }
    }

    pub fn checkpoint<const N: usize>(&self) -> Checkpoint<N> {
        let state = match &self.state {
            SymbolReaderState::None => StateCheckpoint::None,
            SymbolReaderState::Lz77(lz77_state) => {
                let mut window = [0u32; N];
                let start = (lz77_state.num_decoded & Lz77State::WINDOW_MASK) as usize;
                let end = ((lz77_state.num_decoded + N as u32) & Lz77State::WINDOW_MASK) as usize;
                if start < end {
                    let window_first = &lz77_state.window[start..];
                    let actual_size = window_first.len().min(N);
                    window[..actual_size].copy_from_slice(&window_first[..actual_size]);
                } else {
                    let window_first = &lz77_state.window[start..];
                    let first_len = window_first
                        .len()
                        .min((1 << Lz77State::LOG_WINDOW_SIZE) - start);
                    window[..first_len].copy_from_slice(&window_first[..first_len]);
                    window[N - end..].copy_from_slice(&lz77_state.window[..end]);
                }
                StateCheckpoint::Lz77 {
                    num_to_copy: lz77_state.num_to_copy,
                    copy_pos: lz77_state.copy_pos,
                    num_decoded: lz77_state.num_decoded,
                    window,
                }
            }
            SymbolReaderState::Rle(rle_state) => StateCheckpoint::Rle {
                last_sym: rle_state.last_sym,
                repeat_count: rle_state.repeat_count,
            },
        };

        Checkpoint {
            state,
            ans_reader: self.ans_reader.checkpoint(),
            errors: self.errors.clone(),
        }
    }

    pub fn restore<const N: usize>(&mut self, checkpoint: Checkpoint<N>) {
        match checkpoint.state {
            StateCheckpoint::None => {
                if !matches!(self.state, SymbolReaderState::None) {
                    panic!("checkpoint type mismatch");
                }
            }
            StateCheckpoint::Lz77 {
                num_to_copy,
                copy_pos,
                num_decoded,
                window,
            } => {
                let SymbolReaderState::Lz77(lz77_state) = &mut self.state else {
                    panic!("checkpoint type mismatch");
                };

                let num_rewind = lz77_state.num_decoded - num_decoded;
                let rewind_window = &window[..num_rewind as usize];

                let start = (num_decoded & Lz77State::WINDOW_MASK) as usize;
                let end = ((num_decoded + num_rewind) & Lz77State::WINDOW_MASK) as usize;
                if start < end {
                    lz77_state.window[start..end].copy_from_slice(rewind_window);
                } else {
                    let window_first = &mut lz77_state.window[start..];
                    let first_len = window_first.len();
                    window_first.copy_from_slice(&rewind_window[..first_len]);
                    lz77_state.window[..end].copy_from_slice(&rewind_window[first_len..]);
                }

                lz77_state.num_to_copy = num_to_copy;
                lz77_state.copy_pos = copy_pos;
                lz77_state.num_decoded = num_decoded;
            }
            StateCheckpoint::Rle {
                last_sym,
                repeat_count,
            } => {
                let SymbolReaderState::Rle(rle_state) = &mut self.state else {
                    panic!("checkpoint type mismatch");
                };

                rle_state.last_sym = last_sym;
                rle_state.repeat_count = repeat_count;
            }
        }

        self.ans_reader = checkpoint.ans_reader;
        self.errors = checkpoint.errors;
    }
}

impl Histograms {
    pub fn decode(num_contexts: usize, br: &mut BitReader, allow_lz77: bool) -> Result<Histograms> {
        let lz77_params = Lz77Params::read_unconditional(&(), br, &Empty {})?;
        if !allow_lz77 && lz77_params.enabled {
            return Err(Error::Lz77Disallowed);
        }
        let (num_contexts, lz77_length_uint) = if lz77_params.enabled {
            (
                num_contexts + 1,
                Some(HybridUint::decode(/*log_alpha_size=*/ 8, br)?),
            )
        } else {
            (num_contexts, None)
        };

        let context_map = if num_contexts > 1 {
            decode_context_map(num_contexts, br)?
        } else {
            vec![0]
        };
        assert_eq!(context_map.len(), num_contexts);

        let use_prefix_code = br.read(1)? != 0;
        let log_alpha_size = if use_prefix_code {
            HUFFMAN_MAX_BITS
        } else {
            br.read(2)? as usize + 5
        };
        let num_histograms = *context_map.iter().max().unwrap() + 1;
        let uint_configs = ((0..num_histograms).map(|_| HybridUint::decode(log_alpha_size, br)))
            .collect::<Result<_>>()?;

        let codes = if use_prefix_code {
            Codes::Huffman(HuffmanCodes::decode(num_histograms as usize, br)?)
        } else {
            Codes::Ans(AnsCodes::decode(
                num_histograms as usize,
                log_alpha_size,
                br,
            )?)
        };

        Ok(Histograms {
            lz77_params,
            lz77_length_uint,
            context_map,
            log_alpha_size,
            uint_configs,
            codes,
        })
    }

    pub fn map_context_to_cluster(&self, context: usize) -> usize {
        self.context_map[context] as usize
    }

    pub fn num_histograms(&self) -> usize {
        *self.context_map.iter().max().unwrap() as usize + 1
    }
}

#[cfg(test)]
impl Histograms {
    /// Builds a decoder that reads an octet at a time and emits its bit-reversed value.
    pub fn reverse_octet(num_contexts: usize) -> Self {
        let d = HuffmanCodes::byte_histogram();
        let codes = Codes::Huffman(d);
        let uint_configs = vec![HybridUint::new(8, 0, 0)];
        Self {
            lz77_params: Lz77Params {
                enabled: false,
                min_symbol: None,
                min_length: None,
            },
            lz77_length_uint: None,
            uint_configs,
            log_alpha_size: 15,
            context_map: vec![0u8; num_contexts],
            codes,
        }
    }

    pub fn rle(num_contexts: usize, min_symbol: u32, min_length: u32) -> Self {
        let d = HuffmanCodes::byte_histogram_rle();
        let codes = Codes::Huffman(d);
        let uint_configs = vec![HybridUint::new(8, 0, 0), HybridUint::new(0, 0, 0)];
        let mut context_map = vec![0u8; num_contexts + 1];
        *context_map.last_mut().unwrap() = 1;
        Self {
            lz77_params: Lz77Params {
                enabled: true,
                min_symbol: Some(min_symbol),
                min_length: Some(min_length),
            },
            lz77_length_uint: Some(HybridUint::new(8, 0, 0)),
            uint_configs,
            log_alpha_size: 15,
            context_map,
            codes,
        }
    }
}

#[derive(Debug)]
enum StateCheckpoint<const N: usize> {
    None,
    Lz77 {
        num_to_copy: u32,
        copy_pos: u32,
        num_decoded: u32,
        window: [u32; N],
    },
    Rle {
        last_sym: Option<u32>,
        repeat_count: u32,
    },
}

#[derive(Debug)]
pub struct Checkpoint<const N: usize> {
    state: StateCheckpoint<N>,
    ans_reader: AnsReader,
    errors: ErrorState,
}

#[cfg(test)]
mod test {
    use std::ops::ControlFlow;

    use test_log::test;

    use super::*;

    #[test]
    fn rle_arb() {
        let histograms = Histograms::rle(1, 240, 3);

        arbtest::arbtest(|u| {
            let width = u.int_in_range(1usize..=256)?;

            let mut bitstream = Vec::new();
            let mut expected_bytes = Vec::new();
            u.arbitrary_loop(None, None, |u| {
                let do_repeat = !expected_bytes.is_empty() && u.ratio(1, 4)?;
                let range = if do_repeat { 240u8..=255 } else { 0u8..=239 };
                let byte = u.int_in_range(range)?;
                bitstream.push(byte);

                if do_repeat {
                    let count = byte as usize - 237;
                    let sym = *expected_bytes.last().unwrap();
                    for _ in 0..count {
                        expected_bytes.push(sym);
                    }
                } else {
                    expected_bytes.push(byte);
                }

                Ok(if expected_bytes.len() >= 256 {
                    ControlFlow::Break(())
                } else {
                    ControlFlow::Continue(())
                })
            })?;
            for b in &mut bitstream {
                *b = b.reverse_bits();
            }

            // Read RLE
            let mut br = BitReader::new(&bitstream);
            let mut reader = SymbolReader::new(&histograms, &mut br, Some(width)).unwrap();

            for expected in expected_bytes {
                let actual = reader.read_unsigned_clustered(&histograms, &mut br, 0);
                assert_eq!(actual, expected as u32);
            }

            let SymbolReaderState::Rle(rle_state) = &reader.state else {
                panic!()
            };
            assert_eq!(rle_state.repeat_count, 0);
            assert!(reader.check_final_state(&histograms, &mut br).is_ok());
            assert_eq!(br.total_bits_available(), 0);

            Ok(())
        });
    }
}