// address.rs
use super::{chain::UserChainData, transaction::UserTransactionData};
use crate::error::UserStateError;

use ic_cdk::export::{candid::CandidType, serde::Deserialize};
use std::collections::HashMap;

/// Represents the user address data structure, including the chain_data.
#[derive(Debug, CandidType, Deserialize, Clone)]
pub struct UserAddressData {
    pub name: String,
    pub hidden: bool,
    pub disabled: bool,
    pub public_key: Vec<u8>,
    pub chain_data: HashMap<u64, UserChainData>,
}

impl UserAddressData {
    /// Creates a new AddressData with the given name and public_key.
    pub fn new(public_key: Vec<u8>, name: String) -> Self {
        Self {
            name,
            public_key,
            hidden: false,
            disabled: false,
            chain_data: HashMap::default(),
        }
    }

    // Safe access method to get a chain by chain_id
    pub fn get_chain(&self, chain_id: u64) -> Result<&UserChainData, UserStateError> {
        self.chain_data
            .get(&chain_id)
            .ok_or(UserStateError::ChainNotFound)
    }

    /// Adds a transaction to the specified chain and updates the ChainData.
    /// Returns an error if the chain_id is not found.
    pub fn add_transaction(
        &mut self,
        chain_id: u64,
        nonce: u64,
        transaction: UserTransactionData,
    ) -> Result<&UserChainData, UserStateError> {
        if let Some(chain_data) = self.chain_data.get_mut(&chain_id) {
            chain_data.add(nonce, transaction);

            Ok(chain_data)
        } else {
            Err(UserStateError::ChainNotFound)
        }
    }

    /// Gets the ChainData for a specific chain_id.
    pub fn get_transactions(&self, chain_id: u64) -> Result<&UserChainData, UserStateError> {
        self.chain_data
            .get(&chain_id)
            .ok_or(UserStateError::ChainNotFound)
    }

    /// Clears the transactions vector for the specified chain.
    /// Returns an error if the chain_id is not found.
    pub fn clear_transactions(&mut self, chain_id: u64) -> Result<&UserChainData, UserStateError> {
        if let Some(chain_data) = self.chain_data.get_mut(&chain_id) {
            chain_data.transactions.clear();

            Ok(chain_data)
        } else {
            Err(UserStateError::ChainNotFound)
        }
    }

    /// Adds a new chain to the address data.
    /// Returns an error if the chain_id is already in use.
    pub fn add_chain(
        &mut self,
        chain_id: u64,
        chain_data: UserChainData,
    ) -> Result<&UserChainData, UserStateError> {
        if self.chain_data.contains_key(&chain_id) {
            Err(UserStateError::ChainAlreadyExists)
        } else {
            self.chain_data.insert(chain_id, chain_data);

            if let Some(chain_data) = self.chain_data.get(&chain_id) {
                Ok(chain_data)
            } else {
                Err(UserStateError::ChainNotFound)
            }
        }
    }

    /// Removes a chain from the address data.
    /// Returns an error if the chain_id is not found.
    pub fn remove_chain(&mut self, chain_id: u64) -> Result<UserChainData, UserStateError> {
        if let Some(chain_data) = self.chain_data.remove(&chain_id) {
            Ok(chain_data)
        } else {
            Err(UserStateError::ChainNotFound)
        }
    }

    /// Returns the number of chains for the address.
    /// This is used for the UI to determine if the address is empty.
    pub fn chain_count(&self) -> usize {
        self.chain_data.len()
    }

    /// Returns the number of transactions for the address.
    /// This is used for the UI to determine if the address is empty.
    pub fn transaction_count(&self) -> usize {
        self.chain_data
            .values()
            .map(|chain_data| chain_data.transactions.len())
            .sum()
    }

    /// Returns the number of transactions for the specified chain.
    /// This is used for the UI to determine if the chain is empty.
    pub fn chain_transaction_count(&self, chain_id: u64) -> usize {
        if let Some(chain_data) = self.chain_data.get(&chain_id) {
            chain_data.transactions.len()
        } else {
            0
        }
    }

    /// Set Nonce for a specific chain_id.
    /// Returns an error if the chain_id is not found.
    pub fn set_nonce(
        &mut self,
        chain_id: u64,
        nonce: u64,
    ) -> Result<&UserChainData, UserStateError> {
        if let Some(chain_data) = self.chain_data.get_mut(&chain_id) {
            chain_data.nonce = nonce;

            Ok(chain_data)
        } else {
            Err(UserStateError::ChainNotFound)
        }
    }

    /// Get Nonce for a specific chain_id.
    /// Returns an error if the chain_id is not found.
    pub fn get_nonce(&self, chain_id: u64) -> Result<u64, UserStateError> {
        if let Some(chain_data) = self.chain_data.get(&chain_id) {
            Ok(chain_data.nonce)
        } else {
            Err(UserStateError::ChainNotFound)
        }
    }

    /// Get Transaction for a specific chain_id and index.
    /// Returns an error if the chain_id is not found.
    pub fn get_transaction(
        &self,
        chain_id: u64,
        index: usize,
    ) -> Result<&UserTransactionData, UserStateError> {
        let chain_data = self.get_chain(chain_id)?;
        chain_data.get_transaction(index)
    }
}

#[derive(Clone, Debug, CandidType, Default, Deserialize)]
pub struct UserAddressArgs {
    pub public_key: Vec<u8>,
    pub name: Option<String>,
}

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

    proptest! {
        #[test]
        fn test_add_and_clear_transactions(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            nonce: u64,
            transactions: Vec<UserTransactionData>,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);
            address_data.add_chain(chain_id, UserChainData::default()).unwrap();

            for (index, transaction) in transactions.iter().enumerate() {
                address_data.add_transaction(chain_id, nonce + index as u64, transaction.clone()).unwrap();
            }

            let chain_data = address_data.get_transactions(chain_id).unwrap();
            assert_eq!(chain_data.transactions.len(), transactions.len());

            address_data.clear_transactions(chain_id).unwrap();

            let chain_data = address_data.get_transactions(chain_id).unwrap();
            assert_eq!(chain_data.transactions.len(), 0);
        }

        #[test]
        fn test_add_chain_error(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            chain_data: UserChainData,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);
            address_data.add_chain(chain_id, chain_data.clone()).unwrap();

            let result = address_data.add_chain(chain_id, chain_data);

            match result {
                Err(UserStateError::ChainAlreadyExists) => assert!(true),
                _ => panic!("Expected ChainAlreadyExists error"),
            }
        }

        #[test]
        fn test_add_transaction_error(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            nonce: u64,
            transaction: UserTransactionData,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);

            let result = address_data.add_transaction(chain_id, nonce, transaction);

            match result {
                Err(UserStateError::ChainNotFound) => assert!(true),
                _ => panic!("Expected ChainNotFound error"),
            }
        }

        #[test]
        fn test_clear_transactions_error(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);

            let result = address_data.clear_transactions(chain_id);

            match result {
                Err(UserStateError::ChainNotFound) => assert!(true),
                _ => panic!("Expected ChainNotFound error"),
            }
        }

        #[test]
        fn test_get_transaction(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            nonce: u64,
            transactions: Vec<UserTransactionData>,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);
            address_data.add_chain(chain_id, UserChainData::default()).unwrap();

            for (index, transaction) in transactions.iter().enumerate() {
                address_data.add_transaction(chain_id, nonce + index as u64, transaction.clone()).unwrap();
            }

            for (index, transaction) in transactions.iter().enumerate() {
                let result = address_data.get_transaction(chain_id, index).unwrap();
                assert_eq!(result, transaction);
            }
        }

        #[test]
        fn test_get_transaction_error(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            index: usize,
        ) {
            let address_data = UserAddressData::new(public_key, name);

            let result = address_data.get_transaction(chain_id, index);

            match result {
                Err(UserStateError::ChainNotFound) => assert!(true),
                _ => panic!("Expected ChainNotFound error"),
            }
        }

        #[test]
        fn test_get_transactions(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            nonce: u64,
            transactions: Vec<UserTransactionData>,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);
            address_data.add_chain(chain_id, UserChainData::default()).unwrap();

            for (index, transaction) in transactions.iter().enumerate() {
                address_data.add_transaction(chain_id, nonce + index as u64, transaction.clone()).unwrap();
            }

            let chain_data = address_data.get_transactions(chain_id).unwrap();
            assert_eq!(chain_data.transactions.len(), transactions.len());
        }

        #[test]
        fn test_get_transactions_error(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
        ) {
            let address_data = UserAddressData::new(public_key, name);

            let result = address_data.get_transactions(chain_id);

            match result {
                Err(UserStateError::ChainNotFound) => assert!(true),
                _ => panic!("Expected ChainNotFound error"),
            }
        }

        #[test]
        fn test_get_chain(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            chain_data: UserChainData,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);
            address_data.add_chain(chain_id, chain_data.clone()).unwrap();

            let result = address_data.get_chain(chain_id).unwrap();

            assert_eq!(result.clone(), chain_data);
        }

        #[test]
        fn test_get_chain_error(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
        ) {
            let address_data = UserAddressData::new(public_key, name);

            let result = address_data.get_chain(chain_id);

            match result {
                Err(UserStateError::ChainNotFound) => assert!(true),
                _ => panic!("Expected ChainNotFound error"),
            }
        }

        #[test]
        fn test_get_nonce(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            nonce: u64,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);
            address_data.add_chain(chain_id, UserChainData::default()).unwrap();
            address_data.set_nonce(chain_id, nonce).unwrap();

            let result = address_data.get_nonce(chain_id).unwrap();
            assert_eq!(result, nonce);
        }

        #[test]
        fn test_get_nonce_error(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
        ) {
            let address_data = UserAddressData::new(public_key, name);

            let result = address_data.get_nonce(chain_id);

            match result {
                Err(UserStateError::ChainNotFound) => assert!(true),
                _ => panic!("Expected ChainNotFound error"),
            }
        }

        #[test]
        fn test_set_nonce(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            nonce: u64,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);
            address_data.add_chain(chain_id, UserChainData::default()).unwrap();

            let result = address_data.set_nonce(chain_id, nonce);
            assert!(result.is_ok());
        }

        #[test]
        fn test_set_nonce_error(
            public_key: Vec<u8>,
            name: String,
            chain_id: u64,
            nonce: u64,
        ) {
            let mut address_data = UserAddressData::new(public_key, name);

            let result = address_data.set_nonce(chain_id, nonce);

            match result {
                Err(UserStateError::ChainNotFound) => assert!(true),
                _ => panic!("Expected ChainNotFound error"),
            }
        }

    }
}