// NVM-L1 状态管理

use crate::contract::Contract;
use crate::types::{Address, Hash};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// 账户状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AccountState {
    /// 地址
    pub address: Address,
    /// 余额
    pub holdings: u128,
    /// Nonce
    pub nonce: u64,
    /// 代码哈希 (如果是合约账户)
    pub code_hash: Option<Hash>,
    /// 存储根哈希
    pub storage_root: Hash,
}

impl AccountState {
    /// 创建新账户
    pub fn new(address: Address) -> Self {
        Self {
            address,
            holdings: 0,
            nonce: 0,
            code_hash: None,
            storage_root: Hash::zero(),
        }
    }

    /// 检查是否为合约账户
    pub fn is_contract(&self) -> bool {
        self.code_hash.is_some()
    }
}

/// 存储键值对
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageEntry {
    /// 键
    pub key: Hash,
    /// 值
    pub value: Vec<u8>,
}

/// 世界状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorldState {
    /// 账户映射 (address -> AccountState)
    accounts: HashMap<Address, AccountState>,
    /// 合约映射 (address -> Contract)
    contracts: HashMap<Address, Contract>,
    /// 存储映射 (address -> (key -> value))
    storage: HashMap<Address, HashMap<Hash, Vec<u8>>>,
    /// 状态根哈希
    state_root: Hash,
}

impl WorldState {
    /// 创建新的世界状态
    pub fn new() -> Self {
        Self {
            accounts: HashMap::new(),
            contracts: HashMap::new(),
            storage: HashMap::new(),
            state_root: Hash::zero(),
        }
    }

    /// 获取账户状态
    pub fn get_account(&self, address: &Address) -> Option<&AccountState> {
        self.accounts.get(address)
    }

    /// 获取可变账户状态
    pub fn get_account_mut(&mut self, address: &Address) -> Option<&mut AccountState> {
        self.accounts.get_mut(address)
    }

    /// 创建账户
    pub fn create_account(&mut self, address: Address) -> Result<(), String> {
        if self.accounts.contains_key(&address) {
            return Err("Account already exists".to_string());
        }
        self.accounts.insert(address, AccountState::new(address));
        Ok(())
    }

    /// 获取账户余额
    pub fn get_balance(&self, address: &Address) -> u128 {
        self.accounts
            .get(address)
            .map(|acc| acc.holdings)
            .unwrap_or(0)
    }

    /// 设置账户余额
    pub fn set_balance(&mut self, address: &Address, holdings: u128) -> Result<(), String> {
        let account = self
            .accounts
            .get_mut(address)
            .ok_or("Account not found")?;
        account.holdings = holdings;
        Ok(())
    }

    /// 转账
    pub fn transfer(
        &mut self,
        from: &Address,
        to: &Address,
        amount: u128,
    ) -> Result<(), String> {
        // 检查发送方余额
        let from_balance = self.get_balance(from);
        if from_balance < amount {
            return Err("Insufficient holdings".to_string());
        }

        // 确保接收方账户存在
        if !self.accounts.contains_key(to) {
            self.create_account(*to)?;
        }

        // 执行转账
        self.set_balance(from, from_balance - amount)?;
        let to_balance = self.get_balance(to);
        self.set_balance(to, to_balance + amount)?;

        Ok(())
    }

    /// 获取账户nonce
    pub fn get_nonce(&self, address: &Address) -> u64 {
        self.accounts
            .get(address)
            .map(|acc| acc.nonce)
            .unwrap_or(0)
    }

    /// 增加账户nonce
    pub fn increment_nonce(&mut self, address: &Address) -> Result<(), String> {
        let account = self
            .accounts
            .get_mut(address)
            .ok_or("Account not found")?;
        account.nonce = account.nonce.saturating_add(1);
        Ok(())
    }

    /// 部署合约
    pub fn deploy_contract(&mut self, contract: Contract) -> Result<(), String> {
        let address = contract.address;

        // 检查地址是否已被占用
        if self.contracts.contains_key(&address) {
            return Err("Contract already exists at this address".to_string());
        }

        // 创建合约账户
        if !self.accounts.contains_key(&address) {
            self.create_account(address)?;
        }

        // 设置代码哈希
        let account = self.accounts.get_mut(&address).unwrap();
        account.code_hash = Some(*contract.code.code_hash());

        // 保存合约
        self.contracts.insert(address, contract);

        Ok(())
    }

    /// 获取合约
    pub fn get_contract(&self, address: &Address) -> Option<&Contract> {
        self.contracts.get(address)
    }

    /// 获取可变合约
    pub fn get_contract_mut(&mut self, address: &Address) -> Option<&mut Contract> {
        self.contracts.get_mut(address)
    }

    /// 读取存储
    pub fn get_storage(&self, address: &Address, key: &Hash) -> Option<Vec<u8>> {
        self.storage
            .get(address)
            .and_then(|storage| storage.get(key))
            .cloned()
    }

    /// 写入存储
    pub fn set_storage(&mut self, address: &Address, key: Hash, value: Vec<u8>) {
        self.storage
            .entry(*address)
            .or_insert_with(HashMap::new)
            .insert(key, value);
    }

    /// 删除存储
    pub fn remove_storage(&mut self, address: &Address, key: &Hash) {
        if let Some(storage) = self.storage.get_mut(address) {
            storage.remove(key);
        }
    }

    /// 获取状态根
    pub fn state_root(&self) -> &Hash {
        &self.state_root
    }

    /// 更新状态根
    pub fn update_state_root(&mut self) {
        // 简化版本：计算所有账户状态的哈希
        // 实际实现应该使用Merkle Patricia Tree
        let mut data = Vec::new();
        for (addr, account) in &self.accounts {
            data.extend_from_slice(addr.as_bytes());
            data.extend_from_slice(&account.holdings.to_le_bytes());
            data.extend_from_slice(&account.nonce.to_le_bytes());
        }
        self.state_root = Hash::sha3_384(&data);
    }
}

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

/// 状态管理器
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateManager {
    /// 当前世界状态
    current_state: WorldState,
    /// 状态快照栈
    snapshots: Vec<WorldState>,
}

impl StateManager {
    /// 创建新的状态管理器
    pub fn new() -> Self {
        Self {
            current_state: WorldState::new(),
            snapshots: Vec::new(),
        }
    }

    /// 获取当前状态
    pub fn current_state(&self) -> &WorldState {
        &self.current_state
    }

    /// 获取可变当前状态
    pub fn current_state_mut(&mut self) -> &mut WorldState {
        &mut self.current_state
    }

    /// 创建快照
    pub fn snapshot(&mut self) {
        self.snapshots.push(self.current_state.clone());
    }

    /// 恢复快照
    pub fn revert(&mut self) -> Result<(), String> {
        let snapshot = self.snapshots.pop().ok_or("No snapshot to revert")?;
        self.current_state = snapshot;
        Ok(())
    }

    /// 提交当前状态
    pub fn commit(&mut self) {
        self.snapshots.clear();
        self.current_state.update_state_root();
    }
}

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

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

    #[test]
    fn test_account_creation() {
        let mut state = WorldState::new();
        let addr = Address::new([1u8; 20]);
        assert!(state.create_account(addr).is_ok());
        assert!(state.get_account(&addr).is_some());
    }

    #[test]
    fn test_balance_operations() {
        let mut state = WorldState::new();
        let addr = Address::new([1u8; 20]);
        state.create_account(addr).unwrap();

        assert!(state.set_balance(&addr, 1000).is_ok());
        assert_eq!(state.get_balance(&addr), 1000);
    }

    #[test]
    fn test_transfer() {
        let mut state = WorldState::new();
        let from = Address::new([1u8; 20]);
        let to = Address::new([2u8; 20]);

        state.create_account(from).unwrap();
        state.set_balance(&from, 1000).unwrap();

        assert!(state.transfer(&from, &to, 500).is_ok());
        assert_eq!(state.get_balance(&from), 500);
        assert_eq!(state.get_balance(&to), 500);
    }

    #[test]
    fn test_nonce() {
        let mut state = WorldState::new();
        let addr = Address::new([1u8; 20]);
        state.create_account(addr).unwrap();

        assert_eq!(state.get_nonce(&addr), 0);
        state.increment_nonce(&addr).unwrap();
        assert_eq!(state.get_nonce(&addr), 1);
    }

    #[test]
    fn test_storage() {
        let mut state = WorldState::new();
        let addr = Address::new([1u8; 20]);
        let key = Hash::new([1u8; 32]);
        let value = vec![1, 2, 3, 4];

        state.set_storage(&addr, key, value.clone());
        assert_eq!(state.get_storage(&addr, &key), Some(value));
    }

    #[test]
    fn test_state_manager_snapshot() {
        let mut manager = StateManager::new();
        let addr = Address::new([1u8; 20]);

        manager.current_state_mut().create_account(addr).unwrap();
        manager.current_state_mut().set_balance(&addr, 1000).unwrap();

        manager.snapshot();

        manager.current_state_mut().set_balance(&addr, 500).unwrap();
        assert_eq!(manager.current_state().get_balance(&addr), 500);

        manager.revert().unwrap();
        assert_eq!(manager.current_state().get_balance(&addr), 1000);
    }
}