//! WASM execution engine using wasmtime.
//!
//! The engine compiles and executes WASM smart contracts with:
//! - JIT compilation for performance
//! - Fuel-based execution limits
//! - Host function imports

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

use parking_lot::RwLock;
use wasmtime::{
    Config, Engine, Instance, Linker, Memory, Module, Store, StoreLimits, StoreLimitsBuilder,
};

use crate::context::ExecutionContext;
use crate::gas::GasMeter;
use crate::storage::{ContractStorage, MemoryStorage};
use crate::{ContractId, ExecutionResult, VmError, MAX_CONTRACT_SIZE, MAX_MEMORY_PAGES};

/// Compiled contract module.
#[derive(Clone)]
pub struct ContractModule {
    /// Module hash (contract ID).
    pub id: ContractId,
    /// Original bytecode.
    pub bytecode: Vec<u8>,
    /// Compiled wasmtime module.
    module: Module,
}

impl ContractModule {
    /// Compiles a contract from WASM bytecode.
    pub fn compile(engine: &Engine, bytecode: Vec<u8>) -> Result<Self, VmError> {
        if bytecode.len() > MAX_CONTRACT_SIZE {
            return Err(VmError::BytecodeTooLarge {
                size: bytecode.len(),
                max: MAX_CONTRACT_SIZE,
            });
        }

        // Compute contract ID from bytecode hash
        let hash: [u8; 32] = blake3::hash(&bytecode).into();
        let id = ContractId::from_bytes(hash);

        // Compile the module
        let module =
            Module::new(engine, &bytecode).map_err(|e| VmError::InvalidBytecode(e.to_string()))?;

        Ok(ContractModule {
            id,
            bytecode,
            module,
        })
    }

    /// Returns the module's exports.
    pub fn exports(&self) -> impl Iterator<Item = (&str, wasmtime::ExternType)> + '_ {
        self.module.exports().map(|e| (e.name(), e.ty()))
    }

    /// Checks if the module exports a function.
    pub fn has_export(&self, name: &str) -> bool {
        self.module
            .exports()
            .any(|e| e.name() == name && matches!(e.ty(), wasmtime::ExternType::Func(_)))
    }
}

/// Store data for WASM execution.
pub struct StoreData {
    /// Gas meter.
    pub gas: GasMeter,
    /// Execution context.
    pub context: ExecutionContext<MemoryStorage>,
    /// Store limits.
    pub limits: StoreLimits,
    /// Return data buffer.
    pub return_data: Vec<u8>,
    /// Error message if reverted.
    pub error: Option<String>,
}

/// Instance of a running contract.
pub struct VmInstance {
    /// The store containing execution state.
    store: Store<StoreData>,
    /// The instantiated module.
    instance: Instance,
    /// Memory export.
    memory: Memory,
}

impl VmInstance {
    /// Calls a function on the contract.
    pub fn call(&mut self, method: &str, args: &[u8]) -> Result<ExecutionResult, VmError> {
        // Get the function export
        let func = self
            .instance
            .get_typed_func::<(i32, i32), i32>(&mut self.store, method)
            .map_err(|e| VmError::InvalidMethod(e.to_string()))?;

        // Write args to memory
        let args_ptr = self.write_to_memory(args)?;
        let args_len = args.len() as i32;

        // Call the function
        let result = func
            .call(&mut self.store, (args_ptr, args_len))
            .map_err(|e| {
                // Check for gas exhaustion
                if self.store.data().gas.remaining() == 0 {
                    VmError::OutOfGas {
                        used: self.store.data().gas.consumed(),
                        limit: self.store.data().gas.limit(),
                    }
                } else if let Some(err) = &self.store.data().error {
                    VmError::Revert(err.clone())
                } else {
                    VmError::ExecutionError(e.to_string())
                }
            })?;

        // Read return data
        let return_data = if result >= 0 {
            self.read_from_memory(result as u32, self.store.data().return_data.len())?
        } else {
            Vec::new()
        };

        // Extract storage changes from the pending storage operations
        let storage_changes = self.store.data().context.storage.pending_changes();

        Ok(ExecutionResult {
            return_data,
            gas_used: self.store.data().gas.consumed(),
            logs: self.store.data().context.logs.clone(),
            storage_changes,
            internal_calls: Vec::new(),
        })
    }

    /// Writes data to WASM memory.
    fn write_to_memory(&mut self, data: &[u8]) -> Result<i32, VmError> {
        let memory = self.memory;
        let data_ptr = memory.data_size(&self.store) as i32;

        // Grow memory if needed
        let pages_needed = ((data_ptr as usize + data.len()) / 65536) + 1;
        let current_pages = memory.size(&self.store) as usize;
        if pages_needed > current_pages {
            memory
                .grow(&mut self.store, (pages_needed - current_pages) as u64)
                .map_err(|e| VmError::MemoryViolation(e.to_string()))?;
        }

        // Write data
        memory.data_mut(&mut self.store)[data_ptr as usize..data_ptr as usize + data.len()]
            .copy_from_slice(data);

        Ok(data_ptr)
    }

    /// Reads data from WASM memory.
    fn read_from_memory(&self, ptr: u32, len: usize) -> Result<Vec<u8>, VmError> {
        let memory = self.memory;
        let data = memory.data(&self.store);

        if ptr as usize + len > data.len() {
            return Err(VmError::MemoryViolation(format!(
                "Read out of bounds: {} + {} > {}",
                ptr,
                len,
                data.len()
            )));
        }

        Ok(data[ptr as usize..ptr as usize + len].to_vec())
    }

    /// Returns gas consumed.
    pub fn gas_consumed(&self) -> u64 {
        self.store.data().gas.consumed()
    }

    /// Returns remaining gas.
    pub fn gas_remaining(&self) -> u64 {
        self.store.data().gas.remaining()
    }
}

/// The main VM engine.
pub struct VmEngine {
    /// Wasmtime engine with configuration.
    engine: Engine,
    /// Compiled contract cache.
    modules: RwLock<HashMap<ContractId, Arc<ContractModule>>>,
    /// Linker with host functions.
    linker: Linker<StoreData>,
}

impl VmEngine {
    /// Creates a new VM engine.
    pub fn new() -> Result<Self, VmError> {
        let mut config = Config::new();
        config.consume_fuel(true); // Enable fuel-based execution limits
        config.wasm_bulk_memory(true);
        config.wasm_multi_value(true);
        config.wasm_reference_types(true);
        config.cranelift_opt_level(wasmtime::OptLevel::Speed);

        let engine = Engine::new(&config).map_err(|e| VmError::ExecutionError(e.to_string()))?;

        let mut linker = Linker::new(&engine);

        // Register host functions
        Self::register_host_functions(&mut linker)?;

        Ok(VmEngine {
            engine,
            modules: RwLock::new(HashMap::new()),
            linker,
        })
    }

    /// Registers host functions with the linker.
    fn register_host_functions(linker: &mut Linker<StoreData>) -> Result<(), VmError> {
        // Storage read
        linker
            .func_wrap(
                "env",
                "synor_storage_read",
                |mut caller: wasmtime::Caller<'_, StoreData>, key_ptr: i32, out_ptr: i32| -> i32 {
                    // Read key from memory
                    let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
                    let mut key = [0u8; 32];
                    key.copy_from_slice(
                        &memory.data(&caller)[key_ptr as usize..key_ptr as usize + 32],
                    );

                    // Read from storage
                    let data = caller.data_mut();
                    let storage_key = crate::storage::StorageKey::new(key);
                    match data
                        .context
                        .storage
                        .get(&data.context.call.contract, &storage_key)
                    {
                        Some(value) => {
                            let len = value.0.len();
                            // Write to output buffer
                            memory.data_mut(&mut caller)[out_ptr as usize..out_ptr as usize + len]
                                .copy_from_slice(&value.0);
                            len as i32
                        }
                        None => -1,
                    }
                },
            )
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        // Storage write
        linker
            .func_wrap(
                "env",
                "synor_storage_write",
                |mut caller: wasmtime::Caller<'_, StoreData>,
                 key_ptr: i32,
                 value_ptr: i32,
                 value_len: i32|
                 -> i32 {
                    let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
                    let mut key = [0u8; 32];
                    key.copy_from_slice(
                        &memory.data(&caller)[key_ptr as usize..key_ptr as usize + 32],
                    );
                    let value = memory.data(&caller)
                        [value_ptr as usize..value_ptr as usize + value_len as usize]
                        .to_vec();

                    let data = caller.data_mut();
                    if data.context.is_static() {
                        return -1;
                    }

                    let storage_key = crate::storage::StorageKey::new(key);
                    data.context.storage.set(
                        &data.context.call.contract,
                        storage_key,
                        crate::storage::StorageValue::new(value),
                    );
                    0
                },
            )
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        // Get caller
        linker
            .func_wrap(
                "env",
                "synor_get_caller",
                |mut caller: wasmtime::Caller<'_, StoreData>, out_ptr: i32| {
                    let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
                    let data = caller.data();
                    let address = data.context.caller();
                    memory.data_mut(&mut caller)[out_ptr as usize..out_ptr as usize + 32]
                        .copy_from_slice(address.payload());
                },
            )
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        // Get value
        linker
            .func_wrap(
                "env",
                "synor_get_value",
                |caller: wasmtime::Caller<'_, StoreData>| -> i64 {
                    caller.data().context.value() as i64
                },
            )
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        // Get timestamp
        linker
            .func_wrap(
                "env",
                "synor_get_timestamp",
                |caller: wasmtime::Caller<'_, StoreData>| -> i64 {
                    caller.data().context.timestamp() as i64
                },
            )
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        // Get block height
        linker
            .func_wrap(
                "env",
                "synor_get_block_height",
                |caller: wasmtime::Caller<'_, StoreData>| -> i64 {
                    caller.data().context.block_height() as i64
                },
            )
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        // Revert
        linker
            .func_wrap(
                "env",
                "synor_revert",
                |mut caller: wasmtime::Caller<'_, StoreData>, msg_ptr: i32, msg_len: i32| {
                    let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
                    let msg_bytes = &memory.data(&caller)
                        [msg_ptr as usize..msg_ptr as usize + msg_len as usize];
                    let msg = String::from_utf8_lossy(msg_bytes).to_string();
                    caller.data_mut().error = Some(msg);
                },
            )
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        // Return
        linker
            .func_wrap(
                "env",
                "synor_return",
                |mut caller: wasmtime::Caller<'_, StoreData>, data_ptr: i32, data_len: i32| {
                    let memory = caller.get_export("memory").unwrap().into_memory().unwrap();
                    let data = memory.data(&caller)
                        [data_ptr as usize..data_ptr as usize + data_len as usize]
                        .to_vec();
                    caller.data_mut().return_data = data;
                },
            )
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        Ok(())
    }

    /// Compiles a contract module.
    pub fn compile(&self, bytecode: Vec<u8>) -> Result<ContractModule, VmError> {
        ContractModule::compile(&self.engine, bytecode)
    }

    /// Caches a compiled module.
    pub fn cache_module(&self, module: ContractModule) {
        self.modules.write().insert(module.id, Arc::new(module));
    }

    /// Gets a cached module.
    pub fn get_module(&self, id: &ContractId) -> Option<Arc<ContractModule>> {
        self.modules.read().get(id).cloned()
    }

    /// Instantiates a contract for execution.
    pub fn instantiate(
        &self,
        module: &ContractModule,
        context: ExecutionContext<MemoryStorage>,
        gas_limit: u64,
    ) -> Result<VmInstance, VmError> {
        // Create store with limits
        let limits = StoreLimitsBuilder::new()
            .memory_size(MAX_MEMORY_PAGES as usize * 65536)
            .build();

        let store_data = StoreData {
            gas: GasMeter::new(gas_limit),
            context,
            limits,
            return_data: Vec::new(),
            error: None,
        };

        let mut store = Store::new(&self.engine, store_data);
        store.limiter(|data| &mut data.limits);
        store
            .set_fuel(gas_limit)
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        // Instantiate
        let instance = self
            .linker
            .instantiate(&mut store, &module.module)
            .map_err(|e| VmError::ExecutionError(e.to_string()))?;

        // Get memory export
        let memory = instance
            .get_memory(&mut store, "memory")
            .ok_or_else(|| VmError::ExecutionError("No memory export".into()))?;

        Ok(VmInstance {
            store,
            instance,
            memory,
        })
    }

    /// Executes a contract call.
    pub fn execute(
        &self,
        module: &ContractModule,
        method: &str,
        args: &[u8],
        context: ExecutionContext<MemoryStorage>,
        gas_limit: u64,
    ) -> Result<ExecutionResult, VmError> {
        let mut instance = self.instantiate(module, context, gas_limit)?;
        instance.call(method, args)
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::context::CallContext;
    use crate::storage::MemoryStorage;
    use synor_types::{Address, Network};

    fn test_context() -> ExecutionContext<MemoryStorage> {
        let contract = ContractId::from_bytes([0x42; 32]);
        let caller = Address::from_ed25519_pubkey(Network::Mainnet, &[0x11; 32]);
        let call = CallContext::new(contract, caller, 1000, vec![]);
        ExecutionContext::<MemoryStorage>::testing(call)
    }

    #[test]
    fn test_engine_creation() {
        let engine = VmEngine::new();
        assert!(engine.is_ok());
    }

    #[test]
    fn test_compile_invalid() {
        let engine = VmEngine::new().unwrap();
        let result = engine.compile(vec![0x00, 0x01, 0x02]); // Invalid WASM
        assert!(result.is_err());
    }

    // Note: Testing actual WASM execution would require compiling Rust to WASM
    // which is beyond the scope of unit tests
}