//! # RustyTorch++ WASM Inference Runtime
//!
//! WebAssembly inference runtime for deploying ML models in browsers and Node.js.
//!
//! ## Features
//!
//! - **Browser deployment**: Run inference directly in web browsers
//! - **Node.js support**: Server-side WASM with Wasmtime/Wasmer
//! - **Zero dependencies**: No Python, CUDA, or native libraries required
//! - **Streaming inference**: Token-by-token generation for LLMs
//! - **Quantization**: INT8/INT4 models for reduced memory footprint
//!
//! ## Usage (JavaScript/TypeScript)
//!
//! ```javascript
//! import init, { WasmInferenceEngine, InferenceConfig } from 'rtx-wasm-inference';
//!
//! async function main() {
//!     await init();
//!
//!     const config = InferenceConfig.default();
//!     const engine = await WasmInferenceEngine.new(config);
//!
//!     const result = await engine.infer("Hello, world!");
//!     console.log(result.text);
//! }
//! ```

#![warn(missing_docs)]
#![allow(clippy::new_without_default)]

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use wasm_bindgen::prelude::*;

pub mod model;
pub mod quantization;
pub mod runtime;
pub mod tensor;
pub mod tensor_core;
pub mod tokenizer;
pub mod webgpu;

// Re-export submodules
pub use model::*;
pub use quantization::*;
pub use runtime::*;
pub use tensor::*;
pub use tokenizer::*;
pub use webgpu::*;

// Re-export core tensor types (no_std compatible)
pub use tensor_core::backend::{
    BackendLimits, BackendPreference, BackendType, UnifiedBackend, available_backends,
};
pub use tensor_core::simd::{SIMD_AVAILABLE, SIMD_WIDTH};
pub use tensor_core::wasm_api::{
    WasmBackend, WasmTensor, get_backend_info, get_simd_width, is_simd_available,
};
pub use tensor_core::{DType, TensorCore, TensorError};

/// Initialize panic hook for better error messages in browser console
#[wasm_bindgen(start)]
pub fn init_panic_hook() {
    #[cfg(feature = "console_error_panic_hook")]
    console_error_panic_hook::set_once();
}

/// Compute backend for WASM inference
#[wasm_bindgen]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ComputeBackend {
    /// CPU-only execution (always available)
    Cpu,
    /// WebGPU acceleration (requires browser support)
    WebGpu,
    /// Automatic selection (prefer GPU if available)
    Auto,
}

impl Default for ComputeBackend {
    fn default() -> Self {
        Self::Auto
    }
}

/// Configuration for WASM inference engine
#[wasm_bindgen]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceConfig {
    /// Maximum sequence length
    max_seq_len: usize,
    /// Maximum batch size
    max_batch_size: usize,
    /// Use quantized model
    use_quantization: bool,
    /// Quantization bits (4 or 8)
    quantization_bits: u8,
    /// Enable KV-cache
    use_kv_cache: bool,
    /// Temperature for sampling
    temperature: f32,
    /// Top-p sampling threshold
    top_p: f32,
    /// Top-k sampling limit
    top_k: usize,
    /// Enable streaming output
    streaming: bool,
    /// Compute backend selection
    compute_backend: ComputeBackend,
    /// WebGPU workgroup size (for tuning)
    webgpu_workgroup_size: u32,
}

#[wasm_bindgen]
impl InferenceConfig {
    /// Create default configuration
    #[wasm_bindgen(constructor)]
    pub fn new() -> Self {
        Self::default()
    }

    /// Create configuration optimized for speed
    #[wasm_bindgen]
    pub fn fast() -> Self {
        Self {
            max_seq_len: 512,
            max_batch_size: 1,
            use_quantization: true,
            quantization_bits: 4,
            use_kv_cache: true,
            temperature: 0.7,
            top_p: 0.9,
            top_k: 40,
            streaming: true,
            compute_backend: ComputeBackend::Auto,
            webgpu_workgroup_size: 256,
        }
    }

    /// Create configuration optimized for quality
    #[wasm_bindgen]
    pub fn quality() -> Self {
        Self {
            max_seq_len: 2048,
            max_batch_size: 1,
            use_quantization: false,
            quantization_bits: 8,
            use_kv_cache: true,
            temperature: 0.8,
            top_p: 0.95,
            top_k: 50,
            streaming: false,
            compute_backend: ComputeBackend::Auto,
            webgpu_workgroup_size: 256,
        }
    }

    /// Create configuration optimized for WebGPU
    #[wasm_bindgen]
    pub fn webgpu_optimized() -> Self {
        Self {
            max_seq_len: 2048,
            max_batch_size: 1,
            use_quantization: true,
            quantization_bits: 8,
            use_kv_cache: true,
            temperature: 0.7,
            top_p: 0.9,
            top_k: 50,
            streaming: true,
            compute_backend: ComputeBackend::WebGpu,
            webgpu_workgroup_size: 256,
        }
    }

    /// Create configuration for CPU-only execution
    #[wasm_bindgen]
    pub fn cpu_only() -> Self {
        Self {
            compute_backend: ComputeBackend::Cpu,
            ..Self::default()
        }
    }

    /// Set maximum sequence length
    #[wasm_bindgen]
    pub fn with_max_seq_len(mut self, len: usize) -> Self {
        self.max_seq_len = len;
        self
    }

    /// Set temperature
    #[wasm_bindgen]
    pub fn with_temperature(mut self, temp: f32) -> Self {
        self.temperature = temp;
        self
    }

    /// Enable streaming
    #[wasm_bindgen]
    pub fn with_streaming(mut self, enabled: bool) -> Self {
        self.streaming = enabled;
        self
    }

    /// Enable quantization
    #[wasm_bindgen]
    pub fn with_quantization(mut self, bits: u8) -> Self {
        self.use_quantization = true;
        self.quantization_bits = bits;
        self
    }

    /// Set compute backend
    #[wasm_bindgen]
    pub fn with_compute_backend(mut self, backend: ComputeBackend) -> Self {
        self.compute_backend = backend;
        self
    }

    /// Set WebGPU workgroup size
    #[wasm_bindgen]
    pub fn with_webgpu_workgroup_size(mut self, size: u32) -> Self {
        self.webgpu_workgroup_size = size;
        self
    }

    /// Get compute backend
    #[wasm_bindgen(getter)]
    pub fn compute_backend(&self) -> ComputeBackend {
        self.compute_backend
    }

    /// Get WebGPU workgroup size
    #[wasm_bindgen(getter)]
    pub fn webgpu_workgroup_size(&self) -> u32 {
        self.webgpu_workgroup_size
    }
}

impl Default for InferenceConfig {
    fn default() -> Self {
        Self {
            max_seq_len: 1024,
            max_batch_size: 1,
            use_quantization: true,
            quantization_bits: 8,
            use_kv_cache: true,
            temperature: 0.7,
            top_p: 0.9,
            top_k: 50,
            streaming: false,
            compute_backend: ComputeBackend::Auto,
            webgpu_workgroup_size: 256,
        }
    }
}

/// Result from inference
#[wasm_bindgen]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceResult {
    /// Generated text
    text: String,
    /// Number of tokens generated
    tokens_generated: usize,
    /// Inference time in milliseconds
    inference_time_ms: f64,
    /// Tokens per second
    tokens_per_sec: f64,
    /// Finish reason
    finish_reason: String,
}

#[wasm_bindgen]
impl InferenceResult {
    /// Get the generated text
    #[wasm_bindgen(getter)]
    pub fn text(&self) -> String {
        self.text.clone()
    }

    /// Get the number of tokens generated
    #[wasm_bindgen(getter)]
    pub fn tokens_generated(&self) -> usize {
        self.tokens_generated
    }

    /// Get the inference time in milliseconds
    #[wasm_bindgen(getter)]
    pub fn inference_time_ms(&self) -> f64 {
        self.inference_time_ms
    }

    /// Get tokens per second
    #[wasm_bindgen(getter)]
    pub fn tokens_per_sec(&self) -> f64 {
        self.tokens_per_sec
    }

    /// Get the finish reason
    #[wasm_bindgen(getter)]
    pub fn finish_reason(&self) -> String {
        self.finish_reason.clone()
    }

    /// Convert to JSON string
    #[wasm_bindgen]
    pub fn to_json(&self) -> Result<String, JsError> {
        serde_json::to_string(self).map_err(|e| JsError::new(&e.to_string()))
    }
}

/// Main WASM inference engine
#[wasm_bindgen]
pub struct WasmInferenceEngine {
    config: InferenceConfig,
    runtime_info: WasmRuntimeInfo,
    model: Option<WasmModel>,
    tokenizer: Option<WasmTokenizer>,
    kv_cache: Option<WasmKvCache>,
    stats: EngineStats,
}

/// Engine statistics
#[derive(Debug, Clone, Default)]
struct EngineStats {
    total_requests: u64,
    total_tokens: u64,
    total_time_ms: f64,
}

#[wasm_bindgen]
impl WasmInferenceEngine {
    /// Create a new inference engine
    #[wasm_bindgen(constructor)]
    pub fn new(config: InferenceConfig) -> Result<WasmInferenceEngine, JsError> {
        let runtime_info = WasmRuntimeInfo::detect();

        Ok(Self {
            config,
            runtime_info,
            model: None,
            tokenizer: None,
            kv_cache: None,
            stats: EngineStats::default(),
        })
    }

    /// Load a model from bytes
    #[wasm_bindgen]
    pub async fn load_model(&mut self, model_bytes: &[u8]) -> Result<(), JsError> {
        let model = WasmModel::from_bytes(model_bytes, self.config.use_quantization)?;
        self.model = Some(model);

        if self.config.use_kv_cache {
            self.kv_cache = Some(WasmKvCache::new(
                self.config.max_seq_len,
                32,  // num_layers
                32,  // num_heads
                128, // head_dim
            ));
        }

        Ok(())
    }

    /// Load tokenizer from JSON
    #[wasm_bindgen]
    pub fn load_tokenizer(&mut self, tokenizer_json: &str) -> Result<(), JsError> {
        let tokenizer = WasmTokenizer::from_json(tokenizer_json)?;
        self.tokenizer = Some(tokenizer);
        Ok(())
    }

    /// Run inference on input text
    #[wasm_bindgen]
    pub async fn infer(
        &mut self,
        input: &str,
        max_tokens: Option<usize>,
    ) -> Result<InferenceResult, JsError> {
        let start = self.runtime_info.performance_now();
        let max_tokens = max_tokens.unwrap_or(256);

        // Tokenize input
        let tokenizer = self
            .tokenizer
            .as_ref()
            .ok_or_else(|| JsError::new("Tokenizer not loaded"))?;
        let input_ids = tokenizer.encode(input)?;

        // Run model inference
        let model = self
            .model
            .as_ref()
            .ok_or_else(|| JsError::new("Model not loaded"))?;

        let mut output_ids = input_ids.clone();
        let mut generated_tokens = 0;

        for _ in 0..max_tokens {
            // Get next token prediction
            let logits = model.forward(&output_ids, self.kv_cache.as_mut())?;

            // Sample next token
            let next_token = self.sample_token(&logits)?;

            // Check for EOS
            if next_token == tokenizer.eos_token_id() {
                break;
            }

            output_ids.push(next_token);
            generated_tokens += 1;

            // Update KV cache position
            if let Some(cache) = &mut self.kv_cache {
                cache.step();
            }
        }

        // Decode output
        let output_text = tokenizer.decode(&output_ids[input_ids.len()..])?;

        let end = self.runtime_info.performance_now();
        let inference_time_ms = end - start;
        let tokens_per_sec = if inference_time_ms > 0.0 {
            (generated_tokens as f64 / inference_time_ms) * 1000.0
        } else {
            0.0
        };

        // Update stats
        self.stats.total_requests += 1;
        self.stats.total_tokens += generated_tokens as u64;
        self.stats.total_time_ms += inference_time_ms;

        Ok(InferenceResult {
            text: output_text,
            tokens_generated: generated_tokens,
            inference_time_ms,
            tokens_per_sec,
            finish_reason: if generated_tokens < max_tokens {
                "stop".to_string()
            } else {
                "length".to_string()
            },
        })
    }

    /// Sample a token from logits
    fn sample_token(&self, logits: &[f32]) -> Result<u32, JsError> {
        if logits.is_empty() {
            return Err(JsError::new("Empty logits"));
        }

        // Apply temperature
        let temp = self.config.temperature;
        let scaled: Vec<f32> = logits.iter().map(|&x| x / temp).collect();

        // Softmax
        let max_val = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max);
        let exp_sum: f32 = scaled.iter().map(|&x| (x - max_val).exp()).sum();
        let probs: Vec<f32> = scaled
            .iter()
            .map(|&x| (x - max_val).exp() / exp_sum)
            .collect();

        // Top-p sampling
        let mut sorted_indices: Vec<usize> = (0..probs.len()).collect();
        sorted_indices.sort_by(|&a, &b| probs[b].partial_cmp(&probs[a]).unwrap());

        let mut cumsum = 0.0;
        let mut cutoff_idx = sorted_indices.len();
        for (i, &idx) in sorted_indices.iter().enumerate() {
            cumsum += probs[idx];
            if cumsum >= self.config.top_p {
                cutoff_idx = i + 1;
                break;
            }
        }

        // Apply top-k
        cutoff_idx = cutoff_idx.min(self.config.top_k);

        // Renormalize
        let valid_indices = &sorted_indices[..cutoff_idx];
        let valid_sum: f32 = valid_indices.iter().map(|&i| probs[i]).sum();
        let normalized: Vec<f32> = valid_indices
            .iter()
            .map(|&i| probs[i] / valid_sum)
            .collect();

        // Sample
        let r: f32 = rand::random();
        let mut cumsum = 0.0;
        for (i, &prob) in normalized.iter().enumerate() {
            cumsum += prob;
            if r <= cumsum {
                return Ok(valid_indices[i] as u32);
            }
        }

        Ok(valid_indices[0] as u32)
    }

    /// Get runtime information
    #[wasm_bindgen]
    pub fn runtime_info(&self) -> WasmRuntimeInfo {
        self.runtime_info.clone()
    }

    /// Get engine statistics
    #[wasm_bindgen]
    pub fn stats(&self) -> Result<JsValue, JsError> {
        let stats = HashMap::from([
            (
                "total_requests".to_string(),
                self.stats.total_requests as f64,
            ),
            ("total_tokens".to_string(), self.stats.total_tokens as f64),
            ("total_time_ms".to_string(), self.stats.total_time_ms),
            (
                "avg_tokens_per_sec".to_string(),
                if self.stats.total_time_ms > 0.0 {
                    (self.stats.total_tokens as f64 / self.stats.total_time_ms) * 1000.0
                } else {
                    0.0
                },
            ),
        ]);
        serde_wasm_bindgen::to_value(&stats).map_err(|e| JsError::new(&e.to_string()))
    }

    /// Clear KV cache
    #[wasm_bindgen]
    pub fn clear_cache(&mut self) {
        if let Some(cache) = &mut self.kv_cache {
            cache.clear();
        }
    }

    /// Check if model is loaded
    #[wasm_bindgen]
    pub fn is_model_loaded(&self) -> bool {
        self.model.is_some()
    }

    /// Check if tokenizer is loaded
    #[wasm_bindgen]
    pub fn is_tokenizer_loaded(&self) -> bool {
        self.tokenizer.is_some()
    }

    /// Get memory usage estimate
    #[wasm_bindgen]
    pub fn memory_usage(&self) -> usize {
        let model_mem = self
            .model
            .as_ref()
            .map_or(0, model::WasmModel::memory_usage);
        let cache_mem = self
            .kv_cache
            .as_ref()
            .map_or(0, tensor::WasmKvCache::memory_usage);
        let tokenizer_mem = self
            .tokenizer
            .as_ref()
            .map_or(0, tokenizer::WasmTokenizer::memory_usage);
        model_mem + cache_mem + tokenizer_mem
    }
}

/// Version information
#[wasm_bindgen]
pub fn version() -> String {
    env!("CARGO_PKG_VERSION").to_string()
}

/// Check if SIMD is supported
#[wasm_bindgen]
pub fn simd_supported() -> bool {
    #[cfg(feature = "simd")]
    {
        true
    }
    #[cfg(not(feature = "simd"))]
    {
        false
    }
}

/// Check if threading is supported
#[wasm_bindgen]
pub fn threads_supported() -> bool {
    #[cfg(feature = "threads")]
    {
        true
    }
    #[cfg(not(feature = "threads"))]
    {
        false
    }
}

/// Check if WebGPU is supported in the current browser
#[wasm_bindgen]
pub fn webgpu_supported() -> bool {
    WebGPUContext::is_supported()
}

/// Get recommended compute backend based on browser capabilities
#[wasm_bindgen]
pub fn recommended_backend() -> ComputeBackend {
    if webgpu_supported() {
        ComputeBackend::WebGpu
    } else {
        ComputeBackend::Cpu
    }
}

/// Get browser GPU capabilities summary
#[wasm_bindgen]
pub async fn get_gpu_capabilities() -> Result<JsValue, JsError> {
    if !webgpu_supported() {
        let info: HashMap<String, String> = HashMap::from([
            ("status".to_string(), "not_supported".to_string()),
            (
                "message".to_string(),
                "WebGPU is not available in this browser".to_string(),
            ),
        ]);
        return serde_wasm_bindgen::to_value(&info).map_err(|e| JsError::new(&e.to_string()));
    }

    let ctx = WebGPUContext::new().await?;

    let mut info: HashMap<String, String> = HashMap::new();
    info.insert("status".to_string(), format!("{:?}", ctx.status()));

    if let Some(adapter_info) = ctx.adapter_info() {
        info.insert("vendor".to_string(), adapter_info.vendor);
        info.insert("architecture".to_string(), adapter_info.architecture);
        info.insert("device".to_string(), adapter_info.device);
        info.insert("description".to_string(), adapter_info.description);
        info.insert(
            "is_fallback".to_string(),
            adapter_info.is_fallback.to_string(),
        );
    }

    let limits = ctx.limits();
    info.insert(
        "max_buffer_size_mb".to_string(),
        (limits.max_buffer_size / (1024 * 1024)).to_string(),
    );
    info.insert(
        "max_workgroup_size_x".to_string(),
        limits.max_compute_workgroup_size_x.to_string(),
    );

    serde_wasm_bindgen::to_value(&info).map_err(|e| JsError::new(&e.to_string()))
}

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

    #[test]
    fn test_config_default() {
        let config = InferenceConfig::default();
        assert_eq!(config.max_seq_len, 1024);
        assert!(config.use_kv_cache);
    }

    #[test]
    fn test_config_fast() {
        let config = InferenceConfig::fast();
        assert_eq!(config.max_seq_len, 512);
        assert_eq!(config.quantization_bits, 4);
        assert!(config.streaming);
    }

    #[test]
    fn test_config_builder() {
        let config = InferenceConfig::new()
            .with_max_seq_len(2048)
            .with_temperature(0.5)
            .with_streaming(true);

        assert_eq!(config.max_seq_len, 2048);
        assert_eq!(config.temperature, 0.5);
        assert!(config.streaming);
    }

    #[test]
    fn test_config_webgpu_optimized() {
        let config = InferenceConfig::webgpu_optimized();
        assert_eq!(config.compute_backend, ComputeBackend::WebGpu);
        assert_eq!(config.webgpu_workgroup_size, 256);
        assert!(config.streaming);
    }

    #[test]
    fn test_config_cpu_only() {
        let config = InferenceConfig::cpu_only();
        assert_eq!(config.compute_backend, ComputeBackend::Cpu);
    }

    #[test]
    fn test_config_with_compute_backend() {
        let config = InferenceConfig::new()
            .with_compute_backend(ComputeBackend::WebGpu)
            .with_webgpu_workgroup_size(128);

        assert_eq!(config.compute_backend, ComputeBackend::WebGpu);
        assert_eq!(config.webgpu_workgroup_size, 128);
    }

    #[test]
    fn test_compute_backend_default() {
        let backend = ComputeBackend::default();
        assert_eq!(backend, ComputeBackend::Auto);
    }
}