//! WASM module runtime using wasmtime

use std::path::Path;
use std::sync::Arc;

use regex::Regex;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
use wasmtime::*;

use crate::config::{Module as ModuleConfig, Phase};
use super::types::{ModuleAction, RequestInfo, ResponseInfo};

/// Loaded WASM module
struct LoadedModule {
    /// Module name
    name: String,
    /// Compiled module
    _module: Module,
    /// Instance
    instance: Instance,
    /// Store
    store: Store<ModuleState>,
    /// Route patterns
    routes: Vec<Regex>,
    /// Execution phase
    phase: Phase,
}

/// State passed to WASM modules
struct ModuleState {
    /// Current request info (JSON encoded)
    request_json: String,
    /// Current response info (JSON encoded)
    response_json: String,
    /// Action result (JSON encoded)
    action_json: String,
    /// Pending messages to send (conn_id, payload)
    pending_sends: Vec<(String, String)>,
    /// Pending broadcasts (stream, payload)
    pending_broadcasts: Vec<(String, String)>,
    /// Pending closes (conn_id, code, reason)
    pending_closes: Vec<(String, i32, String)>,
    /// Resource limiter for memory bounds
    limiter: StoreLimits,
}

impl Default for ModuleState {
    fn default() -> Self {
        Self {
            request_json: String::new(),
            response_json: String::new(),
            action_json: String::new(),
            pending_sends: Vec::new(),
            pending_broadcasts: Vec::new(),
            pending_closes: Vec::new(),
            limiter: StoreLimitsBuilder::new()
                .memory_size(WASM_MAX_MEMORY_PAGES as usize * 65536)
                .build(),
        }
    }
}

/// Read a string from WASM memory using a Caller
fn read_wasm_string(caller: &mut Caller<'_, ModuleState>, ptr: i32, len: i32) -> Option<String> {
    if ptr == 0 || len <= 0 {
        return Some(String::new());
    }

    let memory = match caller.get_export("memory") {
        Some(Extern::Memory(mem)) => mem,
        _ => return None,
    };

    let data = memory.data(&*caller);
    let start = ptr as usize;
    let end = start + len as usize;

    if end > data.len() {
        return None;
    }

    std::str::from_utf8(&data[start..end])
        .ok()
        .map(|s| s.to_string())
}


/// WASM module runtime
pub struct ModuleRuntime {
    /// Wasmtime engine
    engine: Engine,
    /// Loaded modules
    modules: Arc<RwLock<Vec<LoadedModule>>>,
}

/// WASM sandbox limits
const WASM_MAX_MEMORY_PAGES: u64 = 256; // 256 * 64KB = 16MB max memory per module
const WASM_FUEL_LIMIT: u64 = 1_000_000; // Instruction budget per call

impl ModuleRuntime {
    /// Create a new module runtime with sandbox limits
    pub fn new() -> Result<Self> {
        let mut config = Config::new();
        config.async_support(true);
        config.consume_fuel(true); // Enable fuel metering for CPU limits

        let engine = Engine::new(&config)?;

        Ok(Self {
            engine,
            modules: Arc::new(RwLock::new(Vec::new())),
        })
    }

    /// Load modules from config
    pub async fn load_modules(&self, configs: &[ModuleConfig]) -> anyhow::Result<()> {
        let mut modules = self.modules.write().await;

        for config in configs {
            match self.load_module(config).await {
                Ok(module) => {
                    info!(name = %config.name, wasm = %config.wasm, "Loaded WASM module");
                    modules.push(module);
                }
                Err(e) => {
                    error!(name = %config.name, error = %e, "Failed to load WASM module");
                }
            }
        }

        Ok(())
    }

    /// Load a single module
    async fn load_module(&self, config: &ModuleConfig) -> anyhow::Result<LoadedModule> {
        let path = Path::new(&config.wasm);
        if !path.exists() {
            anyhow::bail!("Module file not found: {}", config.wasm);
        }

        // Compile routes
        let routes: Vec<Regex> = config
            .routes
            .iter()
            .filter_map(|r| match Regex::new(r) {
                Ok(regex) => Some(regex),
                Err(e) => {
                    warn!(pattern = %r, error = %e, "Invalid route pattern");
                    None
                }
            })
            .collect();

        // Load and compile module
        let module = Module::from_file(&self.engine, path)?;

        // Create store with state
        let mut store = Store::new(&self.engine, ModuleState::default());

        // Set memory limits via ResourceLimiter
        store.limiter(|state| &mut state.limiter);

        // Create linker with host functions
        let mut linker = Linker::new(&self.engine);

        // === HTTP Module Interface (env namespace) ===

        // Host function: get request info
        linker.func_wrap("env", "get_request", |caller: Caller<'_, ModuleState>| -> i32 {
            let state = caller.data();
            state.request_json.len() as i32
        })?;

        // Host function: get response info
        linker.func_wrap("env", "get_response", |caller: Caller<'_, ModuleState>| -> i32 {
            let state = caller.data();
            state.response_json.len() as i32
        })?;

        // Host function: set action
        linker.func_wrap("env", "set_action", |mut caller: Caller<'_, ModuleState>, action: i32| {
            let state = caller.data_mut();
            state.action_json = match action {
                0 => r#"{"Continue":null}"#.to_string(),
                1 => r#"{"Block":{"status":403,"body":"Blocked by module"}}"#.to_string(),
                _ => r#"{"Continue":null}"#.to_string(),
            };
        })?;

        // Host function: log (env namespace)
        linker.func_wrap(
            "env",
            "log",
            |_caller: Caller<'_, ModuleState>, level: i32, ptr: i32, len: i32| {
            // In a real implementation, we'd read from WASM memory
            debug!(level = level, ptr = ptr, len = len, "Module log");
        })?;

        // === WebSocket/ActionCable Interface (warpdrive namespace) ===

        // host_log: Log a message from the plugin
        linker.func_wrap(
            "warpdrive",
            "host_log",
            |mut caller: Caller<'_, ModuleState>, level: i32, msg_ptr: i32, msg_len: i32| -> i32 {
                let message = read_wasm_string(&mut caller, msg_ptr, msg_len)
                    .unwrap_or_else(|| "<invalid utf-8>".to_string());

                match level {
                    0 => debug!(target: "warpdrive::plugin", "{}", message),
                    1 => warn!(target: "warpdrive::plugin", "{}", message),
                    2 => info!(target: "warpdrive::plugin", "{}", message),
                    _ => debug!(target: "warpdrive::plugin", "{}", message),
                }

                0 // Success
            },
        )?;

        // host_send: Send a message to a specific connection
        linker.func_wrap(
            "warpdrive",
            "host_send",
            |mut caller: Caller<'_, ModuleState>,
             conn_id_ptr: i32,
             conn_id_len: i32,
             payload_ptr: i32,
             payload_len: i32|
             -> i32 {
                // Read strings first
                let conn_id = match read_wasm_string(&mut caller, conn_id_ptr, conn_id_len) {
                    Some(s) => s,
                    None => return -2,
                };

                let payload = match read_wasm_string(&mut caller, payload_ptr, payload_len) {
                    Some(s) => s,
                    None => return -3,
                };

                debug!(target: "warpdrive::plugin", conn_id = %conn_id, payload_len = payload.len(), "host_send");

                // Queue the message for sending
                caller.data_mut().pending_sends.push((conn_id, payload));

                0 // Success
            },
        )?;

        // host_broadcast: Broadcast a message to all subscribers of a stream
        linker.func_wrap(
            "warpdrive",
            "host_broadcast",
            |mut caller: Caller<'_, ModuleState>,
             stream_ptr: i32,
             stream_len: i32,
             payload_ptr: i32,
             payload_len: i32|
             -> i32 {
                let stream = match read_wasm_string(&mut caller, stream_ptr, stream_len) {
                    Some(s) => s,
                    None => return -2,
                };

                let payload = match read_wasm_string(&mut caller, payload_ptr, payload_len) {
                    Some(s) => s,
                    None => return -3,
                };

                debug!(target: "warpdrive::plugin", stream = %stream, payload_len = payload.len(), "host_broadcast");

                // Queue the broadcast
                caller.data_mut().pending_broadcasts.push((stream, payload));

                0 // Success
            },
        )?;

        // host_close: Close a WebSocket connection
        linker.func_wrap(
            "warpdrive",
            "host_close",
            |mut caller: Caller<'_, ModuleState>,
             conn_id_ptr: i32,
             conn_id_len: i32,
             code: i32,
             reason_ptr: i32,
             reason_len: i32|
             -> i32 {
                let conn_id = match read_wasm_string(&mut caller, conn_id_ptr, conn_id_len) {
                    Some(s) => s,
                    None => return -2,
                };

                let reason = read_wasm_string(&mut caller, reason_ptr, reason_len)
                    .unwrap_or_default();

                debug!(target: "warpdrive::plugin", conn_id = %conn_id, code = code, reason = %reason, "host_close");

                // Queue the close
                caller.data_mut().pending_closes.push((conn_id, code, reason));

                0 // Success
            },
        )?;

        // Instantiate
        let instance = linker.instantiate(&mut store, &module)?;

        Ok(LoadedModule {
            name: config.name.clone(),
            _module: module,
            instance,
            store,
            routes,
            phase: config.phase,
        })
    }

    /// Process request through modules
    pub async fn process_request(&self, path: &str, request: &RequestInfo) -> ModuleAction {
        let mut modules = self.modules.write().await;

        for module in modules.iter_mut() {
            // Check phase
            if module.phase != Phase::Request {
                continue;
            }

            // Check route match
            if !module.routes.is_empty() && !module.routes.iter().any(|r| r.is_match(path)) {
                continue;
            }

            // Set request info in state
            module.store.data_mut().request_json =
                serde_json::to_string(request).unwrap_or_default();
            module.store.data_mut().action_json = r#"{"Continue":null}"#.to_string();

            // Call on_request if it exists
            if let Some(func) = module.instance.get_func(&mut module.store, "on_request") {
                // Add fuel budget for this call
                let _ = module.store.set_fuel(WASM_FUEL_LIMIT);

                match func.call(&mut module.store, &[], &mut []) {
                    Ok(_) => {
                        // Parse action
                        let json = &module.store.data().action_json;
                        let action: ModuleAction =
                            serde_json::from_str(json).unwrap_or_default();

                        match &action {
                            ModuleAction::Continue => continue,
                            ModuleAction::Modify { .. } => return action,
                            ModuleAction::Block { .. } => return action,
                        }
                    }
                    Err(e) => {
                        warn!(module = %module.name, error = %e, "Module on_request failed");
                    }
                }
            }
        }

        ModuleAction::Continue
    }

    /// Process response through modules
    pub async fn process_response(&self, path: &str, response: &ResponseInfo) -> ModuleAction {
        let mut modules = self.modules.write().await;

        for module in modules.iter_mut() {
            // Check phase
            if module.phase != Phase::Response {
                continue;
            }

            // Check route match
            if !module.routes.is_empty() && !module.routes.iter().any(|r| r.is_match(path)) {
                continue;
            }

            // Set response info in state
            module.store.data_mut().response_json =
                serde_json::to_string(response).unwrap_or_default();
            module.store.data_mut().action_json = r#"{"Continue":null}"#.to_string();

            // Call on_response if it exists
            let func = module.instance.get_func(&mut module.store, "on_response");
            if let Some(func) = func {
                // Add fuel budget for this call
                let _ = module.store.set_fuel(WASM_FUEL_LIMIT);

                match func.call(&mut module.store, &[], &mut []) {
                    Ok(_) => {
                        let json = &module.store.data().action_json;
                        let action: ModuleAction =
                            serde_json::from_str(json).unwrap_or_default();

                        match &action {
                            ModuleAction::Continue => continue,
                            ModuleAction::Modify { .. } => return action,
                            ModuleAction::Block { .. } => return action,
                        }
                    }
                    Err(e) => {
                        warn!(module = %module.name, error = %e, "Module on_response failed");
                    }
                }
            }
        }

        ModuleAction::Continue
    }

    /// Get loaded module names
    pub async fn module_names(&self) -> Vec<String> {
        self.modules.read().await.iter().map(|m| m.name.clone()).collect()
    }
}

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

/// WebSocket event types
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(i32)]
pub enum WsEventType {
    Connect = 0,
    Message = 1,
    Disconnect = 2,
}

/// Result of processing a WebSocket event
#[derive(Debug, Default)]
pub struct WsEventResult {
    /// Messages to send to specific connections
    pub sends: Vec<(String, String)>,
    /// Broadcasts to streams
    pub broadcasts: Vec<(String, String)>,
    /// Connections to close (conn_id, code, reason)
    pub closes: Vec<(String, i32, String)>,
}

/// WebSocket plugin runtime for ActionCable-style plugins
pub struct WsPluginRuntime {
    /// Wasmtime engine
    engine: Engine,
    /// Loaded plugin module
    module: Option<Module>,
    /// Plugin instance
    instance: Option<Instance>,
    /// Store with state
    store: Option<Store<ModuleState>>,
    /// Plugin name
    name: String,
}

impl WsPluginRuntime {
    /// Create a new WebSocket plugin runtime with sandbox limits
    pub fn new() -> Result<Self> {
        let mut config = Config::new();
        config.consume_fuel(true); // Enable fuel metering

        let engine = Engine::new(&config)?;

        Ok(Self {
            engine,
            module: None,
            instance: None,
            store: None,
            name: String::new(),
        })
    }

    /// Load a WebSocket plugin from a WASM file
    pub fn load(&mut self, path: &Path, name: &str) -> anyhow::Result<()> {
        if !path.exists() {
            anyhow::bail!("Plugin file not found: {}", path.display());
        }

        // Load and compile module
        let module = Module::from_file(&self.engine, path)?;

        // Create store with state and memory limits
        let mut store = Store::new(&self.engine, ModuleState::default());
        store.limiter(|state| &mut state.limiter);

        // Create linker with warpdrive host functions
        let mut linker = Linker::new(&self.engine);

        // host_log
        linker.func_wrap(
            "warpdrive",
            "host_log",
            |mut caller: Caller<'_, ModuleState>, level: i32, msg_ptr: i32, msg_len: i32| -> i32 {
                let message = read_wasm_string(&mut caller, msg_ptr, msg_len)
                    .unwrap_or_else(|| "<invalid utf-8>".to_string());

                match level {
                    0 => debug!(target: "warpdrive::plugin", "{}", message),
                    1 => warn!(target: "warpdrive::plugin", "{}", message),
                    2 => info!(target: "warpdrive::plugin", "{}", message),
                    _ => debug!(target: "warpdrive::plugin", "{}", message),
                }

                0
            },
        )?;

        // host_send
        linker.func_wrap(
            "warpdrive",
            "host_send",
            |mut caller: Caller<'_, ModuleState>,
             conn_id_ptr: i32,
             conn_id_len: i32,
             payload_ptr: i32,
             payload_len: i32|
             -> i32 {
                let conn_id = match read_wasm_string(&mut caller, conn_id_ptr, conn_id_len) {
                    Some(s) => s,
                    None => return -2,
                };

                let payload = match read_wasm_string(&mut caller, payload_ptr, payload_len) {
                    Some(s) => s,
                    None => return -3,
                };

                caller.data_mut().pending_sends.push((conn_id, payload));
                0
            },
        )?;

        // host_broadcast
        linker.func_wrap(
            "warpdrive",
            "host_broadcast",
            |mut caller: Caller<'_, ModuleState>,
             stream_ptr: i32,
             stream_len: i32,
             payload_ptr: i32,
             payload_len: i32|
             -> i32 {
                let stream = match read_wasm_string(&mut caller, stream_ptr, stream_len) {
                    Some(s) => s,
                    None => return -2,
                };

                let payload = match read_wasm_string(&mut caller, payload_ptr, payload_len) {
                    Some(s) => s,
                    None => return -3,
                };

                caller.data_mut().pending_broadcasts.push((stream, payload));
                0
            },
        )?;

        // host_close
        linker.func_wrap(
            "warpdrive",
            "host_close",
            |mut caller: Caller<'_, ModuleState>,
             conn_id_ptr: i32,
             conn_id_len: i32,
             code: i32,
             reason_ptr: i32,
             reason_len: i32|
             -> i32 {
                let conn_id = match read_wasm_string(&mut caller, conn_id_ptr, conn_id_len) {
                    Some(s) => s,
                    None => return -2,
                };

                let reason = read_wasm_string(&mut caller, reason_ptr, reason_len)
                    .unwrap_or_default();

                caller.data_mut().pending_closes.push((conn_id, code, reason));
                0
            },
        )?;

        // Instantiate
        let instance = linker.instantiate(&mut store, &module)?;

        self.module = Some(module);
        self.instance = Some(instance);
        self.store = Some(store);
        self.name = name.to_string();

        info!(name = %name, path = %path.display(), "Loaded WebSocket plugin");

        Ok(())
    }

    /// Process a WebSocket event through the plugin
    pub fn on_event(
        &mut self,
        event_type: WsEventType,
        conn_id: &str,
        payload: &str,
    ) -> anyhow::Result<WsEventResult> {
        let store = self.store.as_mut().ok_or_else(|| anyhow::anyhow!("Plugin not loaded"))?;
        let instance = self.instance.as_ref().ok_or_else(|| anyhow::anyhow!("Plugin not loaded"))?;

        // Clear pending actions
        store.data_mut().pending_sends.clear();
        store.data_mut().pending_broadcasts.clear();
        store.data_mut().pending_closes.clear();

        // Get memory and wasm_on_event function
        let memory = instance
            .get_memory(&mut *store, "memory")
            .ok_or_else(|| anyhow::anyhow!("No memory export"))?;

        let on_event = instance
            .get_typed_func::<(i32, i32, i32, i32, i32), i32>(&mut *store, "wasm_on_event")
            .map_err(|e| anyhow::anyhow!("wasm_on_event not found: {}", e))?;

        // Write conn_id and payload to WASM memory
        // We'll use a simple allocation strategy: write at fixed offsets
        let conn_id_bytes = conn_id.as_bytes();
        let payload_bytes = payload.as_bytes();

        // Write to memory starting at offset 1024 (leave some space for stack)
        let conn_id_ptr = 1024i32;
        let payload_ptr = conn_id_ptr + conn_id_bytes.len() as i32 + 16; // 16 byte alignment

        // Ensure memory is large enough
        let required_size = (payload_ptr as usize) + payload_bytes.len() + 1024;
        let current_size = memory.data_size(&*store);
        if required_size > current_size {
            let pages_needed = ((required_size - current_size) / 65536) + 1;
            memory.grow(&mut *store, pages_needed as u64)?;
        }

        // Write strings to memory
        memory.write(&mut *store, conn_id_ptr as usize, conn_id_bytes)?;
        memory.write(&mut *store, payload_ptr as usize, payload_bytes)?;

        // Add fuel budget for this call
        let _ = store.set_fuel(WASM_FUEL_LIMIT);

        // Call wasm_on_event
        let result = on_event.call(
            &mut *store,
            (
                event_type as i32,
                conn_id_ptr,
                conn_id_bytes.len() as i32,
                payload_ptr,
                payload_bytes.len() as i32,
            ),
        )?;

        if result != 0 {
            warn!(
                plugin = %self.name,
                event = ?event_type,
                conn_id = %conn_id,
                result = result,
                "Plugin returned error"
            );
        }

        // Collect pending actions
        let state = store.data();
        Ok(WsEventResult {
            sends: state.pending_sends.clone(),
            broadcasts: state.pending_broadcasts.clone(),
            closes: state.pending_closes.clone(),
        })
    }

    /// Get the plugin name
    pub fn name(&self) -> &str {
        &self.name
    }
}

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