// 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::VecDeque, ops::Range};

use crate::{
    bit_reader::BitReader,
    entropy_coding::decode::{Histograms, SymbolReader},
    error::Result,
    frame::modular::{
        ModularChannel, Predictor, Tree,
        decode::{
            channel::ModularChannelDecoder,
            common::{make_pixel, precompute_references},
        },
        predict::{PredictionData, WeightedPredictorState, clamped_gradient},
        tree::{
            FlatTreeNode, NUM_NONREF_PROPERTIES, PROPERTIES_PER_PREVCHAN, TreeNode, predict_flat,
        },
    },
    headers::modular::GroupHeader,
    image::Image,
};

pub struct NoWpTree {
    flat_nodes: Vec<FlatTreeNode>,
    references: Image<i32>,
    property_buffer: Vec<i32>,
}

impl NoWpTree {
    fn new(
        nodes: Vec<TreeNode>,
        max_property_count: usize,
        channel: usize,
        stream: usize,
        xsize: usize,
    ) -> Result<Self> {
        let num_ref_props = max_property_count
            .saturating_sub(NUM_NONREF_PROPERTIES)
            .next_multiple_of(PROPERTIES_PER_PREVCHAN);
        let references = Image::<i32>::new((num_ref_props, xsize))?;
        let num_properties = NUM_NONREF_PROPERTIES + num_ref_props;
        let mut property_buffer: Vec<i32> = vec![0; num_properties];

        property_buffer[0] = channel as i32;
        property_buffer[1] = stream as i32;

        let flat_nodes = Tree::build_flat_tree(&nodes)?;

        Ok(Self {
            flat_nodes,
            references,
            property_buffer,
        })
    }
}

impl ModularChannelDecoder for NoWpTree {
    const NEEDS_TOP: bool = true;
    const NEEDS_TOPTOP: bool = true;

    fn init_row(&mut self, buffers: &mut [&mut ModularChannel], chan: usize, y: usize) {
        precompute_references(buffers, chan, y, &mut self.references);
        self.property_buffer[2..].fill(0);
    }

    fn decode_one(
        &mut self,
        prediction_data: PredictionData,
        pos: (usize, usize),
        xsize: usize,
        reader: &mut SymbolReader,
        br: &mut BitReader,
        histograms: &Histograms,
    ) -> i32 {
        let prediction_result = predict_flat(
            &self.flat_nodes,
            prediction_data,
            xsize,
            None,
            pos.0,
            pos.1,
            &self.references,
            &mut self.property_buffer,
        );
        let dec = reader.read_signed(histograms, br, prediction_result.context as usize);
        make_pixel(dec, prediction_result.multiplier, prediction_result.guess)
    }
}

pub struct GeneralTree {
    no_wp_tree: NoWpTree,
    wp_state: WeightedPredictorState,
}

impl GeneralTree {
    fn new(
        nodes: Vec<TreeNode>,
        max_property_count: usize,
        header: &GroupHeader,
        channel: usize,
        stream: usize,
        xsize: usize,
    ) -> Result<Self> {
        let wp_state = WeightedPredictorState::new(&header.wp_header, xsize);
        Ok(Self {
            no_wp_tree: NoWpTree::new(nodes, max_property_count, channel, stream, xsize)?,
            wp_state,
        })
    }
}

impl ModularChannelDecoder for GeneralTree {
    const NEEDS_TOP: bool = true;
    const NEEDS_TOPTOP: bool = true;

    fn init_row(&mut self, buffers: &mut [&mut ModularChannel], chan: usize, y: usize) {
        self.no_wp_tree.init_row(buffers, chan, y);
    }

    fn decode_one(
        &mut self,
        prediction_data: PredictionData,
        pos: (usize, usize),
        xsize: usize,
        reader: &mut SymbolReader,
        br: &mut BitReader,
        histograms: &Histograms,
    ) -> i32 {
        let prediction_result = predict_flat(
            &self.no_wp_tree.flat_nodes,
            prediction_data,
            xsize,
            Some(&mut self.wp_state),
            pos.0,
            pos.1,
            &self.no_wp_tree.references,
            &mut self.no_wp_tree.property_buffer,
        );
        let dec = reader.read_signed(histograms, br, prediction_result.context as usize);
        let val = make_pixel(dec, prediction_result.multiplier, prediction_result.guess);
        self.wp_state.update_errors(val, pos, xsize);
        val
    }
}

const LUT_MAX_SPLITVAL: i32 = 1023;
const LUT_MIN_SPLITVAL: i32 = -1024;
const LUT_TABLE_SIZE: usize = (LUT_MAX_SPLITVAL - LUT_MIN_SPLITVAL + 1) as usize;
const _: () = assert!(LUT_TABLE_SIZE.is_power_of_two());

pub struct WpOnlyLookup {
    lut: [u8; LUT_TABLE_SIZE], // Lookup (wp value -> *clustered* context id)
    wp_state: WeightedPredictorState,
}

fn make_lut(tree: &[TreeNode], histograms: &Histograms) -> Option<[u8; LUT_TABLE_SIZE]> {
    struct RangeAndNode {
        range: Range<i32>,
        node: u32,
    }
    let mut stack = vec![RangeAndNode {
        range: LUT_MIN_SPLITVAL..LUT_MAX_SPLITVAL + 1,
        node: 0,
    }];

    let mut ans = [0u8; LUT_TABLE_SIZE];
    while let Some(RangeAndNode { range, node }) = stack.pop() {
        let v = tree[node as usize];
        match v {
            TreeNode::Split {
                val, left, right, ..
            } => {
                let first_left = val + 1;
                if first_left >= range.end || first_left <= range.start {
                    return None;
                }
                stack.push(RangeAndNode {
                    range: first_left..range.end,
                    node: left,
                });
                stack.push(RangeAndNode {
                    range: range.start..first_left,
                    node: right,
                });
            }
            TreeNode::Leaf {
                offset,
                multiplier,
                id,
                ..
            } => {
                if offset != 0 || multiplier != 1 {
                    return None;
                }
                let start = range.start - LUT_MIN_SPLITVAL;
                let end = range.end - LUT_MIN_SPLITVAL;
                ans[start as usize..end as usize]
                    .fill(histograms.map_context_to_cluster(id as usize) as u8);
            }
        }
    }

    Some(ans)
}

impl WpOnlyLookup {
    fn new(
        tree: &[TreeNode],
        histograms: &Histograms,
        header: &GroupHeader,
        xsize: usize,
    ) -> Option<Self> {
        let wp_state = WeightedPredictorState::new(&header.wp_header, xsize);
        let lut = make_lut(tree, histograms)?;
        Some(Self { lut, wp_state })
    }
}

impl ModularChannelDecoder for WpOnlyLookup {
    const NEEDS_TOP: bool = true;
    const NEEDS_TOPTOP: bool = true;

    fn init_row(&mut self, _buffers: &mut [&mut ModularChannel], _chan: usize, _y: usize) {
        // nothing to do
    }

    #[inline(always)]
    fn decode_one(
        &mut self,
        prediction_data: PredictionData,
        pos: (usize, usize),
        xsize: usize,
        reader: &mut SymbolReader,
        br: &mut BitReader,
        histograms: &Histograms,
    ) -> i32 {
        let (wp_pred, property) = self
            .wp_state
            .predict_and_property(pos, xsize, &prediction_data);
        let ctx = self.lut[(property as i64 - LUT_MIN_SPLITVAL as i64)
            .clamp(0, LUT_TABLE_SIZE as i64 - 1) as usize];
        let dec = reader.read_signed_clustered(histograms, br, ctx as usize);
        let val = dec.wrapping_add(wp_pred as i32);
        self.wp_state.update_errors(val, pos, xsize);
        val
    }
}

/// Fast path for trees that split only on property 9 (gradient: left + top - topleft)
/// with Gradient predictor, offset=0, multiplier=1.
/// Maps property 9 values directly to cluster IDs via a LUT.
/// This targets libjxl effort 2 encoding.
pub struct GradientLookup {
    lut: [u8; LUT_TABLE_SIZE],
}

/// Property 9 is the "gradient property": left + top - topleft
const GRADIENT_PROPERTY: u8 = 9;

fn make_gradient_lut(tree: &[TreeNode], histograms: &Histograms) -> Option<GradientLookup> {
    // Verify all splits are on property 9 and all leaves have Gradient predictor
    for node in tree {
        match node {
            TreeNode::Split { property, .. } => {
                if *property != GRADIENT_PROPERTY {
                    return None;
                }
            }
            TreeNode::Leaf { predictor, .. } => {
                if *predictor != Predictor::Gradient {
                    return None;
                }
            }
        }
    }

    // Use existing make_lut which handles offset=0, multiplier=1 checks
    let lut = make_lut(tree, histograms)?;
    Some(GradientLookup { lut })
}

impl ModularChannelDecoder for GradientLookup {
    const NEEDS_TOP: bool = true;
    const NEEDS_TOPTOP: bool = false;

    fn init_row(&mut self, _: &mut [&mut ModularChannel], _: usize, _: usize) {}

    #[inline(always)]
    fn decode_one(
        &mut self,
        prediction_data: PredictionData,
        _: (usize, usize),
        _: usize,
        reader: &mut SymbolReader,
        br: &mut BitReader,
        histograms: &Histograms,
    ) -> i32 {
        let prop9 = prediction_data
            .left
            .wrapping_add(prediction_data.top)
            .wrapping_sub(prediction_data.topleft);

        let index =
            (prop9 as i64 - LUT_MIN_SPLITVAL as i64).clamp(0, LUT_TABLE_SIZE as i64 - 1) as usize;
        let cluster = self.lut[index];

        let pred = clamped_gradient(
            prediction_data.left as i64,
            prediction_data.top as i64,
            prediction_data.topleft as i64,
        );

        let dec = reader.read_signed_clustered(histograms, br, cluster as usize);
        dec.wrapping_add(pred as i32)
    }
}

pub struct SingleGradientOnly {
    ctx: usize,
}

impl ModularChannelDecoder for SingleGradientOnly {
    const NEEDS_TOP: bool = true;
    const NEEDS_TOPTOP: bool = false;

    fn init_row(&mut self, _: &mut [&mut ModularChannel], _: usize, _: usize) {}

    #[inline(always)]
    fn decode_one(
        &mut self,
        prediction_data: PredictionData,
        _: (usize, usize),
        _: usize,
        reader: &mut SymbolReader,
        br: &mut BitReader,
        histograms: &Histograms,
    ) -> i32 {
        let pred = Predictor::Gradient.predict_one(prediction_data, 0);
        let dec = reader.read_signed(histograms, br, self.ctx);
        make_pixel(dec, 1, pred)
    }
}

#[allow(clippy::large_enum_variant)]
pub enum TreeSpecialCase {
    NoWp(NoWpTree),
    WpOnly(WpOnlyLookup),
    GradientLookup(GradientLookup),
    SingleGradientOnly(SingleGradientOnly),
    General(GeneralTree),
}

pub fn specialize_tree(
    tree: &Tree,
    channel: usize,
    stream: usize,
    xsize: usize,
    header: &GroupHeader,
) -> Result<TreeSpecialCase> {
    // TODO(veluca): consider skipping the pruning if header.uses_global_tree is true.
    let mut pruned_tree = Vec::new();
    let mut queue = VecDeque::new();
    pruned_tree.try_reserve(tree.nodes.len())?;
    queue.try_reserve(tree.nodes.len())?;
    queue.push_front(0);

    let mut uses_wp = false;
    let mut uses_non_wp = false;

    // Obtain a pruned tree without nodes that are not relevant in the current channel and stream.
    // Proceed in BFS order, so that we know that the children of  anode will be adjacent.
    while let Some(v) = queue.pop_front() {
        let node = tree.nodes[v as usize];
        match node {
            TreeNode::Split {
                property,
                val,
                left,
                right,
            } if property < 2 => {
                // If the node splits on static properties, re-enqueue its correct child immediately.
                let vv = if property == 0 { channel } else { stream };
                queue.push_front(if vv as i32 > val { left } else { right });
                continue;
            }
            TreeNode::Split {
                property,
                val,
                left,
                right,
            } => {
                // WeightedPredictor property.
                uses_wp |= property == 15;
                uses_non_wp |= property != 15;
                let base = (queue.len() + pruned_tree.len() + 1) as u32;
                pruned_tree.push(TreeNode::Split {
                    property,
                    val,
                    left: base,
                    right: base + 1,
                });
                queue.push_back(left);
                queue.push_back(right);
            }
            TreeNode::Leaf { predictor, .. } => {
                uses_wp |= predictor == Predictor::Weighted;
                uses_non_wp |= predictor != Predictor::Weighted;
                pruned_tree.push(node);
            }
        }
    }

    if let [
        TreeNode::Leaf {
            predictor: Predictor::Gradient,
            multiplier: 1,
            offset: 0,
            id,
        },
    ] = &*pruned_tree
    {
        return Ok(TreeSpecialCase::SingleGradientOnly(SingleGradientOnly {
            ctx: *id as usize,
        }));
    }

    if !uses_non_wp
        && let Some(wp) = WpOnlyLookup::new(&pruned_tree, &tree.histograms, header, xsize)
    {
        return Ok(TreeSpecialCase::WpOnly(wp));
    }

    // Try gradient LUT for non-WP trees (targets effort 2 encoding)
    if !uses_wp {
        if let Some(gl) = make_gradient_lut(&pruned_tree, &tree.histograms) {
            return Ok(TreeSpecialCase::GradientLookup(gl));
        }
        return Ok(TreeSpecialCase::NoWp(NoWpTree::new(
            pruned_tree,
            tree.max_property_count(),
            channel,
            stream,
            xsize,
        )?));
    }

    Ok(TreeSpecialCase::General(GeneralTree::new(
        pruned_tree,
        tree.max_property_count(),
        header,
        channel,
        stream,
        xsize,
    )?))
}