//! State tree management for ZK-rollups.
//!
//! The state tree is a sparse Merkle tree that stores account states.
//! It supports efficient proofs of inclusion and non-inclusion.
//!
//! # Structure
//!
//! ```text
//!                     Root
//!                    /    \
//!                  H1      H2
//!                 /  \    /  \
//!               H3   H4  H5   H6
//!              / \
//!          Acc0  Acc1  ...  AccN  Empty
//! ```
//!
//! - Tree depth: 32 (supports 2^32 accounts)
//! - Hash function: Poseidon (ZK-friendly) or Blake3 (for off-chain)
//! - Empty leaves: Hash of zero

use blake3;
use parking_lot::RwLock;
use rs_merkle::{Hasher, MerkleTree};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use thiserror::Error;

/// State root - the Merkle root of the state tree.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct StateRoot(pub [u8; 32]);

impl StateRoot {
    /// Creates a state root from bytes.
    pub fn from_bytes(bytes: [u8; 32]) -> Self {
        Self(bytes)
    }

    /// Returns the root as bytes.
    pub fn as_bytes(&self) -> &[u8; 32] {
        &self.0
    }

    /// Returns a zero state root (empty tree).
    pub fn zero() -> Self {
        Self([0u8; 32])
    }

    /// Returns the hex representation.
    pub fn to_hex(&self) -> String {
        hex::encode(self.0)
    }
}

impl From<[u8; 32]> for StateRoot {
    fn from(bytes: [u8; 32]) -> Self {
        Self(bytes)
    }
}

impl Default for StateRoot {
    fn default() -> Self {
        Self::zero()
    }
}

/// Account state stored in the state tree.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AccountState {
    /// Account balance
    pub balance: u128,
    /// Account nonce (for replay protection)
    pub nonce: u64,
    /// Public key hash (for signature verification)
    pub pubkey_hash: [u8; 32],
    /// Additional data (e.g., storage root for smart contracts)
    pub data_hash: [u8; 32],
}

impl AccountState {
    /// Creates a new account state.
    pub fn new(pubkey_hash: [u8; 32]) -> Self {
        Self {
            balance: 0,
            nonce: 0,
            pubkey_hash,
            data_hash: [0u8; 32],
        }
    }

    /// Creates an account with initial balance.
    pub fn with_balance(pubkey_hash: [u8; 32], balance: u128) -> Self {
        Self {
            balance,
            nonce: 0,
            pubkey_hash,
            data_hash: [0u8; 32],
        }
    }

    /// Returns the hash of this account state.
    pub fn hash(&self) -> [u8; 32] {
        let mut hasher = blake3::Hasher::new();
        hasher.update(&self.balance.to_le_bytes());
        hasher.update(&self.nonce.to_le_bytes());
        hasher.update(&self.pubkey_hash);
        hasher.update(&self.data_hash);
        *hasher.finalize().as_bytes()
    }

    /// Serializes to bytes.
    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = Vec::with_capacity(16 + 8 + 32 + 32);
        bytes.extend(&self.balance.to_le_bytes());
        bytes.extend(&self.nonce.to_le_bytes());
        bytes.extend(&self.pubkey_hash);
        bytes.extend(&self.data_hash);
        bytes
    }

    /// Deserializes from bytes.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, StateError> {
        if bytes.len() < 16 + 8 + 32 + 32 {
            return Err(StateError::InvalidAccountData);
        }

        let balance = u128::from_le_bytes(bytes[0..16].try_into().unwrap());
        let nonce = u64::from_le_bytes(bytes[16..24].try_into().unwrap());
        let mut pubkey_hash = [0u8; 32];
        pubkey_hash.copy_from_slice(&bytes[24..56]);
        let mut data_hash = [0u8; 32];
        data_hash.copy_from_slice(&bytes[56..88]);

        Ok(Self {
            balance,
            nonce,
            pubkey_hash,
            data_hash,
        })
    }
}

impl Default for AccountState {
    fn default() -> Self {
        Self::new([0u8; 32])
    }
}

/// Account with index and state.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Account {
    /// Account index in the state tree
    pub index: u64,
    /// Account state
    pub state: AccountState,
}

impl Account {
    /// Creates a new account.
    pub fn new(index: u64, pubkey_hash: [u8; 32]) -> Self {
        Self {
            index,
            state: AccountState::new(pubkey_hash),
        }
    }

    /// Returns the account hash.
    pub fn hash(&self) -> [u8; 32] {
        self.state.hash()
    }
}

/// Sparse Merkle tree for account states.
pub struct StateTree {
    /// Tree depth
    depth: usize,
    /// Accounts by index
    accounts: RwLock<HashMap<u64, AccountState>>,
    /// Cached leaf hashes
    leaf_hashes: RwLock<HashMap<u64, [u8; 32]>>,
    /// Current root (cached)
    root: RwLock<StateRoot>,
}

impl StateTree {
    /// Creates a new empty state tree.
    pub fn new(depth: usize) -> Self {
        Self {
            depth,
            accounts: RwLock::new(HashMap::new()),
            leaf_hashes: RwLock::new(HashMap::new()),
            root: RwLock::new(StateRoot::zero()),
        }
    }

    /// Creates a state tree with default depth (32).
    pub fn with_default_depth() -> Self {
        Self::new(crate::constants::STATE_TREE_DEPTH)
    }

    /// Returns the tree depth.
    pub fn depth(&self) -> usize {
        self.depth
    }

    /// Returns the current state root.
    pub fn root(&self) -> StateRoot {
        *self.root.read()
    }

    /// Returns the number of accounts.
    pub fn account_count(&self) -> usize {
        self.accounts.read().len()
    }

    /// Gets an account by index.
    pub fn get_account(&self, index: u64) -> Option<AccountState> {
        self.accounts.read().get(&index).cloned()
    }

    /// Sets an account state.
    pub fn set_account(&self, index: u64, state: AccountState) -> Result<StateRoot, StateError> {
        // Update account
        let hash = state.hash();
        self.accounts.write().insert(index, state);
        self.leaf_hashes.write().insert(index, hash);

        // Recompute root
        let new_root = self.compute_root()?;
        *self.root.write() = new_root;

        Ok(new_root)
    }

    /// Updates an account with a function.
    pub fn update_account<F>(&self, index: u64, f: F) -> Result<StateRoot, StateError>
    where
        F: FnOnce(&mut AccountState),
    {
        let mut accounts = self.accounts.write();
        let state = accounts.entry(index).or_insert_with(AccountState::default);
        f(state);
        let hash = state.hash();
        drop(accounts);

        self.leaf_hashes.write().insert(index, hash);

        let new_root = self.compute_root()?;
        *self.root.write() = new_root;

        Ok(new_root)
    }

    /// Computes the Merkle root from all accounts.
    fn compute_root(&self) -> Result<StateRoot, StateError> {
        let leaf_hashes = self.leaf_hashes.read();

        if leaf_hashes.is_empty() {
            return Ok(StateRoot::zero());
        }

        // For a sparse Merkle tree, we compute the root differently
        // Here we use a simplified approach with rs_merkle

        let mut leaves: Vec<[u8; 32]> = leaf_hashes.values().copied().collect();

        // Pad to power of 2
        let target_size = (leaves.len() as u64).next_power_of_two() as usize;
        while leaves.len() < target_size {
            leaves.push([0u8; 32]);
        }

        // Create Merkle tree
        let tree: MerkleTree<Blake3Algorithm> = MerkleTree::from_leaves(&leaves);
        let root = tree.root().ok_or(StateError::EmptyTree)?;

        let mut root_bytes = [0u8; 32];
        root_bytes.copy_from_slice(&root);

        Ok(StateRoot(root_bytes))
    }

    /// Generates a Merkle proof for an account.
    pub fn generate_proof(&self, index: u64) -> Result<Vec<[u8; 32]>, StateError> {
        let leaf_hashes = self.leaf_hashes.read();

        if !leaf_hashes.contains_key(&index) {
            return Err(StateError::AccountNotFound(index));
        }

        let mut leaves: Vec<_> = leaf_hashes.iter().collect();
        leaves.sort_by_key(|(k, _)| *k);

        let leaf_values: Vec<[u8; 32]> = leaves.iter().map(|(_, v)| **v).collect();

        // Find the index of our account in the sorted list
        let proof_index = leaves
            .iter()
            .position(|(k, _)| **k == index)
            .ok_or(StateError::AccountNotFound(index))?;

        // Pad to power of 2
        let mut padded_leaves = leaf_values.clone();
        let target_size = (padded_leaves.len() as u64).next_power_of_two() as usize;
        while padded_leaves.len() < target_size {
            padded_leaves.push([0u8; 32]);
        }

        let tree: MerkleTree<Blake3Algorithm> = MerkleTree::from_leaves(&padded_leaves);
        let proof = tree.proof(&[proof_index]);

        Ok(proof.proof_hashes().iter().map(|h| {
            let mut arr = [0u8; 32];
            arr.copy_from_slice(h);
            arr
        }).collect())
    }

    /// Verifies a Merkle proof.
    pub fn verify_proof(
        root: &StateRoot,
        index: u64,
        account_hash: [u8; 32],
        proof: &[[u8; 32]],
    ) -> bool {
        // Simplified verification - in production would use proper sparse Merkle verification
        let mut current = account_hash;

        for sibling in proof {
            let mut combined = Vec::with_capacity(64);
            if index & 1 == 0 {
                combined.extend(&current);
                combined.extend(sibling);
            } else {
                combined.extend(sibling);
                combined.extend(&current);
            }
            current = *blake3::hash(&combined).as_bytes();
        }

        current == root.0
    }

    /// Applies a transfer between accounts.
    pub fn apply_transfer(
        &self,
        from_idx: u64,
        to_idx: u64,
        amount: u128,
    ) -> Result<StateRoot, StateError> {
        // Check sender balance
        {
            let accounts = self.accounts.read();
            let from = accounts
                .get(&from_idx)
                .ok_or(StateError::AccountNotFound(from_idx))?;
            if from.balance < amount {
                return Err(StateError::InsufficientBalance {
                    available: from.balance,
                    required: amount,
                });
            }
        }

        // Update sender
        self.update_account(from_idx, |acc| {
            acc.balance -= amount;
            acc.nonce += 1;
        })?;

        // Update recipient
        self.update_account(to_idx, |acc| {
            acc.balance += amount;
        })
    }

    /// Applies a deposit (increases balance).
    pub fn apply_deposit(&self, to_idx: u64, amount: u128) -> Result<StateRoot, StateError> {
        self.update_account(to_idx, |acc| {
            acc.balance += amount;
        })
    }

    /// Applies a withdrawal (decreases balance).
    pub fn apply_withdrawal(&self, from_idx: u64, amount: u128) -> Result<StateRoot, StateError> {
        {
            let accounts = self.accounts.read();
            let from = accounts
                .get(&from_idx)
                .ok_or(StateError::AccountNotFound(from_idx))?;
            if from.balance < amount {
                return Err(StateError::InsufficientBalance {
                    available: from.balance,
                    required: amount,
                });
            }
        }

        self.update_account(from_idx, |acc| {
            acc.balance -= amount;
            acc.nonce += 1;
        })
    }

    /// Clears all accounts and resets to empty state.
    pub fn clear(&self) {
        self.accounts.write().clear();
        self.leaf_hashes.write().clear();
        *self.root.write() = StateRoot::zero();
    }
}

impl std::fmt::Debug for StateTree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StateTree")
            .field("depth", &self.depth)
            .field("account_count", &self.accounts.read().len())
            .field("root", &self.root.read().to_hex())
            .finish()
    }
}

/// Blake3 algorithm for rs_merkle.
#[derive(Clone)]
pub struct Blake3Algorithm;

impl Hasher for Blake3Algorithm {
    type Hash = [u8; 32];

    fn hash(data: &[u8]) -> Self::Hash {
        *blake3::hash(data).as_bytes()
    }
}

/// State errors.
#[derive(Debug, Error)]
pub enum StateError {
    #[error("Account not found: {0}")]
    AccountNotFound(u64),

    #[error("Insufficient balance: available {available}, required {required}")]
    InsufficientBalance { available: u128, required: u128 },

    #[error("Invalid account data")]
    InvalidAccountData,

    #[error("Tree is empty")]
    EmptyTree,

    #[error("Proof verification failed")]
    ProofVerificationFailed,

    #[error("Invalid state transition: {0}")]
    InvalidTransition(String),
}

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

    #[test]
    fn test_account_state() {
        let pubkey_hash = [0xab; 32];
        let state = AccountState::with_balance(pubkey_hash, 1000);

        assert_eq!(state.balance, 1000);
        assert_eq!(state.nonce, 0);
        assert_eq!(state.pubkey_hash, pubkey_hash);
    }

    #[test]
    fn test_account_serialization() {
        let state = AccountState::with_balance([0xab; 32], 1000);
        let bytes = state.to_bytes();
        let decoded = AccountState::from_bytes(&bytes).unwrap();

        assert_eq!(decoded, state);
    }

    #[test]
    fn test_state_tree_creation() {
        let tree = StateTree::with_default_depth();
        assert_eq!(tree.depth(), 32);
        assert_eq!(tree.account_count(), 0);
        assert_eq!(tree.root(), StateRoot::zero());
    }

    #[test]
    fn test_set_account() {
        let tree = StateTree::with_default_depth();
        let state = AccountState::with_balance([0xab; 32], 1000);

        let root = tree.set_account(0, state.clone()).unwrap();
        assert_ne!(root, StateRoot::zero());

        let retrieved = tree.get_account(0).unwrap();
        assert_eq!(retrieved, state);
    }

    #[test]
    fn test_transfer() {
        let tree = StateTree::with_default_depth();

        // Create two accounts
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 500))
            .unwrap();

        // Transfer
        tree.apply_transfer(0, 1, 300).unwrap();

        // Check balances
        assert_eq!(tree.get_account(0).unwrap().balance, 700);
        assert_eq!(tree.get_account(1).unwrap().balance, 800);

        // Check nonce
        assert_eq!(tree.get_account(0).unwrap().nonce, 1);
    }

    #[test]
    fn test_insufficient_balance() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 100))
            .unwrap();
        tree.set_account(1, AccountState::with_balance([0xbb; 32], 0))
            .unwrap();

        let result = tree.apply_transfer(0, 1, 200);
        assert!(matches!(result, Err(StateError::InsufficientBalance { .. })));
    }

    #[test]
    fn test_deposit_withdrawal() {
        let tree = StateTree::with_default_depth();
        tree.set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();

        // Deposit
        tree.apply_deposit(0, 500).unwrap();
        assert_eq!(tree.get_account(0).unwrap().balance, 1500);

        // Withdrawal
        tree.apply_withdrawal(0, 300).unwrap();
        assert_eq!(tree.get_account(0).unwrap().balance, 1200);
    }

    #[test]
    fn test_state_root_changes() {
        let tree = StateTree::with_default_depth();

        let root1 = tree
            .set_account(0, AccountState::with_balance([0xaa; 32], 1000))
            .unwrap();
        let root2 = tree
            .set_account(1, AccountState::with_balance([0xbb; 32], 500))
            .unwrap();

        // Each state change should produce a different root
        assert_ne!(root1, root2);
    }
}
