use std::sync::Arc;
use tracing::{debug, info, warn};
use wasmtime::{Config, Engine, Instance, Module, Store};

use crate::sfu::ForwardPacket;

/// Holds the environment and engine for executing WASM edge functions
pub struct EdgeWasmEngine {
    engine: Engine,
}

impl EdgeWasmEngine {
    pub fn new() -> Result<Self, anyhow::Error> {
        let mut config = Config::new();
        // Allow WASM to consume fuel if we want strict execution limits
        config.consume_fuel(true);
        // Optimize for fast instantiation since these are short-lived packet mutations
        config.cranelift_opt_level(wasmtime::OptLevel::Speed);

        let engine = Engine::new(&config)?;
        info!("[Edge WASM] Engine initialized for edge computations");
        Ok(Self { engine })
    }

    /// Loads a `.wasm` file representing a media transform or custom route logic.
    pub fn load_module(&self, wasm_bytes: &[u8]) -> Result<Module, anyhow::Error> {
        let module = Module::from_binary(&self.engine, wasm_bytes)?;
        Ok(module)
    }

    /// Example: executes a loaded WASM module to transform a packet in-flight.
    /// This is where N-05 (media transforms) logic exists.
    pub fn process_packet(
        &self,
        module: &Module,
        packet: &mut ForwardPacket,
    ) -> Result<(), anyhow::Error> {
        // We give each execution its own store.
        let mut store = Store::new(&self.engine, ());
        store.set_fuel(100_000)?; // Restrict the amount of compute allowed per packet

        // Instantiate the module into the store.
        let instance = Instance::new(&mut store, module, &[])?;

        // Extract memory to pass data into the WASM guest
        let memory = instance
            .get_memory(&mut store, "memory")
            .ok_or_else(|| anyhow::anyhow!("WASM module must export 'memory'"))?;

        // Determine where we can write the payload in the WASM linear memory
        // A real convention would have `alloc(size)` exported by the module
        let alloc = instance.get_typed_func::<u32, u32>(&mut store, "alloc")?;

        let payload_len: u32 = packet.payload.len() as u32;
        let ptr = alloc.call(&mut store, payload_len)?;

        // Write packet payload into caller memory
        memory.write(&mut store, ptr as usize, &packet.payload)?;

        // Execute the transform function (N-05: e.g. 'transform_media', 'censor_audio')
        let transform = instance.get_typed_func::<(u32, u32), u32>(&mut store, "transform")?;

        // Function returns the new payload length
        let new_len = transform.call(&mut store, (ptr, payload_len))?;

        // Read the mutated data back
        let mut new_payload = vec![0u8; new_len as usize];
        memory.read(&store, ptr as usize, &mut new_payload)?;

        // Apply mutation
        packet.payload = new_payload;

        debug!(
            "[Edge WASM] Packet {} transformed by edge function",
            packet.sequence
        );

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sfu::{ForwardPacket, ParticipantId, StreamId};

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

    #[test]
    fn test_wasm_packet_transform() {
        let engine = EdgeWasmEngine::new().expect("Failed to create WASM engine");

        // A simple WAT module that:
        // 1. Exports memory
        // 2. Exports `alloc(size) -> ptr` (mocked to just return 0)
        // 3. Exports `transform(ptr, len) -> new_len`
        // The transform function simply adds 1 to the first byte in memory.
        let wat = r#"
        (module
            (memory (export "memory") 1)
            (func (export "alloc") (param $size i32) (result i32)
                (i32.const 0)
            )
            (func (export "transform") (param $ptr i32) (param $len i32) (result i32)
                (i32.store8
                    (local.get $ptr)
                    (i32.add
                        (i32.load8_u (local.get $ptr))
                        (i32.const 1)
                    )
                )
                (local.get $len)
            )
        )
        "#;
        let wasm_bytes = wat::parse_str(wat).expect("Failed to compile WAT");

        let module = engine
            .load_module(&wasm_bytes)
            .expect("Failed to load Wasm bytes");

        let mut packet = ForwardPacket {
            source: "p1".to_string(),
            stream_id: "s1".to_string(),
            sequence: 42,
            payload: vec![10, 20, 30], // initial first byte is 10
            lossless: false,
        };

        engine
            .process_packet(&module, &mut packet)
            .expect("Failed to execute WASM transform");

        // The WASM function should have mutated the first byte from 10 to 11
        assert_eq!(packet.payload, vec![11, 20, 30]);
        assert_eq!(packet.sequence, 42); // unchanged
    }

    #[test]
    fn test_wasm_fuel_exhaustion() {
        let engine = EdgeWasmEngine::new().expect("Engine created");
        // An infinite loop in WASM to test fuel consumption limit
        let wat = r#"
        (module
            (memory (export "memory") 1)
            (func (export "alloc") (param $size i32) (result i32)
                (i32.const 0)
            )
            (func (export "transform") (param $ptr i32) (param $len i32) (result i32)
                (loop $my_loop
                    br $my_loop
                )
                (local.get $len)
            )
        )
        "#;
        let wasm_bytes = wat::parse_str(wat).expect("Failed to compile WAT");
        let module = engine
            .load_module(&wasm_bytes)
            .expect("Failed to load Wasm bytes");

        let mut packet = ForwardPacket {
            source: "p0".to_string(),
            stream_id: "s0".to_string(),
            sequence: 1,
            payload: vec![0],
            lossless: false,
        };

        let result = engine.process_packet(&module, &mut packet);

        // Ensure execution was aborted by Wasmtime because we ran out of fuel
        assert!(result.is_err());
        let err_str = result.unwrap_err().to_string();
        assert!(err_str.contains("error while executing"));
    }
}
