//! Test environment for contract testing.
//!
//! Provides a complete testing environment that wraps the VM engine
//! with test utilities for deploying and interacting with contracts.

use std::collections::HashMap;
use std::sync::Arc;

use parking_lot::RwLock;
use synor_types::{Address, Hash256, Network};
use synor_vm::{
    context::{BlockInfo, CallContext, ExecutionContext, TransactionInfo},
    engine::ContractModule,
    storage::{ContractStorage, MemoryStorage},
    ContractId, ExecutionResult, StorageKey, StorageValue, VmEngine,
};

use crate::{MockStorage, TestAccount, TestError, TestResult};

/// Configuration for the test environment.
#[derive(Clone, Debug)]
pub struct TestEnvironmentConfig {
    /// Default gas limit for calls.
    pub default_gas_limit: u64,
    /// Chain ID.
    pub chain_id: u64,
    /// Network type.
    pub network: Network,
    /// Whether to auto-commit storage changes.
    pub auto_commit: bool,
    /// Block time advancement per call (ms).
    pub block_time_ms: u64,
}

impl Default for TestEnvironmentConfig {
    fn default() -> Self {
        TestEnvironmentConfig {
            default_gas_limit: 10_000_000,
            chain_id: 31337, // Common test chain ID
            network: Network::Devnet,
            auto_commit: true,
            block_time_ms: 1000,
        }
    }
}

/// The main test environment for contract testing.
///
/// Provides a complete sandbox for testing smart contracts including:
/// - Contract deployment and execution
/// - Account management with balances
/// - Block and time simulation
/// - Storage inspection
///
/// # Example
///
/// ```rust,ignore
/// let mut env = TestEnvironment::builder()
///     .with_account(TestAccount::with_balance(1000))
///     .with_account(TestAccount::with_balance(500))
///     .build();
///
/// let contract = env.deploy_contract(&bytecode, &init_params)?;
/// let result = env.call_contract(contract, "transfer", &args)?;
/// ```
pub struct TestEnvironment {
    /// The VM engine.
    engine: VmEngine,
    /// Mock storage.
    storage: Arc<RwLock<MockStorage>>,
    /// Deployed contracts.
    contracts: HashMap<ContractId, Arc<ContractModule>>,
    /// Test accounts by address.
    accounts: HashMap<Address, TestAccount>,
    /// Account balances (separate tracking).
    balances: HashMap<Address, u64>,
    /// Current block info.
    block: BlockInfo,
    /// Configuration.
    config: TestEnvironmentConfig,
    /// Transaction counter (for unique tx hashes).
    tx_counter: u64,
}

impl TestEnvironment {
    /// Creates a new test environment with default settings.
    pub fn new() -> TestResult<Self> {
        let engine = VmEngine::new().map_err(TestError::VmError)?;
        let coinbase = Address::from_ed25519_pubkey(Network::Devnet, &[0; 32]);

        Ok(TestEnvironment {
            engine,
            storage: Arc::new(RwLock::new(MockStorage::new())),
            contracts: HashMap::new(),
            accounts: HashMap::new(),
            balances: HashMap::new(),
            block: BlockInfo {
                height: 1,
                timestamp: 1704067200000, // Jan 1, 2024 00:00:00 UTC
                hash: Hash256::default(),
                blue_score: 1,
                daa_score: 1,
                coinbase,
            },
            config: TestEnvironmentConfig::default(),
            tx_counter: 0,
        })
    }

    /// Creates a builder for configuring the test environment.
    pub fn builder() -> TestEnvironmentBuilder {
        TestEnvironmentBuilder::new()
    }

    /// Deploys a contract and returns its ID.
    ///
    /// # Arguments
    ///
    /// * `bytecode` - The WASM bytecode of the contract
    /// * `init_params` - Initialization parameters (serialized)
    ///
    /// # Returns
    ///
    /// The contract ID on success.
    pub fn deploy_contract(
        &mut self,
        bytecode: &[u8],
        init_params: &[u8],
    ) -> TestResult<ContractId> {
        self.deploy_contract_as(bytecode, init_params, None, 0)
    }

    /// Deploys a contract as a specific account.
    pub fn deploy_contract_as(
        &mut self,
        bytecode: &[u8],
        init_params: &[u8],
        deployer: Option<&Address>,
        value: u64,
    ) -> TestResult<ContractId> {
        // Compile the contract
        let module = self
            .engine
            .compile(bytecode.to_vec())
            .map_err(|e| TestError::InvalidBytecode(e.to_string()))?;

        let contract_id = module.id;

        // Get deployer address
        let deployer_addr = deployer
            .cloned()
            .unwrap_or_else(|| Address::from_ed25519_pubkey(self.config.network, &[0; 32]));

        // Check balance if value > 0
        if value > 0 {
            let balance = self.balances.get(&deployer_addr).copied().unwrap_or(0);
            if balance < value {
                return Err(TestError::InsufficientBalance {
                    have: balance,
                    need: value,
                });
            }
        }

        // Create execution context
        let call = CallContext::new(contract_id, deployer_addr.clone(), value, init_params.to_vec());

        let mut storage = MemoryStorage::new();
        // Copy existing storage for this contract if any
        for (key, val) in self.storage.read().entries(&contract_id) {
            storage.set(&contract_id, key, val);
        }
        storage.commit();

        let context = ExecutionContext::new(
            self.block.clone(),
            self.make_tx_info(&deployer_addr),
            call,
            self.config.default_gas_limit,
            storage,
            self.config.chain_id,
        );

        // Initialize the contract
        let result = self
            .engine
            .execute(
                &module,
                "__synor_init",
                init_params,
                context,
                self.config.default_gas_limit,
            )
            .map_err(|e| TestError::DeploymentFailed(e.to_string()))?;

        // Apply storage changes
        if self.config.auto_commit {
            self.apply_storage_changes(&result);
        }

        // Debit value from deployer
        if value > 0 {
            if let Some(balance) = self.balances.get_mut(&deployer_addr) {
                *balance -= value;
            }
        }

        // Cache the compiled module
        self.contracts.insert(contract_id, Arc::new(module));

        // Advance block
        self.advance_block_internal();

        Ok(contract_id)
    }

    /// Calls a method on a deployed contract.
    ///
    /// # Arguments
    ///
    /// * `contract` - The contract ID
    /// * `method` - The method name to call
    /// * `args` - Method arguments (serialized)
    ///
    /// # Returns
    ///
    /// The execution result.
    pub fn call_contract(
        &mut self,
        contract: ContractId,
        method: &str,
        args: &[u8],
    ) -> TestResult<ExecutionResult> {
        self.call_contract_as(contract, method, args, None, 0)
    }

    /// Calls a method on a contract as a specific account.
    pub fn call_contract_as(
        &mut self,
        contract: ContractId,
        method: &str,
        args: &[u8],
        caller: Option<&Address>,
        value: u64,
    ) -> TestResult<ExecutionResult> {
        let module = self
            .contracts
            .get(&contract)
            .ok_or_else(|| TestError::CallFailed(format!("Contract not found: {}", contract)))?
            .clone();

        let caller_addr = caller
            .cloned()
            .unwrap_or_else(|| Address::from_ed25519_pubkey(self.config.network, &[1; 32]));

        // Check balance
        if value > 0 {
            let balance = self.balances.get(&caller_addr).copied().unwrap_or(0);
            if balance < value {
                return Err(TestError::InsufficientBalance {
                    have: balance,
                    need: value,
                });
            }
        }

        // Build call data (method selector + args)
        let call_data = crate::encode_call_args(method, args);

        // Create call context
        let call = CallContext::new(contract, caller_addr.clone(), value, call_data.clone());

        // Prepare storage
        let mut storage = MemoryStorage::new();
        for (key, val) in self.storage.read().entries(&contract) {
            storage.set(&contract, key, val);
        }
        storage.commit();

        let context = ExecutionContext::new(
            self.block.clone(),
            self.make_tx_info(&caller_addr),
            call,
            self.config.default_gas_limit,
            storage,
            self.config.chain_id,
        );

        // Execute
        let result = self
            .engine
            .execute(
                &module,
                "__synor_call",
                &call_data,
                context,
                self.config.default_gas_limit,
            )
            .map_err(TestError::VmError)?;

        // Apply storage changes
        if self.config.auto_commit {
            self.apply_storage_changes(&result);
        }

        // Debit value from caller
        if value > 0 {
            if let Some(balance) = self.balances.get_mut(&caller_addr) {
                *balance -= value;
            }
        }

        // Advance block
        self.advance_block_internal();

        Ok(result)
    }

    /// Calls a contract method as a static call (read-only).
    pub fn static_call(
        &self,
        contract: ContractId,
        method: &str,
        args: &[u8],
    ) -> TestResult<ExecutionResult> {
        let module = self
            .contracts
            .get(&contract)
            .ok_or_else(|| TestError::CallFailed(format!("Contract not found: {}", contract)))?
            .clone();

        let caller_addr = Address::from_ed25519_pubkey(self.config.network, &[0; 32]);
        let call_data = crate::encode_call_args(method, args);

        // Create static call context
        let call = CallContext::static_call(contract, caller_addr.clone(), call_data.clone());

        // Prepare read-only storage view
        let mut storage = MemoryStorage::new();
        for (key, val) in self.storage.read().entries(&contract) {
            storage.set(&contract, key, val);
        }
        storage.commit();

        // For static calls, use a dummy tx info since we don't mutate state
        let tx_info = TransactionInfo {
            hash: Hash256::default(),
            origin: caller_addr.clone(),
            gas_price: 1,
            gas_limit: self.config.default_gas_limit,
        };

        let context = ExecutionContext::new(
            self.block.clone(),
            tx_info,
            call,
            self.config.default_gas_limit,
            storage,
            self.config.chain_id,
        );

        self.engine
            .execute(
                &module,
                "__synor_call",
                &call_data,
                context,
                self.config.default_gas_limit,
            )
            .map_err(TestError::VmError)
    }

    /// Sets the balance of an address.
    pub fn set_balance(&mut self, address: &Address, amount: u64) {
        self.balances.insert(address.clone(), amount);
    }

    /// Gets the balance of an address.
    pub fn get_balance(&self, address: &Address) -> u64 {
        self.balances.get(address).copied().unwrap_or(0)
    }

    /// Transfers balance between addresses.
    pub fn transfer(&mut self, from: &Address, to: &Address, amount: u64) -> TestResult<()> {
        let from_balance = self.get_balance(from);
        if from_balance < amount {
            return Err(TestError::InsufficientBalance {
                have: from_balance,
                need: amount,
            });
        }

        self.set_balance(from, from_balance - amount);
        let to_balance = self.get_balance(to);
        self.set_balance(to, to_balance + amount);

        Ok(())
    }

    /// Advances the block number by n blocks.
    pub fn advance_block(&mut self, n: u64) {
        for _ in 0..n {
            self.advance_block_internal();
        }
    }

    /// Gets a storage value for a contract.
    pub fn get_storage(&self, contract: &ContractId, key: &StorageKey) -> Option<StorageValue> {
        self.storage.read().get(contract, key)
    }

    /// Sets a storage value directly (for test setup).
    pub fn set_storage(&self, contract: &ContractId, key: StorageKey, value: StorageValue) {
        self.storage.write().set_direct(contract, key, value);
    }

    /// Checks if a contract has any storage.
    pub fn has_storage(&self, contract: &ContractId) -> bool {
        !self.storage.read().is_empty(contract)
    }

    /// Gets all storage entries for a contract.
    pub fn get_all_storage(&self, contract: &ContractId) -> Vec<(StorageKey, StorageValue)> {
        self.storage.read().entries(contract)
    }

    /// Clears all storage for a contract.
    pub fn clear_storage(&self, contract: &ContractId) {
        self.storage.read().clear_contract(contract);
    }

    /// Takes a storage snapshot.
    pub fn snapshot(&self) -> crate::mock_storage::StorageSnapshot {
        self.storage.read().snapshot()
    }

    /// Restores from a storage snapshot.
    pub fn restore(&self, snapshot: crate::mock_storage::StorageSnapshot) {
        self.storage.read().restore(snapshot);
    }

    /// Gets the current block height.
    pub fn block_height(&self) -> u64 {
        self.block.height
    }

    /// Gets the current block timestamp.
    pub fn block_timestamp(&self) -> u64 {
        self.block.timestamp
    }

    /// Sets the block timestamp.
    pub fn set_timestamp(&mut self, timestamp: u64) {
        self.block.timestamp = timestamp;
    }

    /// Advances time by milliseconds.
    pub fn advance_time(&mut self, ms: u64) {
        self.block.timestamp += ms;
    }

    /// Registers a test account.
    pub fn register_account(&mut self, account: TestAccount) {
        let address = account.address();
        self.balances.insert(address.clone(), account.balance());
        self.accounts.insert(address, account);
    }

    /// Gets a registered account by address.
    pub fn get_account(&self, address: &Address) -> Option<&TestAccount> {
        self.accounts.get(address)
    }

    /// Gets a mutable reference to a registered account.
    pub fn get_account_mut(&mut self, address: &Address) -> Option<&mut TestAccount> {
        self.accounts.get_mut(address)
    }

    /// Gets all registered accounts.
    pub fn accounts(&self) -> &HashMap<Address, TestAccount> {
        &self.accounts
    }

    /// Gets the configuration.
    pub fn config(&self) -> &TestEnvironmentConfig {
        &self.config
    }

    /// Sets the default gas limit.
    pub fn set_gas_limit(&mut self, limit: u64) {
        self.config.default_gas_limit = limit;
    }

    /// Gets storage change history.
    pub fn storage_history(&self) -> Vec<synor_vm::StorageChange> {
        self.storage.read().get_history()
    }

    /// Clears storage history.
    pub fn clear_storage_history(&self) {
        self.storage.read().clear_history();
    }

    // Internal helpers

    fn advance_block_internal(&mut self) {
        self.block.height += 1;
        self.block.timestamp += self.config.block_time_ms;
        self.block.blue_score += 1;
        self.block.daa_score += 1;
        // Update block hash
        let mut hash_input = Vec::new();
        hash_input.extend_from_slice(&self.block.height.to_le_bytes());
        hash_input.extend_from_slice(&self.block.timestamp.to_le_bytes());
        self.block.hash = Hash256::from_bytes(blake3::hash(&hash_input).into());
    }

    fn make_tx_info(&mut self, origin: &Address) -> TransactionInfo {
        self.tx_counter += 1;
        let mut hash_input = self.tx_counter.to_le_bytes().to_vec();
        hash_input.extend_from_slice(origin.payload());

        TransactionInfo {
            hash: Hash256::from_bytes(blake3::hash(&hash_input).into()),
            origin: origin.clone(),
            gas_price: 1,
            gas_limit: self.config.default_gas_limit,
        }
    }

    fn apply_storage_changes(&self, result: &ExecutionResult) {
        let storage = self.storage.write();
        for change in &result.storage_changes {
            match &change.new_value {
                Some(value) => {
                    storage.set_direct(&change.contract, change.key, value.clone());
                }
                None => {
                    storage.delete_direct(&change.contract, &change.key);
                }
            }
        }
    }
}

impl Default for TestEnvironment {
    fn default() -> Self {
        Self::new().expect("Failed to create default TestEnvironment")
    }
}

/// Builder for creating test environments with custom configuration.
#[derive(Default)]
pub struct TestEnvironmentBuilder {
    accounts: Vec<TestAccount>,
    config: TestEnvironmentConfig,
    initial_storage: HashMap<ContractId, Vec<(StorageKey, StorageValue)>>,
    initial_timestamp: Option<u64>,
    initial_block_height: Option<u64>,
}

impl TestEnvironmentBuilder {
    /// Creates a new builder with default settings.
    pub fn new() -> Self {
        TestEnvironmentBuilder {
            accounts: Vec::new(),
            config: TestEnvironmentConfig::default(),
            initial_storage: HashMap::new(),
            initial_timestamp: None,
            initial_block_height: None,
        }
    }

    /// Adds a test account.
    pub fn with_account(mut self, account: TestAccount) -> Self {
        self.accounts.push(account);
        self
    }

    /// Adds multiple test accounts.
    pub fn with_accounts(mut self, accounts: Vec<TestAccount>) -> Self {
        self.accounts.extend(accounts);
        self
    }

    /// Sets the default gas limit.
    pub fn with_gas_limit(mut self, limit: u64) -> Self {
        self.config.default_gas_limit = limit;
        self
    }

    /// Sets the chain ID.
    pub fn with_chain_id(mut self, chain_id: u64) -> Self {
        self.config.chain_id = chain_id;
        self
    }

    /// Sets the network.
    pub fn with_network(mut self, network: Network) -> Self {
        self.config.network = network;
        self
    }

    /// Sets auto-commit behavior.
    pub fn with_auto_commit(mut self, auto_commit: bool) -> Self {
        self.config.auto_commit = auto_commit;
        self
    }

    /// Sets the block time.
    pub fn with_block_time(mut self, ms: u64) -> Self {
        self.config.block_time_ms = ms;
        self
    }

    /// Sets initial storage for a contract.
    pub fn with_storage(
        mut self,
        contract: ContractId,
        entries: Vec<(StorageKey, StorageValue)>,
    ) -> Self {
        self.initial_storage.insert(contract, entries);
        self
    }

    /// Sets the initial timestamp.
    pub fn with_timestamp(mut self, timestamp: u64) -> Self {
        self.initial_timestamp = Some(timestamp);
        self
    }

    /// Sets the initial block height.
    pub fn with_block_height(mut self, height: u64) -> Self {
        self.initial_block_height = Some(height);
        self
    }

    /// Builds the test environment.
    pub fn build(self) -> TestEnvironment {
        let mut env = TestEnvironment::new().expect("Failed to create TestEnvironment");
        env.config = self.config;

        // Register accounts
        for account in self.accounts {
            env.register_account(account);
        }

        // Set initial storage
        for (contract, entries) in self.initial_storage {
            for (key, value) in entries {
                env.set_storage(&contract, key, value);
            }
        }

        // Set initial timestamp if specified
        if let Some(timestamp) = self.initial_timestamp {
            env.block.timestamp = timestamp;
        }

        // Set initial block height if specified
        if let Some(height) = self.initial_block_height {
            env.block.height = height;
        }

        env
    }
}

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

    #[test]
    fn test_environment_creation() {
        let env = TestEnvironment::new().unwrap();
        assert_eq!(env.block_height(), 1);
        assert!(env.block_timestamp() > 0);
    }

    #[test]
    fn test_environment_builder() {
        let alice = predefined::alice().set_balance(1000);
        let bob = predefined::bob().set_balance(500);

        let env = TestEnvironment::builder()
            .with_account(alice.clone())
            .with_account(bob.clone())
            .with_gas_limit(5_000_000)
            .with_chain_id(12345)
            .build();

        assert_eq!(env.get_balance(&alice.address()), 1000);
        assert_eq!(env.get_balance(&bob.address()), 500);
        assert_eq!(env.config().default_gas_limit, 5_000_000);
        assert_eq!(env.config().chain_id, 12345);
    }

    #[test]
    fn test_balance_operations() {
        let mut env = TestEnvironment::new().unwrap();
        let alice = predefined::alice();
        let bob = predefined::bob();

        env.set_balance(&alice.address(), 1000);
        assert_eq!(env.get_balance(&alice.address()), 1000);

        env.transfer(&alice.address(), &bob.address(), 300).unwrap();
        assert_eq!(env.get_balance(&alice.address()), 700);
        assert_eq!(env.get_balance(&bob.address()), 300);

        // Insufficient balance
        let result = env.transfer(&alice.address(), &bob.address(), 1000);
        assert!(result.is_err());
    }

    #[test]
    fn test_block_advancement() {
        let mut env = TestEnvironment::new().unwrap();
        let initial_height = env.block_height();
        let initial_timestamp = env.block_timestamp();

        env.advance_block(5);

        assert_eq!(env.block_height(), initial_height + 5);
        assert_eq!(
            env.block_timestamp(),
            initial_timestamp + 5 * env.config().block_time_ms
        );
    }

    #[test]
    fn test_time_manipulation() {
        let mut env = TestEnvironment::new().unwrap();

        env.set_timestamp(2000000000000);
        assert_eq!(env.block_timestamp(), 2000000000000);

        env.advance_time(5000);
        assert_eq!(env.block_timestamp(), 2000000005000);
    }

    #[test]
    fn test_storage_operations() {
        let env = TestEnvironment::new().unwrap();
        let contract = ContractId::from_bytes([0x42; 32]);
        let key = StorageKey::from_bytes(b"test_key");
        let value = StorageValue::from_u64(12345);

        env.set_storage(&contract, key, value.clone());
        assert_eq!(env.get_storage(&contract, &key), Some(value));
        assert!(env.has_storage(&contract));

        env.clear_storage(&contract);
        assert!(!env.has_storage(&contract));
    }

    #[test]
    fn test_storage_snapshot() {
        let env = TestEnvironment::new().unwrap();
        let contract = ContractId::from_bytes([0x42; 32]);
        let key = StorageKey::from_bytes(b"test_key");

        env.set_storage(&contract, key, StorageValue::from_u64(100));
        let snapshot = env.snapshot();

        env.set_storage(&contract, key, StorageValue::from_u64(200));
        assert_eq!(
            env.get_storage(&contract, &key),
            Some(StorageValue::from_u64(200))
        );

        env.restore(snapshot);
        assert_eq!(
            env.get_storage(&contract, &key),
            Some(StorageValue::from_u64(100))
        );
    }
}