use crate::{account::AccountState, db::StateDb};

/// Read/write access to account state.
/// `set_balance` and `inc_nonce` take `&self` because StateDb uses RocksDB's
/// interior mutability — writes do not require exclusive access at the Rust level.
pub trait StateAccessor {
    fn get_balance(&self, address: &str) -> u64;
    fn get_nonce(&self, address: &str) -> u64;
    fn set_balance(&self, address: &str, balance: u64);
    fn inc_nonce(&self, address: &str);
}

impl StateAccessor for StateDb {
    fn get_balance(&self, address: &str) -> u64 {
        let key = format!("account:{address}");
        self.get::<AccountState>(&key)
            .ok()
            .flatten()
            .map(|a| a.balance)
            .unwrap_or(0)
    }

    fn get_nonce(&self, address: &str) -> u64 {
        let key = format!("account:{address}");
        self.get::<AccountState>(&key)
            .ok()
            .flatten()
            .map(|a| a.nonce)
            .unwrap_or(0)
    }

    fn set_balance(&self, address: &str, balance: u64) {
        let key = format!("account:{address}");
        let mut account = self
            .get::<AccountState>(&key)
            .ok()
            .flatten()
            .unwrap_or_else(|| AccountState::new(address.to_string()));
        account.balance = balance;
        let _ = self.put(&key, &account);
    }

    fn inc_nonce(&self, address: &str) {
        let key = format!("account:{address}");
        let mut account = self
            .get::<AccountState>(&key)
            .ok()
            .flatten()
            .unwrap_or_else(|| AccountState::new(address.to_string()));
        account.nonce += 1;
        let _ = self.put(&key, &account);
    }
}