rave_engine 0.8.0

use crate::types::{
    entries::{global::lanes::UnitIssuerMap, GlobalDefinitionExt, LaneDefinition, UnitDefinition},
    Ledger, TxStateManager,
};
use hdi::prelude::{
    trace, wasm_error, ActionHash, ActionHashB64, AgentPubKeyB64, Deserialize, ExternResult,
    Serialize,
};
use rhai::Map;
use serde::Deserializer;
use serde_json::Value;
use std::collections::{BTreeMap, BTreeSet};
use zfuel::fuel::ZFuel;

// this indicated the index at which the unit is added to the global list of units
pub type UnitIndex = u8;

#[derive(Debug, Serialize, Clone, PartialEq, Eq, Default)]
pub struct UnitIndexMap(pub BTreeMap<String, ActionHashB64>);

impl UnitIndexMap {
    pub fn new() -> Self {
        Self(BTreeMap::new())
    }
    pub fn get_unit_indexes(&self) -> Vec<String> {
        self.0.keys().cloned().collect()
    }
}
impl<'de> Deserialize<'de> for UnitIndexMap {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let map: BTreeMap<String, ActionHashB64> = BTreeMap::deserialize(deserializer)?;
        for key in map.keys() {
            key.parse::<UnitIndex>().map_err(serde::de::Error::custom)?;
        }
        Ok(Self(map))
    }
}
// the string is the index of the unyt definition
// ZFuel is the amount of the unyt
#[derive(Debug, Serialize, Clone, PartialEq, Eq)]
pub struct UnitMap(BTreeMap<String, ZFuel>);

impl<'de> Deserialize<'de> for UnitMap {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let map: BTreeMap<String, ZFuel> = BTreeMap::deserialize(deserializer)?;
        for key in map.keys() {
            key.parse::<UnitIndex>().map_err(serde::de::Error::custom)?;
        }
        Ok(Self(map))
    }
}

impl TryFrom<Value> for UnitMap {
    type Error = String;

    fn try_from(value: Value) -> Result<Self, Self::Error> {
        serde_json::from_value(value).map_err(|e| e.to_string())
    }
}

impl Default for UnitMap {
    fn default() -> Self {
        Self::new()
    }
}

impl From<Vec<(u32, &str)>> for UnitMap {
    fn from(amounts: Vec<(u32, &str)>) -> Self {
        let mut map = std::collections::BTreeMap::new();
        for (k, v) in amounts {
            map.insert(k.to_string(), v.parse().expect("Failed to parse amount"));
        }
        Self(map)
    }
}

impl UnitMap {
    /// No unit entries at all — the shape of a *name-only* allocation, which
    /// carries conservation `sources` but no collectable value (distinct from
    /// an explicit zero amount, which some unit definitions validly allow).
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    // we check if any unit index has amount zero apart from base unit
    // if its zero we dont need to keep it so remove it
    pub fn remove_zero_amounts(&mut self, skip_base_unit: bool) {
        self.0.retain(|key, amount| {
            if skip_base_unit && key == "0" {
                return true;
            }
            amount != &ZFuel::zero()
        });
    }

    /// Keep only strictly positive amounts (e.g. alliance supply totals).
    pub fn retain_positive_only(&mut self) {
        self.0.retain(|_, amount| *amount > ZFuel::zero());
    }

    pub fn invert_signs(
        &self,
        global_definition: &GlobalDefinitionExt,
        lane_definitions: &[ActionHash],
    ) -> ExternResult<Self> {
        let unit_def_map = self.get_unit_definition(global_definition, lane_definitions)?;

        let mut new_amount = Self::new();
        for (key, amount) in self.0.iter() {
            let unit_def = unit_def_map
                .get(key)
                .ok_or(wasm_error!("Unit definition not found for index: {}", key))?;
            if unit_def.unit_type.should_invert_sign() {
                new_amount.0.insert(key.clone(), (ZFuel::zero() - amount)?);
            } else {
                new_amount.0.insert(key.clone(), *amount);
            }
        }
        Ok(new_amount)
    }

    /// Returns the unit indexes whose amounts differ between `self` and `other`,
    /// treating a missing index as a zero amount. The comparison is symmetric over
    /// the union of both key sets, so an index present on only one side is reported
    /// when its amount is non-zero. A `None` other is treated as empty, so every
    /// non-zero index in `self` is reported.
    pub fn differing_unit_indexes(&self, other: &Option<Self>) -> ExternResult<Vec<String>> {
        trace!("comparing self: {:?} with other: {:?}", self, other);
        let empty = Self::new();
        let other = other.as_ref().unwrap_or(&empty);
        let all_keys: BTreeSet<String> = self.0.keys().chain(other.0.keys()).cloned().collect();
        Ok(all_keys
            .into_iter()
            .filter(|key| self.get_safe(key) != other.get_safe(key))
            .collect())
    }
    pub fn should_run_credit_check(
        &self,
        expected_ledger: &Ledger,
        previous_ledger: &Ledger,
        unit_def_map: &BTreeMap<String, UnitDefinition>,
    ) -> bool {
        for key in self.0.keys() {
            let expected_value = expected_ledger.balance.get_safe(key);
            let expected_proposed = expected_ledger.proposed_balance.get_safe(key);
            let previous_value = previous_ledger.balance.get_safe(key);
            let previous_proposed = previous_ledger.proposed_balance.get_safe(key);

            // proposal only track debit and so it would be a negative value
            // so we need to add(not subtract) the two values to get the net value
            let expected_net = match expected_value + expected_proposed {
                Ok(value) => value,
                Err(_) => continue,
            };
            let previous_net = match previous_value + previous_proposed {
                Ok(value) => value,
                Err(_) => continue,
            };

            // only credit check units whose net balance decreased and went negative
            if expected_net < previous_net && expected_net < ZFuel::zero() {
                if let Some(unit_def) = unit_def_map.get(key) {
                    if unit_def.unit_type.first_law_of_thermodynamics() {
                        return true;
                    }
                }
            }
        }
        false
    }
    pub fn summarize(
        &mut self,
        handle: &TxStateManager,
        new_amount: Self,
        unit_def_map: &BTreeMap<String, UnitDefinition>,
    ) -> ExternResult<()> {
        for (key, unit) in new_amount.0.iter() {
            let unit_def = unit_def_map
                .get(key)
                .ok_or(wasm_error!("Unit definition not found for index: {}", key))?;
            // based on the unit def we will choose to sumarize/ update the ledger appropriately
            if let Some(updated_balance) = unit_def
                .unit_type
                .summarize(handle, key, self.0.get(key).unwrap_or(&ZFuel::zero()), unit)
                .unwrap_or_else(|_| panic!("Failed to summarize {handle:?}"))
            {
                // update the balance and not just insert new key
                self.0.insert(key.clone(), updated_balance);
            }
        }
        Ok(())
    }
    fn check_validity(&self, unit_def_map: &BTreeMap<String, UnitDefinition>) -> ExternResult<()> {
        // check that if this value can be a valid `0`
        // so check if any of the Unit def is valid for 0
        // get the indexes that is used

        // iter through used index check if zero and then see if its valid
        for index in self.get_unit_indexes() {
            if self.0.get(&index).unwrap_or(&ZFuel::zero()) == &ZFuel::zero() {
                let unit_def = unit_def_map.get(&index).ok_or(wasm_error!(
                    "Unit definition not found for index: {}",
                    index
                ))?;
                if !unit_def.is_valid_for_zero() {
                    return Err(wasm_error!(
                        "unit index {} ({}) is not valid for zero",
                        index,
                        unit_def.unit_symbol
                    ));
                }
            }
        }

        Ok(())
    }
    /// Normalizes the precision of all unit amounts to match their unit definitions.
    ///
    /// This function aligns the precision of each amount to its unit definition's expected precision.
    /// It only succeeds if the conversion can be done without data loss:
    /// - ✅ 3.100 → 3.1 (precision 3 → 1): No data loss, trailing zeros removed
    /// - ❌ 3.45 → 3.? (precision 2 → 1): Data loss, would lose the .05
    ///
    /// # Errors
    /// Returns an error if:
    /// - Unit definition is not found for an index
    /// - Precision conversion would lose significant digits (via ZFuel::to_precision)
    ///
    /// # Example
    /// ```ignore
    /// // If unit definition expects precision 2:
    /// amount = 100.000  // precision 3
    /// normalize_precision() // → 100.00 (precision 2) ✓
    ///
    /// amount = 100.456  // precision 3
    /// normalize_precision() // → Error! Would lose .006
    /// ```
    pub fn normalize_precision(
        &mut self,
        global_definition: &GlobalDefinitionExt,
        lane_definitions: &[ActionHash],
    ) -> ExternResult<BTreeMap<String, UnitDefinition>> {
        let unit_def_map = self.get_unit_definition(global_definition, lane_definitions)?;

        for (index, amount) in self.0.iter_mut() {
            let unit_def = unit_def_map.get(index).ok_or(wasm_error!(
                "Unit definition not found for index: {}",
                index
            ))?;

            // Normalize precision to match unit definition
            // This will error if conversion would lose data
            unit_def.unit_type.normalize_precision(amount)?;
        }
        self.check_validity(&unit_def_map)?;
        Ok(unit_def_map)
    }
    // validate the units in the map against the unit definitions
    pub fn validate_units(
        &self,
        unit_def_map: &BTreeMap<String, UnitDefinition>,
    ) -> ExternResult<()> {
        for (index, amount) in self.0.iter() {
            let unit_def = unit_def_map.get(index).ok_or(wasm_error!(
                "Unit definition not found for index: {}",
                index
            ))?;
            unit_def.unit_type.validate_unit(amount)?
        }
        Ok(())
    }

    /// Blocks rated/nameable transfers in negotiated flow (proposal/counter-proposal/commitment/accept).
    /// Agreement flow (parked spend + receipt) is validated elsewhere via state gating.
    pub fn disallow_rated_and_nameable_units(
        &self,
        unit_def_map: &BTreeMap<String, UnitDefinition>,
    ) -> ExternResult<()> {
        for (index, amount) in self.0.iter() {
            if amount == &ZFuel::zero() {
                continue;
            }
            let unit_def = unit_def_map.get(index).ok_or(wasm_error!(
                "Unit definition not found for index: {}",
                index
            ))?;
            if matches!(
                unit_def.unit_type,
                crate::types::entries::UnytType::RatedUnit(_)
                    | crate::types::entries::UnytType::NameableUnit(_)
            ) {
                return Err(wasm_error!(
                    "Rated and Nameable units may only be transferred via agreements (parked spend and receipt), not in proposal/commitment/accept flow"
                ));
            }
        }
        Ok(())
    }

    /// Validates who may issue units.
    /// For positive amounts:
    /// - if unit index is in `unit_issuers`, author must be listed
    /// - if unit index is missing, anyone may issue
    pub fn validate_unit_issuers(
        &self,
        author: &AgentPubKeyB64,
        unit_issuers: &UnitIssuerMap,
        unit_def_map: &BTreeMap<String, UnitDefinition>,
    ) -> ExternResult<()> {
        for (index, amount) in self.0.iter() {
            if amount <= &ZFuel::zero() {
                continue;
            }
            let unit_def = unit_def_map.get(index).ok_or(wasm_error!(
                "Unit definition not found for index: {}",
                index
            ))?;
            // Issuer permissions apply only to rated/nameable units.
            if !matches!(
                unit_def.unit_type,
                crate::types::entries::UnytType::RatedUnit(_)
                    | crate::types::entries::UnytType::NameableUnit(_)
            ) {
                continue;
            }
            if let Some(allowed_issuers) = unit_issuers.get(index) {
                if !allowed_issuers.contains(author) {
                    return Err(wasm_error!(format!(
                        "Agent {} is not allowed to issue unit index {}",
                        author, index
                    )));
                }
            }
        }
        Ok(())
    }

    pub fn get_unit_definition(
        &self,
        global_definition: &GlobalDefinitionExt,
        lane_id: &[ActionHash],
    ) -> ExternResult<BTreeMap<String, UnitDefinition>> {
        Self::get_unit_definitions(&self.get_unit_indexes(), global_definition, lane_id)
    }
    /// Resolve the unit definitions for an explicit set of indexes — the global
    /// definition's service units first, then the lanes'. Errors if any index has
    /// no service-unit definition. Resolution only needs the index keys, so a
    /// caller holding just an index list (e.g. a diff result) uses this directly
    /// rather than fabricating a `UnitMap`.
    pub fn get_unit_definitions(
        indexes: &[String],
        global_definition: &GlobalDefinitionExt,
        lane_id: &[ActionHash],
    ) -> ExternResult<BTreeMap<String, UnitDefinition>> {
        let mut lanes = vec![];
        for hash in lane_id {
            lanes.push(LaneDefinition::must_get(hash)?);
        }
        let mut result = BTreeMap::new();
        for index in indexes {
            // global definition first, then the lanes; an unknown index is invalid
            let service_unit_hash = global_definition
                .lane_def
                .service_units
                .0
                .get(index)
                .or_else(|| lanes.iter().find_map(|l| l.service_units.0.get(index)))
                .ok_or(wasm_error!(
                    "unable to find the service unit definition for the index: {}",
                    index
                ))?;
            result.insert(
                index.clone(),
                UnitDefinition::must_get(&service_unit_hash.clone().into())?,
            );
        }
        Ok(result)
    }
    pub fn new() -> Self {
        Self(BTreeMap::new())
    }
    pub fn get(&self, key: &str) -> Option<ZFuel> {
        self.0.get(key).cloned()
    }
    pub fn get_safe(&self, key: &String) -> ZFuel {
        self.0.get(key).cloned().unwrap_or(ZFuel::zero())
    }
    pub fn get_unit_indexes(&self) -> Vec<String> {
        self.0.keys().cloned().collect()
    }
    pub fn into_iter(&self) -> BTreeMap<String, ZFuel> {
        self.0.clone()
    }
    pub fn load(map: BTreeMap<String, ZFuel>) -> Self {
        Self(map)
    }
    // Add two UnitMap together
    pub fn add(&mut self, amount: Self) -> ExternResult<()> {
        for (key, amount) in amount.0.iter() {
            if self.0.contains_key(key) {
                self.0.insert(
                    key.clone(),
                    (self.0.get(key).unwrap_or(&ZFuel::zero()) + amount)?,
                );
            } else {
                self.0.insert(key.clone(), *amount);
            }
        }
        Ok(())
    }
    // Subtract two UnitMap together
    pub fn sub(&mut self, amount: Self) -> ExternResult<()> {
        for (key, amount) in amount.0.iter() {
            if self.0.contains_key(key) {
                self.0.insert(
                    key.clone(),
                    (self.0.get(key).unwrap_or(&ZFuel::zero()) - amount)?,
                );
            } else {
                let new_amount = (ZFuel::zero() - amount)?;
                self.0.insert(key.clone(), new_amount);
            }
        }
        Ok(())
    }
    // Convert the UnitMap to a Value
    pub fn to_value(&self) -> Value {
        let json_map: serde_json::Map<String, Value> = self
            .0
            .iter()
            .map(|(k, v)| (k.clone(), v.to_string().into()))
            .collect();
        json_map.into()
    }
    pub fn to_map(&self) -> Map {
        let mut map = Map::new();
        for (k, v) in self.0.iter() {
            map.insert(k.clone().into(), v.to_string().into());
        }
        map
    }
    pub fn is_zero(&self) -> bool {
        // check that all values are zero
        if self.0.is_empty() {
            return true;
        }
        self.0.iter().all(|(_, value)| value == &ZFuel::zero())
    }
    pub fn is_negative(&self) -> bool {
        if self.0.is_empty() {
            return false;
        }
        self.0.iter().any(|(_, value)| value < &ZFuel::zero())
    }
    /// Indexes whose amount is negative. Conservation scans allocation
    /// components for these: inside a summed total a negative conserved
    /// component cancels against a padded positive one (`+150` and `−50` net
    /// to `+100`), so the sum alone cannot see the mint it masks.
    pub fn negative_unit_indexes(&self) -> Vec<String> {
        self.0
            .iter()
            .filter(|(_, value)| *value < &ZFuel::zero())
            .map(|(key, _)| key.clone())
            .collect()
    }
    // /// converts the service UnitMap to negative values
    // /// this is used to the external API's
    // fn handle_service_unit_sign(mut amount: Self) -> ExternResult<Self> {
    //     for amount_value in amount.iter_mut().skip(1) {
    //         let negative_one = Fraction::new(-1, 1)?;
    //         *amount_value = (*amount_value * negative_one)?;
    //     }
    //     Ok(amount)
    // }

    // add all the UnitMap in the vec, and for this we need to add each position to the same position in the  other Unit
    pub fn sum_vec(units: Vec<UnitMap>) -> ExternResult<UnitMap> {
        let mut sum = Self::new();
        for unit in units {
            sum.add(unit)?;
        }
        Ok(sum)
    }

    pub fn push_base_unit(&mut self, amount: ZFuel) {
        self.0.insert("0".to_string(), amount);
    }

    pub fn get_base_unyt(&self) -> ZFuel {
        match self.0.get("0") {
            Some(amount) => *amount,
            None => ZFuel::zero(),
        }
    }
    // check if the current unit map is less than the other unit map
    // this is used to check if the balance is less than the previous ledger
    pub fn less_than(&self, other: &Self) -> bool {
        for (key, amount) in self.0.iter() {
            if amount < other.0.get(key).unwrap_or(&ZFuel::zero()) {
                return true;
            }
        }
        false
    }
    pub fn contains_base_unit(&self) -> bool {
        self.0.contains_key("0")
    }
    pub fn get_service_units(&self) -> Vec<String> {
        self.0.keys().filter(|k| k != &"0").cloned().collect()
    }
}

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

    #[test]
    fn is_zero_true_for_empty() {
        assert!(UnitMap::new().is_zero());
    }

    #[test]
    fn is_zero_true_for_all_zero_values() {
        let m = UnitMap::from(vec![(0u32, "0"), (1u32, "0")]);
        assert!(m.is_zero());
    }

    #[test]
    fn is_zero_false_when_any_nonzero() {
        let m = UnitMap::from(vec![(0u32, "0"), (1u32, "10")]);
        assert!(!m.is_zero());
    }

    #[test]
    fn is_negative_false_for_empty() {
        assert!(!UnitMap::new().is_negative());
    }

    #[test]
    fn is_negative_false_for_all_positive() {
        let m = UnitMap::from(vec![(0u32, "5"), (1u32, "10")]);
        assert!(!m.is_negative());
    }

    #[test]
    fn is_negative_true_when_any_negative() {
        let m = UnitMap::from(vec![(0u32, "5"), (1u32, "-1")]);
        assert!(m.is_negative());
    }

    #[test]
    fn is_negative_false_when_all_zero() {
        let m = UnitMap::from(vec![(0u32, "0"), (1u32, "0")]);
        assert!(!m.is_negative());
    }

    #[test]
    fn differing_indexes_reports_value_mismatch() {
        let a = UnitMap::from(vec![(0u32, "5"), (1u32, "3")]);
        let b = UnitMap::from(vec![(0u32, "5"), (1u32, "4")]);
        assert_eq!(
            a.differing_unit_indexes(&Some(b)).unwrap(),
            vec!["1".to_string()]
        );
    }

    #[test]
    fn differing_indexes_reports_keys_only_in_other() {
        // symmetric: an index present only on the `other` side is still reported
        let a = UnitMap::from(vec![(0u32, "5")]);
        let b = UnitMap::from(vec![(0u32, "5"), (2u32, "7")]);
        assert_eq!(
            a.differing_unit_indexes(&Some(b)).unwrap(),
            vec!["2".to_string()]
        );
    }

    #[test]
    fn differing_indexes_treats_missing_as_zero() {
        // a zero amount on one side and a missing key on the other are not a difference
        let a = UnitMap::from(vec![(0u32, "5"), (1u32, "0")]);
        let b = UnitMap::from(vec![(0u32, "5")]);
        assert!(a.differing_unit_indexes(&Some(b)).unwrap().is_empty());
    }

    #[test]
    fn differing_indexes_none_other_returns_nonzero_keys() {
        let a = UnitMap::from(vec![(0u32, "5"), (1u32, "0")]);
        assert_eq!(
            a.differing_unit_indexes(&None).unwrap(),
            vec!["0".to_string()]
        );
    }
}