sdk_common/src/process.rs

use std::{
    collections::HashMap,
    sync::{Mutex, MutexGuard, OnceLock},
};

use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use sp_client::bitcoin::OutPoint;

use crate::{
    pcd::{AnkPcdHash, Pcd, RoleDefinition},
    prd::{Prd, PrdType},
    signature::Proof,
    MutexExt,
};

#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct ProcessState {
    pub commited_in: OutPoint,
    pub pcd_commitment: Value, // If we can't modify a field, we just copy the previous value
    pub encrypted_pcd: Value, // Some fields may be clear, if the owner of the process decides so
    pub keys: Map<String, Value>, // We may not always have all the keys
    pub validation_tokens: Vec<Proof>, // Signature of the hash of the encrypted pcd tagged with some decision like "yes" or "no"
}

impl ProcessState {
    pub fn decrypt_pcd(&self) -> Value {
        // TODO add real error management
        let mut fields2plain = Map::new();
        let _ = self.encrypted_pcd.decrypt_fields(&self.keys, &mut fields2plain);
        Value::Object(fields2plain)
    }

    fn compute_modified_fields(&self, previous_state: Option<&ProcessState>) -> Vec<String> {
        let new_state = &self.encrypted_pcd;

        // Ensure the new state is a JSON object
        let new_state_obj = new_state
            .as_object()
            .expect("New state should be a JSON object");

        if let Some(prev_state) = previous_state {
            // Previous state exists; compute differences
            let previous_state_obj = prev_state
                .encrypted_pcd
                .as_object()
                .expect("Previous state should be a JSON object");

            // Compute modified fields by comparing with previous state
            new_state_obj
                .iter()
                .filter_map(|(key, value)| {
                    let previous_value = previous_state_obj.get(key);
                    if previous_value.is_none() || value != previous_value.unwrap() {
                        Some(key.clone())
                    } else {
                        None
                    }
                })
                .collect()
        } else {
            // No previous state; all fields are considered modified
            new_state_obj.keys().cloned().collect()
        }
    }

    pub fn is_valid(&self, previous_state: Option<&ProcessState>) -> anyhow::Result<()> {
        if self.validation_tokens.is_empty() {
            return Err(anyhow::anyhow!(
                "Can't validate a state with no proofs attached"
            ));
        }

        // Compute modified fields
        let modified_fields = self.compute_modified_fields(previous_state);

        if modified_fields.is_empty() {
            return Err(anyhow::anyhow!("State is identical to the previous state"));
        }

        println!("modified fields: {:?}", modified_fields);

        // Compute the hash of the new state
        let new_state_hash = AnkPcdHash::from_value(&self.encrypted_pcd);

        let mut fields2plains = Map::new();
        self.encrypted_pcd
            .decrypt_fields(&self.keys, &mut fields2plains)?;

        let mut roles2rules: HashMap<String, RoleDefinition> = HashMap::new();
        if let Some(roles) = fields2plains.get("roles") {
            if let Some(roles_map) = roles.as_object() {
                for (role, conditions) in roles_map {
                    let role_def = serde_json::from_value::<RoleDefinition>(conditions.clone())?;
                    roles2rules.insert(role.to_string(), role_def);
                }
            } else {
                return Err(anyhow::anyhow!("Roles is not an object"));
            }
        } else {
            return Err(anyhow::anyhow!("Missing roles in the encrypted pcd"));
        }

        // Check if each modified field satisfies at least one applicable rule across all roles
        let all_fields_validated = modified_fields.iter().all(|field| {
            // Collect applicable rules from all roles for the current field
            let applicable_roles: Vec<RoleDefinition> = roles2rules
                .iter()
                .filter_map(|(_, role_def)| {
                    let mut filtered_role_def = role_def.clone();
                    let rules = filtered_role_def.get_applicable_rules(field);
                    filtered_role_def.validation_rules =
                        rules.into_iter().map(|r| r.clone()).collect();
                    if filtered_role_def.validation_rules.is_empty() {
                        None
                    } else {
                        Some(filtered_role_def)
                    }
                })
                .collect();

            if applicable_roles.is_empty() {
                return false; // No rules apply to this field, consider it invalid
            }

            applicable_roles.into_iter().any(|role_def| {
                role_def.validation_rules.iter().any(|rule| {
                    rule.is_satisfied(
                        field,
                        new_state_hash.clone(),
                        &self.validation_tokens,
                        &role_def.members,
                    )
                })
            })
        });

        if all_fields_validated {
            Ok(())
        } else {
            Err(anyhow::anyhow!("Not enough valid proofs"))
        }
    }

    pub fn is_empty(&self) -> bool {
        self.encrypted_pcd == Value::Null || 
            self.pcd_commitment == Value::Null
    }
}

/// A process is basically a succession of states
/// If a process has nothing to do with us, impending_requests will be empty
/// The latest state MUST be an empty state with only the commited_in field set at the last unspent outpoint
#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
pub struct Process {
    states: Vec<ProcessState>,
    impending_requests: Vec<Prd>,
}

impl Process {
    pub fn new(
        commited_in: OutPoint
    ) -> Self {
        let empty_state = ProcessState {
            commited_in,
            ..Default::default()
        };
        Self {
            states: vec![empty_state],
            impending_requests: vec![],
        }
    }

    pub fn get_last_unspent_outpoint(&self) -> anyhow::Result<OutPoint> {
        if self.states.is_empty() { return Err(anyhow::Error::msg("Empty Process")); }
        let last_state = self.states.last().unwrap();
        Ok(last_state.commited_in)
    }

    pub fn update_states_tip(&mut self, new_commitment: OutPoint) -> anyhow::Result<()> {
        if self.states.is_empty() { return Err(anyhow::Error::msg("Empty Process")); }
        let last_value = self.states.last().unwrap();
        if !last_value.is_empty() {
            return Err(anyhow::Error::msg("Last value should be empty"));
        }
        if last_value.commited_in == new_commitment {
            return Err(anyhow::Error::msg("new_commitment is the same than existing tip"));
        }

        // Before updating we make sure that we only have one concurrent state
        let concurrent_states = self.get_latest_concurrent_states()?;
        if concurrent_states.len() != 2 {
            return Err(anyhow::Error::msg("We must have exactly one state for the current tip"));
        }

        // Replace the existing tip
        let new_tip = ProcessState {
            commited_in: new_commitment,
            ..Default::default()
        };
        let _ = self.states.pop().unwrap();
        self.states.push(new_tip);

        Ok(())
    }

    /// We want to insert a new state that would be commited by the last UTXO
    /// The new state *must* have the same commited_in than the last empty one
    /// We want to always keep an empty state with only the latest unspent commited_in value at the last position
    pub fn insert_concurrent_state(&mut self, new_state: ProcessState) -> anyhow::Result<()> {
        if self.states.is_empty() { return Err(anyhow::Error::msg("Empty Process")); }
        let last_value = self.states.last().unwrap();
        if !last_value.is_empty() {
            return Err(anyhow::Error::msg("Last value should be empty"));
        }
        if last_value.commited_in != new_state.commited_in {
            return Err(anyhow::Error::msg("A concurrent state must have the same commited in than the tip of the states"));
        }
        let empty_state = self.states.pop().unwrap();
        self.states.push(new_state);
        // We always keep an empty state at the end
        self.states.push(empty_state);
        Ok(())
    }

    pub fn get_state_at(&self, index: usize) -> Option<&ProcessState> {
        self.states.get(index)
    }

    pub fn get_state_at_mut(&mut self, index: usize) -> Option<&mut ProcessState> {
        self.states.get_mut(index)
    }

    pub fn get_latest_state(&self) -> Option<&ProcessState> {
        self.states.last()
    }

    pub fn get_latest_state_mut(&mut self) -> Option<&mut ProcessState> {
        self.states.last_mut()
    }

    pub fn get_previous_state(&self, current_state: &ProcessState) -> Option<&ProcessState> {
        // Find the index of the current state
        let current_index = self
            .states
            .iter()
            .position(|state| state == current_state)?;

        // Check if there is a previous state
        if current_index > 0 {
            // Create a new Process with the previous state
            let previous_state = self.get_state_at(current_index - 1).unwrap();
            Some(&previous_state)
        } else {
            None // No previous state exists
        }
    }

    pub fn get_state_for_commitments_root(&mut self, merkle_root: [u8; 32]) -> anyhow::Result<&mut ProcessState> {
        if self.get_number_of_states() == 0 {
            // This should never happen, but we better get rid of it now
            return Err(anyhow::Error::msg("process is empty".to_owned()));
        }

        for p in self.get_latest_concurrent_states_mut()? {
            if p.is_empty() {
                continue;
            }
            let root = <Value as Pcd>::create_merkle_tree(&p.pcd_commitment).unwrap().root().unwrap();
            if merkle_root == root { 
                return Ok(p);
            }
        }

        return Err(anyhow::Error::msg("No state for this merkle root"));
    }

    /// This is useful when multiple unvalidated states are pending waiting for enough validations
    /// It returns the latest state and all the previous states that have the same commited_in
    /// It means that all the returned states are unvalidated and except for the one that get validated they will be pruned
    pub fn get_latest_concurrent_states(&self) -> anyhow::Result<Vec<&ProcessState>> {
        if self.get_number_of_states() == 0 {
            // This should never happen, but we better get rid of it now
            return Err(anyhow::Error::msg("process is empty".to_owned()));
        }

        let mut states = vec![];
        let mut previous_commited_in = OutPoint::null();
        // We iterate backwards until we find a state that has a different commited_in
        for state in self.states.iter().rev() {
            log::debug!("state: {:#?}", state);
            if previous_commited_in == OutPoint::null() {
                previous_commited_in = state.commited_in;
            } else if previous_commited_in != state.commited_in {
                break;
            }
            states.push(state);
        }

        Ok(states)
    }

    pub fn get_latest_concurrent_states_mut(&mut self) -> anyhow::Result<Vec<&mut ProcessState>> {
        if self.get_number_of_states() == 0 {
            // This should never happen, but we better get rid of it now
            return Err(anyhow::Error::msg("process is empty".to_owned()));
        }

        let mut states = vec![];
        let mut previous_commited_in = OutPoint::null();
        // We iterate backwards until we find a state that has a different commited_in
        for state in self.states.iter_mut().rev() {
            if previous_commited_in == OutPoint::null() {
                previous_commited_in = state.commited_in;
            } else if previous_commited_in != state.commited_in {
                break;
            }
            states.push(state);
        }

        Ok(states)
    }

    pub fn remove_all_concurrent_states(&mut self) -> anyhow::Result<Vec<ProcessState>> {
        if self.get_number_of_states() == 0 {
            // This should never happen, but we better get rid of it now
            return Err(anyhow::Error::msg("process is empty".to_owned()));
        }

        let empty_state = self.states.pop().unwrap();
        let last_commitment_outpoint = empty_state.commited_in;
        let split_index = self.states.iter().position(|state| state.commited_in == last_commitment_outpoint).unwrap();

        let removed = self.states.split_off(split_index);

        // We make sure we always have an empty state at the end
        self.states.push(empty_state);

        Ok(removed)
    }

    pub fn get_latest_commited_state(&self) -> Option<&ProcessState> {
        if self.states.is_empty() {
            return None;
        }

        let last_state = self.states.last().unwrap();

        debug_assert!(last_state.is_empty()); // Last state must always be empty

        // We look for the last commited in before all the pending states
        let latest_outpoint = last_state.commited_in;

        return self
            .states
            .iter()
            .rev()
            .find(|s| s.commited_in != latest_outpoint);
    }

    pub fn insert_impending_request(&mut self, request: Prd) {
        self.impending_requests.push(request);
    }

    pub fn get_impending_requests(&self) -> Vec<&Prd> {
        self.impending_requests.iter().collect()
    }

    pub fn get_impending_requests_mut(&mut self) -> Vec<&mut Prd> {
        self.impending_requests.iter_mut().collect()
    }

    pub fn prune_impending_requests(&mut self) {
        self.impending_requests = self.impending_requests.clone().into_iter()
            .filter(|r| r.prd_type != PrdType::None)
            .collect();
    }

    pub fn get_state_index(&self, state: &ProcessState) -> Option<usize> {
        // Get the commited_in value of the provided state
        let target_commited_in = state.commited_in;

        // Find the index of the first state with the same commited_in value
        self.states
            .iter()
            .position(|s| s.commited_in == target_commited_in)
    }

    pub fn get_number_of_states(&self) -> usize {
        self.states.len()
    }
}

pub static CACHEDPROCESSES: OnceLock<Mutex<HashMap<OutPoint, Process>>> = OnceLock::new();

pub fn lock_processes() -> Result<MutexGuard<'static, HashMap<OutPoint, Process>>, anyhow::Error> {
    CACHEDPROCESSES
        .get_or_init(|| Mutex::new(HashMap::new()))
        .lock_anyhow()
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;

    use serde_json::json;
    use sp_client::{
        bitcoin::{secp256k1::SecretKey, Network}, silentpayments::utils::SilentPaymentAddress, spclient::{SpClient, SpWallet, SpendKey}
    };

    use crate::pcd::{Member, ValidationRule};
    use crate::signature::{AnkValidationNoHash, AnkValidationYesHash};

    use super::*;

    fn create_alice_wallet() -> SpWallet {
        SpWallet::new(
            SpClient::new(
                "default".to_owned(),
                SecretKey::from_str(
                    "a67fb6bf5639efd0aeb19c1c584dd658bceda87660ef1088d4a29d2e77846973",
                )
                .unwrap(),
                SpendKey::Secret(
                    SecretKey::from_str(
                        "a1e4e7947accf33567e716c9f4d186f26398660e36cf6d2e711af64b3518e65c",
                    )
                    .unwrap(),
                ),
                None,
                Network::Signet,
            )
            .unwrap(),
            None,
            vec![],
        )
        .unwrap()
    }

    fn create_bob_wallet() -> SpWallet {
        SpWallet::new(
            SpClient::new(
                "default".to_owned(),
                SecretKey::from_str(
                    "4d9f62b2340de3f0bafd671b78b19edcfded918c4106baefd34512f12f520e9b",
                )
                .unwrap(),
                SpendKey::Secret(
                    SecretKey::from_str(
                        "dafb99602721577997a6fe3da54f86fd113b1b58f0c9a04783d486f87083a32e",
                    )
                    .unwrap(),
                ),
                None,
                Network::Signet,
            )
            .unwrap(),
            None,
            vec![],
        )
        .unwrap()
    }

    fn create_carol_wallet() -> SpWallet {
        SpWallet::new(
            SpClient::new(
                "default".to_owned(),
                SecretKey::from_str(
                    "e4a5906eaa1a7ab24d5fc8d9b600d47f79caa6511c056c111677b7a33e62c5e9",
                )
                .unwrap(),
                SpendKey::Secret(
                    SecretKey::from_str(
                        "e4c282e14668af1435e39df78403a7b406a791e3c6e666295496a6a865ade162",
                    )
                    .unwrap(),
                ),
                None,
                Network::Signet,
            )
            .unwrap(),
            None,
            vec![],
        )
        .unwrap()
    }

    fn dummy_process_state() -> ProcessState {
        let alice_wallet = create_alice_wallet();
        let bob_wallet = create_bob_wallet();
        let carol_wallet = create_carol_wallet();

        let alice_address =
            SilentPaymentAddress::try_from(alice_wallet.get_client().get_receiving_address())
                .unwrap();
        let bob_address =
            SilentPaymentAddress::try_from(bob_wallet.get_client().get_receiving_address())
                .unwrap();
        let carol_address =
            SilentPaymentAddress::try_from(carol_wallet.get_client().get_receiving_address())
                .unwrap();

        let alice_bob = Member::new(vec![alice_address, bob_address]).unwrap();
        let carol = Member::new(vec![carol_address]).unwrap();

        let validation_rule1 =
            ValidationRule::new(1.0, vec!["field1".to_owned(), "roles".to_owned()], 0.5).unwrap();
        let validation_rule2 = ValidationRule::new(1.0, vec!["field2".to_owned()], 0.5).unwrap();

        let encrypted_pcd = json!({
            "field1": "value1",
            "field2": "value2",
            "roles": {
                "role1": {
                    "members": [alice_bob],
                    "validation_rules": [validation_rule1]
                },
                "role2": {
                    "members": [carol],
                    "validation_rules": [validation_rule2]
                }
            }
        });

        let outpoint = OutPoint::null();

        let pcd_commitment = encrypted_pcd.hash_fields(outpoint).unwrap();

        let mut fields2keys = Map::new();
        let mut fields2cipher = Map::new();
        // let field_to_encrypt: Vec<String> = encrypted_pcd.as_object().unwrap().keys().map(|k| k.clone()).collect();
        let field_to_encrypt = vec!["field1".to_string(), "field2".to_string()];

        encrypted_pcd
            .encrypt_fields(&field_to_encrypt, &mut fields2keys, &mut fields2cipher)
            .unwrap();

        ProcessState {
            commited_in: outpoint,
            pcd_commitment: Value::Object(pcd_commitment),
            encrypted_pcd,
            keys: Map::new(),
            validation_tokens: vec![],
        }
    }

    fn add_validation_token(state: &mut ProcessState, signing_key: SecretKey, accept: bool) {
        let pcd_hash = AnkPcdHash::from_value(&state.encrypted_pcd);
        if accept {
            let validation_hash = AnkValidationYesHash::from_commitment(pcd_hash);
            let proof = Proof::new(
                crate::signature::AnkHash::ValidationYes(validation_hash),
                signing_key,
            );
            state.validation_tokens.push(proof);
        } else {
            let validation_hash = AnkValidationNoHash::from_commitment(pcd_hash);
            let proof = Proof::new(
                crate::signature::AnkHash::ValidationNo(validation_hash),
                signing_key,
            );
            state.validation_tokens.push(proof);
        }
    }

    #[test]
    fn test_error_no_proofs() {
        let state = dummy_process_state();
        let result = state.is_valid(None);
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err().to_string(),
            "Can't validate a state with no proofs attached"
        );
    }

    #[test]
    fn test_error_identical_previous_state() {
        let mut state = dummy_process_state();
        // We sign with a random key
        let signing_key =
            SecretKey::from_str("39b2a765dc93e02da04a0e9300224b4f99fa7b83cfae49036dff58613fd3277c")
                .unwrap();
        add_validation_token(&mut state, signing_key, true);
        let result = state.is_valid(Some(&state));
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err().to_string(),
            "State is identical to the previous state"
        );
    }

    #[test]
    /// We provide a proof signed with a key that is not the spending key for either alice or bob
    fn test_error_invalid_proof() {
        let mut state = dummy_process_state();
        // We sign with a random key
        let signing_key =
            SecretKey::from_str("39b2a765dc93e02da04a0e9300224b4f99fa7b83cfae49036dff58613fd3277c")
                .unwrap();
        add_validation_token(&mut state, signing_key, true);
        let result = state.is_valid(None);
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().to_string(), "Not enough valid proofs");
    }

    #[test]
    /// Carol signs alone for an init state
    fn test_error_not_enough_signatures() {
        let mut state = dummy_process_state();
        // We sign with Carol key
        let carol_key: SecretKey = create_carol_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        add_validation_token(&mut state, carol_key, true);
        let result = state.is_valid(None);
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().to_string(), "Not enough valid proofs");
    }

    #[test]
    /// Alice signs alone for her fields in an init state
    fn test_valid_just_enough_signatures() {
        let mut state = dummy_process_state();
        // We sign with Alice and Carol keys
        let alice_key: SecretKey = create_alice_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        let carol_key: SecretKey = create_carol_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        add_validation_token(&mut state, alice_key, true);
        add_validation_token(&mut state, carol_key, true);
        let result = state.is_valid(None);
        assert!(result.is_ok());
    }

    #[test]
    /// everyone signs for everything
    fn test_valid_all_signatures() {
        let mut state = dummy_process_state();
        let alice_key: SecretKey = create_alice_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        let bob_key: SecretKey = create_bob_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        let carol_key: SecretKey = create_carol_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        add_validation_token(&mut state, alice_key, true);
        add_validation_token(&mut state, bob_key, true);
        add_validation_token(&mut state, carol_key, true);
        let result = state.is_valid(None);
        assert!(result.is_ok());
    }

    #[test]
    /// Carol refuses change for her part
    fn test_error_carol_votes_no() {
        let mut state = dummy_process_state();
        let alice_key: SecretKey = create_alice_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        let bob_key: SecretKey = create_bob_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        let carol_key: SecretKey = create_carol_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        add_validation_token(&mut state, alice_key, true);
        add_validation_token(&mut state, bob_key, true);
        add_validation_token(&mut state, carol_key, false);
        let result = state.is_valid(None);
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().to_string(), "Not enough valid proofs");
    }

    #[test]
    /// Bob refuses change for his part, but Alice is enough to reach quorum
    fn test_valid_bob_votes_no() {
        let mut state = dummy_process_state();
        let alice_key: SecretKey = create_alice_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        let bob_key: SecretKey = create_bob_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        let carol_key: SecretKey = create_carol_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        add_validation_token(&mut state, alice_key, true);
        add_validation_token(&mut state, bob_key, false);
        add_validation_token(&mut state, carol_key, true);
        let result = state.is_valid(None);
        assert!(result.is_ok());
    }

    #[test]
    /// Everyone signs, and we have a previous state
    fn test_valid_everyone_signs_with_prev_state() {
        let state = dummy_process_state();
        let mut new_state = state.clone();
        if let Value::Object(ref mut map) = new_state.encrypted_pcd {
            // Modify the field
            map.insert("field1".to_string(), Value::String("new_value1".to_owned()));
        } else {
            // Handle the case where encrypted_pcd is not an object
            panic!("encrypted_pcd is not a JSON object.");
        }
        let alice_key: SecretKey = create_alice_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        let bob_key: SecretKey = create_bob_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        let carol_key: SecretKey = create_carol_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        add_validation_token(&mut new_state, alice_key, true);
        add_validation_token(&mut new_state, bob_key, true);
        add_validation_token(&mut new_state, carol_key, true);
        let result = new_state.is_valid(Some(&state));
        assert!(result.is_ok());
    }

    #[test]
    /// Only Carol signs, but she doesn't have modification rights on modified field
    fn test_error_not_right_signatures_with_prev_state() {
        let state = dummy_process_state();
        let mut new_state = state.clone();
        if let Value::Object(ref mut map) = new_state.encrypted_pcd {
            // Modify the field
            map.insert("field1".to_string(), Value::String("new_value1".to_owned()));
        } else {
            // Handle the case where encrypted_pcd is not an object
            panic!("encrypted_pcd is not a JSON object.");
        }
        let carol_key: SecretKey = create_carol_wallet()
            .get_client()
            .get_spend_key()
            .try_into()
            .unwrap();
        add_validation_token(&mut new_state, carol_key, true);
        let result = new_state.is_valid(Some(&state));
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().to_string(), "Not enough valid proofs");
    }
}