use std::borrow::Cow;
use std::collections::BTreeMap;
use std::fmt;
use std::ops::Index;

use compact_str::CompactString;
use serde::de::{MapAccess, Visitor};
use serde::ser::SerializeMap;
use serde::{Deserialize, Deserializer, Serialize, Serializer};

/// Sorted numeric fields stored contiguously, without a tree node per snapshot.
#[derive(Clone, Default, PartialEq)]
pub struct NumericFactorMap {
    entries: Vec<(CompactString, f64)>,
}

fn compact_key(key: Cow<'static, str>) -> CompactString {
    match key {
        Cow::Borrowed(value) => CompactString::const_new(value),
        Cow::Owned(value) => CompactString::from(value),
    }
}

impl NumericFactorMap {
    pub const fn new() -> Self {
        Self {
            entries: Vec::new(),
        }
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }
    pub fn clear(&mut self) {
        self.entries.clear();
    }

    /// Reserve known new fields without geometric spare capacity per snapshot.
    pub fn reserve_exact(&mut self, additional: usize) {
        self.entries.reserve_exact(additional);
    }

    pub fn get(&self, key: &str) -> Option<&f64> {
        self.entries
            .binary_search_by(|(name, _)| name.as_str().cmp(key))
            .ok()
            .map(|index| &self.entries[index].1)
    }

    pub fn get_mut(&mut self, key: &str) -> Option<&mut f64> {
        self.entries
            .binary_search_by(|(name, _)| name.as_str().cmp(key))
            .ok()
            .map(|index| &mut self.entries[index].1)
    }

    pub fn contains_key(&self, key: &str) -> bool {
        self.get(key).is_some()
    }

    pub fn insert(&mut self, key: Cow<'static, str>, value: f64) -> Option<f64> {
        self.insert_compact(compact_key(key), value)
    }

    pub fn insert_compact(&mut self, key: CompactString, value: f64) -> Option<f64> {
        if self
            .entries
            .last()
            .is_none_or(|(last, _)| last.as_str() < key.as_str())
        {
            self.entries.push((key, value));
            return None;
        }
        match self
            .entries
            .binary_search_by(|(name, _)| name.as_str().cmp(key.as_str()))
        {
            Ok(index) => Some(std::mem::replace(&mut self.entries[index].1, value)),
            Err(index) => {
                self.entries.insert(index, (key, value));
                None
            }
        }
    }

    pub fn remove(&mut self, key: &str) -> Option<f64> {
        self.entries
            .binary_search_by(|(name, _)| name.as_str().cmp(key))
            .ok()
            .map(|index| self.entries.remove(index).1)
    }

    pub fn retain(&mut self, mut keep: impl FnMut(&CompactString, &mut f64) -> bool) {
        self.entries.retain_mut(|(key, value)| keep(key, value));
    }

    pub fn iter(&self) -> Iter<'_> {
        Iter(self.entries.iter())
    }
    pub fn keys(&self) -> impl DoubleEndedIterator<Item = &CompactString> + ExactSizeIterator {
        self.entries.iter().map(|(key, _)| key)
    }
    pub fn values(&self) -> impl DoubleEndedIterator<Item = &f64> + ExactSizeIterator {
        self.entries.iter().map(|(_, value)| value)
    }
}

impl fmt::Debug for NumericFactorMap {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_map().entries(self).finish()
    }
}

impl Index<&str> for NumericFactorMap {
    type Output = f64;
    fn index(&self, key: &str) -> &Self::Output {
        self.get(key).expect("numeric factor key not found")
    }
}

pub struct Iter<'a>(std::slice::Iter<'a, (CompactString, f64)>);
impl<'a> Iterator for Iter<'a> {
    type Item = (&'a CompactString, &'a f64);
    fn next(&mut self) -> Option<Self::Item> {
        self.0.next().map(|(k, v)| (k, v))
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.0.size_hint()
    }
}
impl DoubleEndedIterator for Iter<'_> {
    fn next_back(&mut self) -> Option<Self::Item> {
        self.0.next_back().map(|(k, v)| (k, v))
    }
}
impl ExactSizeIterator for Iter<'_> {}
impl<'a> IntoIterator for &'a NumericFactorMap {
    type Item = (&'a CompactString, &'a f64);
    type IntoIter = Iter<'a>;
    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}
impl IntoIterator for NumericFactorMap {
    type Item = (CompactString, f64);
    type IntoIter = std::vec::IntoIter<Self::Item>;
    fn into_iter(self) -> Self::IntoIter {
        self.entries.into_iter()
    }
}

impl FromIterator<(Cow<'static, str>, f64)> for NumericFactorMap {
    fn from_iter<T: IntoIterator<Item = (Cow<'static, str>, f64)>>(iter: T) -> Self {
        iter.into_iter().map(|(key, value)| (compact_key(key), value)).collect()
    }
}
impl FromIterator<(CompactString, f64)> for NumericFactorMap {
    fn from_iter<T: IntoIterator<Item = (CompactString, f64)>>(iter: T) -> Self {
        let mut entries: Vec<_> = iter.into_iter().collect();
        // Stable sorting preserves last-value-wins for repeated input keys.
        if !entries.windows(2).all(|pair| pair[0].0 <= pair[1].0) {
            entries.sort_by(|left, right| left.0.cmp(&right.0));
        }
        entries.dedup_by(|later, earlier| {
            if later.0 == earlier.0 {
                earlier.1 = later.1;
                true
            } else {
                false
            }
        });
        Self { entries }
    }
}
impl Extend<(Cow<'static, str>, f64)> for NumericFactorMap {
    fn extend<T: IntoIterator<Item = (Cow<'static, str>, f64)>>(&mut self, iter: T) {
        self.extend(iter.into_iter().map(|(key, value)| (compact_key(key), value)));
    }
}
impl Extend<(CompactString, f64)> for NumericFactorMap {
    fn extend<T: IntoIterator<Item = (CompactString, f64)>>(&mut self, iter: T) {
        let mut incoming: Self = iter.into_iter().collect();
        if incoming.is_empty() {
            return;
        }
        if self.is_empty() {
            *self = incoming;
            return;
        }
        if self.entries.last().unwrap().0 < incoming.entries[0].0 {
            self.entries.append(&mut incoming.entries);
            return;
        }
        // Merge sorted sets in linear time; wide factor batches must not shift
        // the existing vector once per field. Existing keys keep their identity.
        let mut merged = Vec::with_capacity(self.len() + incoming.len());
        let mut old = std::mem::take(&mut self.entries).into_iter().peekable();
        let mut new = incoming.entries.into_iter().peekable();
        while let (Some(left), Some(right)) = (old.peek(), new.peek()) {
            match left.0.cmp(&right.0) {
                std::cmp::Ordering::Less => merged.push(old.next().unwrap()),
                std::cmp::Ordering::Greater => merged.push(new.next().unwrap()),
                std::cmp::Ordering::Equal => {
                    let (key, _) = old.next().unwrap();
                    merged.push((key, new.next().unwrap().1));
                }
            }
        }
        merged.extend(old);
        merged.extend(new);
        self.entries = merged;
    }
}
impl<const N: usize> From<[(Cow<'static, str>, f64); N]> for NumericFactorMap {
    fn from(entries: [(Cow<'static, str>, f64); N]) -> Self {
        entries.into_iter().collect()
    }
}
impl From<BTreeMap<Cow<'static, str>, f64>> for NumericFactorMap {
    fn from(entries: BTreeMap<Cow<'static, str>, f64>) -> Self {
        entries.into_iter().collect()
    }
}

impl Serialize for NumericFactorMap {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let mut map = serializer.serialize_map(Some(self.len()))?;
        for (key, value) in self {
            map.serialize_entry(key, value)?;
        }
        map.end()
    }
}
impl<'de> Deserialize<'de> for NumericFactorMap {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        struct Fields;
        impl<'de> Visitor<'de> for Fields {
            type Value = NumericFactorMap;
            fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                f.write_str("a numeric factor map")
            }
            fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Self::Value, A::Error> {
                let mut entries = Vec::new();
                while let Some((key, value)) = map.next_entry::<CompactString, f64>()? {
                    entries.push((key, value));
                }
                Ok(entries.into_iter().collect())
            }
        }
        deserializer.deserialize_map(Fields)
    }
}

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

    #[test]
    fn exact_reservation_preserves_values_and_avoids_growth_during_known_inserts() {
        let mut map = NumericFactorMap::from([
            (Cow::Borrowed("amount"), 125.25),
            (Cow::Borrowed("nullable_value"), f64::from_bits(0x7ff8000000000042)),
            (Cow::Borrowed("signal"), -0.0),
        ]);
        let original = map.iter().map(|(key, value)| (key.to_string(), value.to_bits())).collect::<Vec<_>>();
        map.reserve_exact(2);
        assert_eq!(map.iter().map(|(key, value)| (key.to_string(), value.to_bits())).collect::<Vec<_>>(), original);
        let buffer = map.entries.as_ptr();
        map.insert(Cow::Borrowed("pre_close"), 12.5);
        map.insert(Cow::Borrowed("no_limit"), 0.0);
        assert_eq!(map.entries.as_ptr(), buffer);
        assert_eq!(map.len(), 5);
        assert_eq!(map["signal"].to_bits(), (-0.0_f64).to_bits());
        assert_eq!(map["nullable_value"].to_bits(), 0x7ff8000000000042);
        let before = map.entries.as_ptr();
        map.reserve_exact(0);
        assert_eq!(map.entries.as_ptr(), before);
    }

    #[test]
    fn compact_keys_inline_dynamic_names_and_keep_long_static_storage() {
        const LONG: &str = "a_long_static_factor_identifier_that_must_remain_borrowed";
        let map = NumericFactorMap::from([
            (Cow::Owned("dynamic_factor_20".to_owned()), -0.0),
            (Cow::Borrowed(LONG), 1.0),
        ]);
        let cloned = map.clone();
        let short = cloned.keys().find(|key| key.as_str() == "dynamic_factor_20").unwrap();
        assert!(!short.is_heap_allocated());
        let long = cloned.keys().find(|key| key.as_str() == LONG).unwrap();
        assert_eq!(long.as_static_str(), Some(LONG));
        assert_eq!(cloned["dynamic_factor_20"].to_bits(), (-0.0_f64).to_bits());
        assert_eq!(std::mem::size_of::<(CompactString, f64)>(), std::mem::size_of::<(Cow<'static, str>, f64)>());
    }

    #[test]
    fn long_dynamic_unicode_and_short_keys_keep_the_same_json_map() {
        let entries = ["", "a", "a_field_longer_than_the_inline_string_capacity", "价格因子", "ths_up_days_stock"]
            .into_iter().enumerate().map(|(index, key)| (Cow::Owned(key.to_string()), index as f64 + 0.25))
            .collect::<BTreeMap<_, _>>();
        let map = NumericFactorMap::from(entries.clone());
        assert_eq!(serde_json::to_string(&map).unwrap(), serde_json::to_string(&entries).unwrap());
        let decoded: NumericFactorMap = serde_json::from_str(&serde_json::to_string(&map).unwrap()).unwrap();
        assert_eq!(decoded, map);
        assert!(!decoded.keys().find(|key| key.as_str() == "ths_up_days_stock").unwrap().is_heap_allocated());
    }

    #[test]
    fn updates_order_removal_and_values_match_tree_map() {
        let mut flat = NumericFactorMap::new();
        let mut tree = BTreeMap::new();
        let mut seed = 71_u64;
        for index in 0..10000 {
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            let key: Cow<'static, str> = Cow::Owned(format!("factor_{:04}", (seed >> 32) % 1000));
            let value = (index as f64 - 5000.0) / 7.0;
            if index % 11 == 0 {
                assert_eq!(flat.remove(key.as_ref()), tree.remove(key.as_ref()));
            } else {
                assert_eq!(flat.insert(key.clone(), value), tree.insert(key, value));
            }
        }
        assert_eq!(
            flat.iter().map(|(key, value)| (key.as_str(), value)).collect::<Vec<_>>(),
            tree.iter().map(|(key, value)| (key.as_ref(), value)).collect::<Vec<_>>()
        );
        flat.retain(|_, value| *value > 100.0);
        tree.retain(|_, value| *value > 100.0);
        assert_eq!(
            flat.iter().map(|(key, value)| (key.as_str(), value)).collect::<Vec<_>>(),
            tree.iter().map(|(key, value)| (key.as_ref(), value)).collect::<Vec<_>>()
        );
        assert_eq!(
            std::mem::size_of::<NumericFactorMap>(),
            std::mem::size_of_val(&tree)
        );
    }

    #[test]
    fn bulk_load_is_sorted_and_keeps_last_value_for_each_field() {
        let input = vec![
            (Cow::Borrowed("z"), 2.0),
            (Cow::Borrowed("a"), 1.0),
            (Cow::Borrowed("z"), 3.0),
            (Cow::Borrowed("z"), 4.0),
        ];
        let flat: NumericFactorMap = input.clone().into_iter().collect();
        let tree: BTreeMap<_, _> = input.into_iter().collect();
        assert_eq!(
            flat.iter().map(|(key, value)| (key.as_str(), value)).collect::<Vec<_>>(),
            tree.iter().map(|(key, value)| (key.as_ref(), value)).collect::<Vec<_>>()
        );
        assert_eq!(flat["z"], 4.0);
    }

    #[test]
    fn serialization_keeps_the_map_contract_and_precise_numbers() {
        let input = [
            (Cow::Borrowed("zero"), -0.0),
            (Cow::Borrowed("tiny"), 1.0000000000000002),
            (Cow::Borrowed("large"), 9.123456789123456e20),
        ];
        let flat = NumericFactorMap::from(input.clone());
        let tree = BTreeMap::from(input);
        let json = serde_json::to_string(&flat).unwrap();
        assert_eq!(json, serde_json::to_string(&tree).unwrap());
        let decoded: NumericFactorMap = serde_json::from_str(&json).unwrap();
        for (key, value) in &flat {
            assert_eq!(value.to_bits(), decoded[key.as_ref()].to_bits());
        }
        let duplicate: NumericFactorMap = serde_json::from_str(r#"{"x":1,"x":2}"#).unwrap();
        assert_eq!(duplicate["x"], 2.0);
    }

    #[test]
    fn clone_does_not_share_mutable_values() {
        let original = NumericFactorMap::from([(Cow::Borrowed("signal"), 1.0)]);
        let mut copy = original.clone();
        *copy.get_mut("signal").unwrap() = 0.0;
        copy.insert(Cow::Borrowed("other"), 2.0);
        assert_eq!(original["signal"], 1.0);
        assert!(!original.contains_key("other"));
    }

    #[test]
    fn wide_batch_merge_matches_tree_and_preserves_old_key_ownership() {
        let entries = (0..4096)
            .map(|index| (Cow::Owned(format!("f_{index:05}")), index as f64))
            .collect::<Vec<_>>();
        let mut flat: NumericFactorMap = entries.clone().into_iter().collect();
        let mut tree = BTreeMap::from_iter(entries);
        flat.insert(Cow::Borrowed("shared"), -0.0);
        tree.insert(Cow::Borrowed("shared"), -0.0);
        let incoming = (2048..8192)
            .rev()
            .map(|index| (Cow::Owned(format!("f_{index:05}")), -(index as f64)))
            .chain([(Cow::Owned("shared".to_owned()), 1.0)])
            .collect::<Vec<_>>();
        flat.extend(incoming.clone());
        tree.extend(incoming);
        assert_eq!(
            flat.iter().map(|(key, value)| (key.as_str(), value)).collect::<Vec<_>>(),
            tree.iter().map(|(key, value)| (key.as_ref(), value)).collect::<Vec<_>>()
        );
        assert_eq!(flat.keys().last().map(CompactString::as_str), Some("shared"));
        assert!(!flat.keys().last().unwrap().is_heap_allocated());
        flat.extend([(Cow::Borrowed("zz"), f64::NAN)]);
        assert!(flat["zz"].is_nan());
        flat.extend(std::iter::empty::<(CompactString, f64)>());
        assert_eq!(flat.len(), tree.len() + 1);
    }
}