//! PDF Number Tree helper (PDF spec section 7.9.7)
//!
//! Number trees are like Name Trees but use integers as keys. They have:
//! - **Nums**: array of alternating integer-value pairs (leaf nodes)
//! - **Kids**: array of references to child nodes (internal nodes)
//! - **Limits**: array of [min_num, max_num] describing key range

use std::collections::HashMap;

use crate::pdf::object::{Array, Dict, Name, ObjRef, Object};

/// Helper for reading PDF Number Trees.
///
/// Wraps a dictionary that represents a Number Tree node (leaf or internal).
pub struct NumberTreeHelper<'a> {
    dict: &'a Dict,
}

impl<'a> NumberTreeHelper<'a> {
    /// Create a helper for the given Number Tree dictionary.
    pub fn new(dict: &'a Dict) -> Self {
        Self { dict }
    }

    /// Look up a value by numeric key.
    ///
    /// For leaf nodes, searches the Nums array. For internal nodes,
    /// finds the child whose Limits contain the key and recurses.
    pub fn get(&self, key: i64, objects: &HashMap<ObjRef, Object>) -> Option<Object> {
        if let Some(nums) = self.dict.get(&Name::new("Nums")).and_then(Object::as_array) {
            Self::search_nums(nums, key)
        } else if let Some(kids) = self.dict.get(&Name::new("Kids")).and_then(Object::as_array) {
            for kid in kids {
                let kid_ref = kid.as_obj_ref()?;
                let kid_obj = objects.get(&kid_ref)?;
                let kid_dict = kid_obj.as_dict()?;
                let limits = NumberTreeHelper::new(kid_dict).limits()?;
                if key >= limits.0 && key <= limits.1 {
                    return NumberTreeHelper::new(kid_dict).get(key, objects);
                }
            }
            None
        } else {
            None
        }
    }

    fn search_nums(nums: &Array, key: i64) -> Option<Object> {
        let mut i = 0;
        while i + 1 < nums.len() {
            let k = nums[i].as_int()?;
            if k == key {
                return Some(nums[i + 1].clone());
            }
            if k > key {
                break;
            }
            i += 2;
        }
        None
    }

    /// Get all key-value pairs from the tree (in key order).
    pub fn get_all_entries(&self, objects: &HashMap<ObjRef, Object>) -> Vec<(i64, Object)> {
        let mut entries = Vec::new();
        self.collect_entries(objects, &mut entries);
        entries.sort_by_key(|(k, _)| *k);
        entries
    }

    fn collect_entries(&self, objects: &HashMap<ObjRef, Object>, out: &mut Vec<(i64, Object)>) {
        if let Some(nums) = self.dict.get(&Name::new("Nums")).and_then(Object::as_array) {
            let mut i = 0;
            while i + 1 < nums.len() {
                if let (Some(k), Some(v)) = (nums[i].as_int(), nums.get(i + 1)) {
                    out.push((k, v.clone()));
                }
                i += 2;
            }
        } else if let Some(kids) = self.dict.get(&Name::new("Kids")).and_then(Object::as_array) {
            for kid in kids {
                if let Some(kid_ref) = kid.as_obj_ref() {
                    if let Some(kid_obj) = objects.get(&kid_ref) {
                        if let Some(kid_dict) = kid_obj.as_dict() {
                            NumberTreeHelper::new(kid_dict).collect_entries(objects, out);
                        }
                    }
                }
            }
        }
    }

    /// Check whether the tree contains the given key.
    pub fn has_key(&self, key: i64, objects: &HashMap<ObjRef, Object>) -> bool {
        self.get(key, objects).is_some()
    }

    /// Get the Limits array as (min, max) if present.
    pub fn limits(&self) -> Option<(i64, i64)> {
        let arr = self
            .dict
            .get(&Name::new("Limits"))
            .and_then(Object::as_array)?;
        if arr.len() >= 2 {
            let min = arr[0].as_int()?;
            let max = arr[1].as_int()?;
            Some((min, max))
        } else {
            None
        }
    }
}

/// Helper for mutating PDF Number Trees (flat Nums arrays only).
///
/// Insert and remove operate on the Nums array directly. For trees with Kids,
/// use the appropriate leaf node's helper.
pub struct NumberTreeHelperMut<'a> {
    dict: &'a mut Dict,
}

impl<'a> NumberTreeHelperMut<'a> {
    /// Create a mutable helper for the given Number Tree dictionary.
    pub fn new(dict: &'a mut Dict) -> Self {
        Self { dict }
    }

    /// Insert a key-value pair into the Nums array, maintaining sorted order.
    ///
    /// If the key already exists, the value is replaced.
    pub fn insert(&mut self, key: i64, value: Object) {
        let nums_key = Name::new("Nums");
        let mut nums = self
            .dict
            .get(&nums_key)
            .and_then(Object::as_array)
            .cloned()
            .unwrap_or_default();

        let mut i = 0;
        let mut found = false;
        while i < nums.len() {
            if let Some(k) = nums[i].as_int() {
                if k == key {
                    nums[i + 1] = value.clone();
                    found = true;
                    break;
                }
                if k > key {
                    nums.insert(i, Object::Int(key));
                    nums.insert(i + 1, value.clone());
                    found = true;
                    break;
                }
            }
            i += 2;
        }
        if !found {
            nums.push(Object::Int(key));
            nums.push(value);
        }
        self.dict.insert(nums_key, Object::Array(nums));
    }

    /// Remove a key from the Nums array. Returns the removed value if present.
    pub fn remove(&mut self, key: i64) -> Option<Object> {
        let nums_key = Name::new("Nums");
        let nums = self
            .dict
            .get_mut(&nums_key)
            .and_then(Object::as_array_mut)?;

        let mut i = 0;
        while i + 1 < nums.len() {
            if nums[i].as_int() == Some(key) {
                let removed = nums.remove(i + 1);
                nums.remove(i);
                return Some(removed);
            }
            i += 2;
        }
        None
    }
}

/// Build a flat Number Tree dictionary from a list of (key, value) pairs.
///
/// Entries are sorted by key. The result is a Dict with Nums and Limits.
pub fn build_number_tree(entries: Vec<(i64, Object)>) -> Object {
    let mut sorted = entries;
    sorted.sort_by_key(|(k, _)| *k);

    let first_key = sorted.first().map(|(k, _)| *k);
    let last_key = sorted.last().map(|(k, _)| *k);

    let mut nums = Array::new();
    for (k, v) in sorted {
        nums.push(Object::Int(k));
        nums.push(v);
    }

    let mut dict = Dict::new();
    dict.insert(Name::new("Nums"), Object::Array(nums));

    if let (Some(first), Some(last)) = (first_key, last_key) {
        dict.insert(
            Name::new("Limits"),
            Object::Array(vec![Object::Int(first), Object::Int(last)]),
        );
    }

    Object::Dict(dict)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_flat_number_tree() -> Dict {
        let mut dict = Dict::new();
        dict.insert(
            Name::new("Nums"),
            Object::Array(vec![
                Object::Int(10),
                Object::Int(100),
                Object::Int(20),
                Object::Int(200),
                Object::Int(30),
                Object::Int(300),
            ]),
        );
        dict.insert(
            Name::new("Limits"),
            Object::Array(vec![Object::Int(10), Object::Int(30)]),
        );
        dict
    }

    #[test]
    fn test_number_tree_helper_new() {
        let dict = make_flat_number_tree();
        let helper = NumberTreeHelper::new(&dict);
        assert_eq!(helper.limits(), Some((10, 30)));
    }

    #[test]
    fn test_number_tree_helper_get_flat() {
        let dict = make_flat_number_tree();
        let helper = NumberTreeHelper::new(&dict);
        let objects = HashMap::new();

        assert_eq!(helper.get(10, &objects).and_then(|o| o.as_int()), Some(100));
        assert_eq!(helper.get(20, &objects).and_then(|o| o.as_int()), Some(200));
        assert_eq!(helper.get(30, &objects).and_then(|o| o.as_int()), Some(300));
        assert_eq!(helper.get(15, &objects), None);
        assert_eq!(helper.get(5, &objects), None);
        assert_eq!(helper.get(40, &objects), None);
    }

    #[test]
    fn test_number_tree_helper_get_all_entries_flat() {
        let dict = make_flat_number_tree();
        let helper = NumberTreeHelper::new(&dict);
        let objects = HashMap::new();

        let entries = helper.get_all_entries(&objects);
        assert_eq!(entries.len(), 3);
        assert_eq!(entries[0], (10, Object::Int(100)));
        assert_eq!(entries[1], (20, Object::Int(200)));
        assert_eq!(entries[2], (30, Object::Int(300)));
    }

    #[test]
    fn test_number_tree_helper_has_key() {
        let dict = make_flat_number_tree();
        let helper = NumberTreeHelper::new(&dict);
        let objects = HashMap::new();

        assert!(helper.has_key(10, &objects));
        assert!(helper.has_key(20, &objects));
        assert!(helper.has_key(30, &objects));
        assert!(!helper.has_key(15, &objects));
        assert!(!helper.has_key(0, &objects));
    }

    #[test]
    fn test_number_tree_helper_limits() {
        let dict = make_flat_number_tree();
        let helper = NumberTreeHelper::new(&dict);
        assert_eq!(helper.limits(), Some((10, 30)));
    }

    #[test]
    fn test_number_tree_helper_limits_empty() {
        let mut dict = Dict::new();
        dict.insert(Name::new("Nums"), Object::Array(vec![]));
        let helper = NumberTreeHelper::new(&dict);
        assert_eq!(helper.limits(), None);
    }

    #[test]
    fn test_number_tree_helper_mut_insert_new() {
        let mut dict = Dict::new();
        dict.insert(Name::new("Nums"), Object::Array(vec![]));

        let mut helper = NumberTreeHelperMut::new(&mut dict);
        helper.insert(5, Object::Int(50));
        helper.insert(1, Object::Int(10));
        helper.insert(3, Object::Int(30));

        let nums = dict
            .get(&Name::new("Nums"))
            .and_then(Object::as_array)
            .unwrap();
        assert_eq!(nums.len(), 6);
        assert_eq!(nums[0].as_int(), Some(1));
        assert_eq!(nums[1].as_int(), Some(10));
        assert_eq!(nums[2].as_int(), Some(3));
        assert_eq!(nums[3].as_int(), Some(30));
        assert_eq!(nums[4].as_int(), Some(5));
        assert_eq!(nums[5].as_int(), Some(50));
    }

    #[test]
    fn test_number_tree_helper_mut_insert_replace() {
        let mut dict = make_flat_number_tree();
        let mut helper = NumberTreeHelperMut::new(&mut dict);

        helper.insert(20, Object::Int(999));

        let nums = dict
            .get(&Name::new("Nums"))
            .and_then(Object::as_array)
            .unwrap();
        let mut i = 0;
        while i < nums.len() {
            if nums[i].as_int() == Some(20) {
                assert_eq!(nums[i + 1].as_int(), Some(999));
                break;
            }
            i += 2;
        }
    }

    #[test]
    fn test_number_tree_helper_mut_remove() {
        let mut dict = make_flat_number_tree();
        let mut helper = NumberTreeHelperMut::new(&mut dict);

        let removed = helper.remove(20);
        assert_eq!(removed, Some(Object::Int(200)));

        let nums = dict
            .get(&Name::new("Nums"))
            .and_then(Object::as_array)
            .unwrap();
        assert_eq!(nums.len(), 4);
        assert_eq!(nums[0].as_int(), Some(10));
        assert_eq!(nums[2].as_int(), Some(30));
    }

    #[test]
    fn test_number_tree_helper_mut_remove_missing() {
        let mut dict = make_flat_number_tree();
        let mut helper = NumberTreeHelperMut::new(&mut dict);

        let removed = helper.remove(99);
        assert_eq!(removed, None);

        let nums = dict
            .get(&Name::new("Nums"))
            .and_then(Object::as_array)
            .unwrap();
        assert_eq!(nums.len(), 6);
    }

    #[test]
    fn test_number_tree_helper_mut_remove_from_empty() {
        let mut dict = Dict::new();
        dict.insert(Name::new("Nums"), Object::Array(vec![]));
        let mut helper = NumberTreeHelperMut::new(&mut dict);
        assert_eq!(helper.remove(1), None);
    }

    #[test]
    fn test_build_number_tree() {
        let entries = vec![
            (30, Object::Int(300)),
            (10, Object::Int(100)),
            (20, Object::Int(200)),
        ];
        let obj = build_number_tree(entries);

        let dict = obj.as_dict().unwrap();
        let nums = dict
            .get(&Name::new("Nums"))
            .and_then(Object::as_array)
            .unwrap();
        assert_eq!(nums[0].as_int(), Some(10));
        assert_eq!(nums[1].as_int(), Some(100));
        assert_eq!(nums[2].as_int(), Some(20));
        assert_eq!(nums[3].as_int(), Some(200));
        assert_eq!(nums[4].as_int(), Some(30));
        assert_eq!(nums[5].as_int(), Some(300));

        let limits = dict
            .get(&Name::new("Limits"))
            .and_then(Object::as_array)
            .unwrap();
        assert_eq!(limits[0].as_int(), Some(10));
        assert_eq!(limits[1].as_int(), Some(30));
    }

    #[test]
    fn test_build_number_tree_empty() {
        let obj = build_number_tree(vec![]);
        let dict = obj.as_dict().unwrap();
        assert!(dict.get(&Name::new("Nums")).is_some());
        assert!(dict.get(&Name::new("Limits")).is_none());
    }

    #[test]
    fn test_build_number_tree_single() {
        let obj = build_number_tree(vec![(42, Object::Bool(true))]);
        let dict = obj.as_dict().unwrap();
        let limits = dict
            .get(&Name::new("Limits"))
            .and_then(Object::as_array)
            .unwrap();
        assert_eq!(limits[0].as_int(), Some(42));
        assert_eq!(limits[1].as_int(), Some(42));
    }

    #[test]
    fn test_number_tree_with_kids() {
        let mut leaf = Dict::new();
        leaf.insert(
            Name::new("Nums"),
            Object::Array(vec![
                Object::Int(100),
                Object::Int(1000),
                Object::Int(200),
                Object::Int(2000),
            ]),
        );
        leaf.insert(
            Name::new("Limits"),
            Object::Array(vec![Object::Int(100), Object::Int(200)]),
        );

        let leaf_ref = ObjRef::new(10, 0);
        let mut objects = HashMap::new();
        objects.insert(leaf_ref, Object::Dict(leaf.clone()));

        let mut root = Dict::new();
        root.insert(
            Name::new("Kids"),
            Object::Array(vec![Object::Ref(leaf_ref)]),
        );
        root.insert(
            Name::new("Limits"),
            Object::Array(vec![Object::Int(100), Object::Int(200)]),
        );

        let helper = NumberTreeHelper::new(&root);
        assert_eq!(
            helper.get(100, &objects).and_then(|o| o.as_int()),
            Some(1000)
        );
        assert_eq!(
            helper.get(200, &objects).and_then(|o| o.as_int()),
            Some(2000)
        );
        assert_eq!(helper.get(150, &objects), None);

        let entries = helper.get_all_entries(&objects);
        assert_eq!(entries.len(), 2);
        assert_eq!(entries[0], (100, Object::Int(1000)));
        assert_eq!(entries[1], (200, Object::Int(2000)));
    }
}
