//! LLMHostsTransformer: complete forward pass for Llama-family architectures.
//!
//! # Concurrency Model
//!
//! `LLMHostsModel` holds **shared, immutable** state after loading:
//! - `weights` — the 4-40 GB tensor data, read-only during inference
//! - `config`, `cos_table`, `sin_table` — immutable metadata
//! - `backend` — `Send + Sync`, all methods take `&self`
//!
//! **Mutable, per-request** state lives in `InferenceSession`:
//! - `kv_caches` — grows during generation, one per concurrent request
//!
//! This allows `LLMHostsEngine` in `lib.rs` to hold the model behind
//! `Arc<RwLock<Option<LLMHostsModel>>>` (write lock only for load/unload,
//! read lock during inference — fully concurrent).

use crate::attention;
use crate::compute::auto::detect_best_backend;
use crate::compute::ComputeBackend;
use crate::kv_cache::KvCache;
use crate::memory::GpuAllocator;
use crate::model::config::LLMHostsModelConfig;
use crate::model::weights::{load_tensor, WeightError};
use crate::gguf::{GgufHeader, MappedWeights};
use crate::tensor::{gather_rows, Tensor};
use crate::ffn::transpose_2d;

// ---------------------------------------------------------------------------
// Per-request mutable state — one InferenceSession per concurrent request
// ---------------------------------------------------------------------------

/// Per-request KV cache state. Cheap to allocate; one per concurrent request.
///
/// Create via [`LLMHostsModel::new_session`]. Pass to every `forward()` call.
/// Drop when the request is complete.
pub struct InferenceSession {
    pub kv_caches: Vec<KvCache>,
}

impl InferenceSession {
    /// Allocate fresh KV caches for all transformer layers.
    pub fn new(config: &LLMHostsModelConfig) -> Self {
        let kv_caches = (0..config.num_layers)
            .map(|_| KvCache::new(config.max_seq_len, config.num_kv_heads, config.head_dim))
            .collect();
        Self { kv_caches }
    }

    /// Reset all KV caches (clear generation state between calls).
    pub fn reset(&mut self) {
        for cache in &mut self.kv_caches {
            cache.reset();
        }
    }

    /// Total bytes used by all KV cache layers in this session.
    pub fn kv_cache_bytes(&self) -> usize {
        self.kv_caches.iter().map(|c| c.size_bytes()).sum()
    }
}

/// All weights for a LLMHostsTransformer model, dequantized to f32.
pub struct LLMHostsModelWeights {
    pub token_embd: Tensor,
    pub output_norm: Tensor,
    pub output: Tensor,
    pub layers: Vec<LLMHostsLayerWeights>,
}

pub struct LLMHostsLayerWeights {
    pub attn_norm: Tensor,
    pub attn_q: Tensor,
    pub attn_k: Tensor,
    pub attn_v: Tensor,
    pub attn_q_bias: Option<Tensor>,
    pub attn_k_bias: Option<Tensor>,
    pub attn_v_bias: Option<Tensor>,
    pub attn_output: Tensor,
    pub ffn_norm: Tensor,
    pub ffn_gate: Tensor,
    pub ffn_up: Tensor,
    pub ffn_down: Tensor,
}

impl LLMHostsModelWeights {
    pub fn from_gguf(
        config: &LLMHostsModelConfig,
        header: &GgufHeader,
        mapped: &MappedWeights,
    ) -> Result<Self, WeightError> {
        let token_embd = load_tensor("token_embd.weight", header, mapped)?;
        let output_norm = load_tensor("output_norm.weight", header, mapped)?;
        // Tied embeddings: if output.weight is not found, fall back to token_embd.weight
        let output = match load_tensor("output.weight", header, mapped) {
            Ok(t) => t,
            Err(WeightError::NotFound(_)) => load_tensor("token_embd.weight", header, mapped)?,
            Err(e) => return Err(e),
        };

        // Pre-transpose output weight at load time (GGUF stores [out, in], we need [in, out])
        let output = transpose_2d(&output);

        let mut layers = Vec::with_capacity(config.num_layers);
        for l in 0..config.num_layers {
            // Pre-transpose projection weights at load time to avoid per-token transposes
            let attn_q = transpose_2d(&load_tensor(&format!("blk.{l}.attn_q.weight"), header, mapped)?);
            let attn_k = transpose_2d(&load_tensor(&format!("blk.{l}.attn_k.weight"), header, mapped)?);
            let attn_v = transpose_2d(&load_tensor(&format!("blk.{l}.attn_v.weight"), header, mapped)?);
            let attn_output = transpose_2d(&load_tensor(&format!("blk.{l}.attn_output.weight"), header, mapped)?);
            let ffn_gate = transpose_2d(&load_tensor(&format!("blk.{l}.ffn_gate.weight"), header, mapped)?);
            let ffn_up = transpose_2d(&load_tensor(&format!("blk.{l}.ffn_up.weight"), header, mapped)?);
            let ffn_down = transpose_2d(&load_tensor(&format!("blk.{l}.ffn_down.weight"), header, mapped)?);
            layers.push(LLMHostsLayerWeights {
                attn_norm: load_tensor(&format!("blk.{l}.attn_norm.weight"), header, mapped)?,
                attn_q,
                attn_k,
                attn_v,
                attn_q_bias: load_tensor(&format!("blk.{l}.attn_q.bias"), header, mapped).ok(),
                attn_k_bias: load_tensor(&format!("blk.{l}.attn_k.bias"), header, mapped).ok(),
                attn_v_bias: load_tensor(&format!("blk.{l}.attn_v.bias"), header, mapped).ok(),
                attn_output,
                ffn_norm: load_tensor(&format!("blk.{l}.ffn_norm.weight"), header, mapped)?,
                ffn_gate,
                ffn_up,
                ffn_down,
            });
        }

        Ok(Self { token_embd, output_norm, output, layers })
    }
}

/// Loaded model weights and shared state — safe to hold behind `Arc<>`.
///
/// All fields are **read-only** after construction. Per-request mutable state
/// (KV cache) lives in [`InferenceSession`] — create one per concurrent request
/// via [`LLMHostsModel::new_session`].
pub struct LLMHostsModel {
    pub config: LLMHostsModelConfig,
    pub weights: LLMHostsModelWeights,
    pub cos_table: Vec<f32>,
    pub sin_table: Vec<f32>,
    /// Compute backend (CPU, CUDA, or WGPU) — auto-detected at model load time.
    backend: Box<dyn ComputeBackend>,
    /// VRAM bookkeeping — tracks memory allocated for weights and KV caches.
    allocator: GpuAllocator,
}

impl LLMHostsModel {
    pub fn new(config: LLMHostsModelConfig, weights: LLMHostsModelWeights) -> Self {
        let backend = detect_best_backend();
        tracing::info!(backend = backend.name(), "Model using compute backend");
        Self::build(config, weights, backend)
    }

    /// Create a model with a specific compute backend (for testing or explicit selection).
    pub fn with_backend(
        config: LLMHostsModelConfig,
        weights: LLMHostsModelWeights,
        backend: Box<dyn ComputeBackend>,
    ) -> Self {
        Self::build(config, weights, backend)
    }

    /// Shared constructor: sets up RoPE tables and weight memory allocator.
    ///
    /// KV cache is NOT pre-allocated here — each [`InferenceSession`] allocates
    /// its own cache on demand. Weight memory is tracked for VRAM budgeting.
    fn build(
        config: LLMHostsModelConfig,
        weights: LLMHostsModelWeights,
        backend: Box<dyn ComputeBackend>,
    ) -> Self {
        let (cos_table, sin_table) =
            crate::rope::precompute_freqs(config.head_dim, config.max_seq_len, config.rope_freq_base);

        // Track only weight memory. KV cache memory is tracked per-session.
        let weight_bytes = GpuAllocator::estimate_weight_bytes(&config);
        let total_estimate = weight_bytes;
        let mut allocator = GpuAllocator::new(backend.name(), total_estimate);
        let _ = allocator.allocate_weights(weight_bytes);
        tracing::info!(
            backend = backend.name(),
            usage = %allocator.usage(),
            "Model weights loaded"
        );

        Self { config, weights, cos_table, sin_table, backend, allocator }
    }

    /// Create a model with a specific VRAM budget for OOM-before-load checking.
    ///
    /// Returns `Err(MemoryError)` if the weight tensors won't fit in the given
    /// VRAM budget. KV cache is per-session and not counted here.
    pub fn with_vram_budget(
        config: LLMHostsModelConfig,
        weights: LLMHostsModelWeights,
        backend: Box<dyn ComputeBackend>,
        vram_bytes: usize,
    ) -> Result<Self, crate::memory::MemoryError> {
        let weight_bytes = GpuAllocator::estimate_weight_bytes(&config);

        let mut allocator = GpuAllocator::new(backend.name(), vram_bytes);
        allocator.allocate_weights(weight_bytes)?;

        let (cos_table, sin_table) =
            crate::rope::precompute_freqs(config.head_dim, config.max_seq_len, config.rope_freq_base);

        tracing::info!(
            backend = backend.name(),
            vram_budget = vram_bytes,
            usage = %allocator.usage(),
            "Model weights loaded within VRAM budget"
        );

        Ok(Self { config, weights, cos_table, sin_table, backend, allocator })
    }

    /// Allocate a fresh [`InferenceSession`] for this model.
    ///
    /// Call once per concurrent request. The session holds the KV cache and
    /// must be passed to every [`forward`](Self::forward) call for that request.
    pub fn new_session(&self) -> InferenceSession {
        InferenceSession::new(&self.config)
    }

    /// Name of the active compute backend.
    pub fn backend_name(&self) -> &str {
        self.backend.name()
    }

    /// Current memory usage snapshot.
    pub fn memory_usage(&self) -> crate::memory::MemoryUsage {
        self.allocator.usage()
    }

    /// Forward pass: token_ids -> logits.
    ///
    /// Takes `&self` (weights are read-only) and `&mut session` (KV cache is
    /// per-request mutable state). Concurrent calls with different sessions are
    /// fully safe — the model itself is never mutated.
    pub fn forward(&self, token_ids: &[u32], session: &mut InferenceSession) -> Tensor {
        let seq_len = token_ids.len();
        let start_pos = session.kv_caches[0].len;
        let c = &self.config;

        // 1. Embedding lookup
        let mut hidden = gather_rows(&self.weights.token_embd, token_ids)
            .expect("embedding gather");

        // 2. Transformer layers
        for l in 0..self.weights.layers.len() {
            let layer_w = &self.weights.layers[l];

            // Attention norm (via backend)
            let normed = self.backend.rmsnorm(&hidden, &layer_w.attn_norm, c.rms_norm_eps);

            // Q, K, V projections (via backend matmul, weights pre-transposed at load time)
            let mut q = self.backend.matmul(&normed, &layer_w.attn_q);
            let mut k_new = self.backend.matmul(&normed, &layer_w.attn_k);
            let mut v_new = self.backend.matmul(&normed, &layer_w.attn_v);

            // Add attention biases if present (e.g. Qwen2.5)
            if let Some(ref bias) = layer_w.attn_q_bias {
                q.add_inplace(bias).expect("q bias");
            }
            if let Some(ref bias) = layer_w.attn_k_bias {
                k_new.add_inplace(bias).expect("k bias");
            }
            if let Some(ref bias) = layer_w.attn_v_bias {
                v_new.add_inplace(bias).expect("v bias");
            }

            // Apply RoPE to Q and K (via backend)
            self.backend.rope(
                &mut q,
                &mut k_new,
                &self.cos_table,
                &self.sin_table,
                c.num_attention_heads,
                c.num_kv_heads,
                c.head_dim,
                start_pos,
            );

            // Update session KV cache
            session.kv_caches[l].append(&k_new.data, &v_new.data, seq_len);
            let kv_len = session.kv_caches[l].len;

            // Build full K and V tensors from cache
            let k_full = Tensor {
                data: session.kv_caches[l].k_slice().to_vec(),
                shape: vec![kv_len, c.num_kv_heads * c.head_dim],
            };
            let v_full = Tensor {
                data: session.kv_caches[l].v_slice().to_vec(),
                shape: vec![kv_len, c.num_kv_heads * c.head_dim],
            };

            // Attention (CPU — Flash Attention kernel in F2)
            let attn_out = attention::attention(
                &q,
                &k_full,
                &v_full,
                &attention::AttentionParams {
                    num_heads: c.num_attention_heads,
                    num_kv_heads: c.num_kv_heads,
                    head_dim: c.head_dim,
                    seq_len,
                    kv_len,
                    causal: true,
                    start_pos,
                },
            );

            // Output projection + residual (via backend matmul)
            let projected = self.backend.matmul(&attn_out, &layer_w.attn_output);
            hidden.add_inplace(&projected).expect("attn residual");

            // FFN norm (via backend)
            let ffn_normed = self.backend.rmsnorm(&hidden, &layer_w.ffn_norm, c.rms_norm_eps);

            // FFN: gate and up projections (via backend matmul)
            let mut gate = self.backend.matmul(&ffn_normed, &layer_w.ffn_gate);
            let up = self.backend.matmul(&ffn_normed, &layer_w.ffn_up);

            // SiLU gate (via backend)
            self.backend.silu_elementwise_mul(&mut gate, &up);

            // Down projection (via backend matmul)
            let ffn_out = self.backend.matmul(&gate, &layer_w.ffn_down);
            hidden.add_inplace(&ffn_out).expect("ffn residual");
        }

        // 3. Final norm (via backend)
        let normed = self.backend.rmsnorm(&hidden, &self.weights.output_norm, c.rms_norm_eps);

        // 4. LM head: hidden -> logits (via backend matmul)
        self.backend.matmul(&normed, &self.weights.output)
    }
}

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

    fn load_test_model() -> (LLMHostsModelConfig, LLMHostsModel) {
        let dir = tempfile::tempdir().unwrap();
        let config_gen = crate::test_utils::TestModelConfig::default();
        let path = crate::test_utils::generate_test_gguf(dir.path(), &config_gen);
        let header = crate::gguf::parse_header(&path).unwrap();
        let mapped = crate::gguf::MappedWeights::open(&path, header.data_offset).unwrap();
        let config = crate::model::config::LLMHostsModelConfig::from_gguf(&header).unwrap();
        let weights = LLMHostsModelWeights::from_gguf(&config, &header, &mapped).unwrap();
        let model = LLMHostsModel::new(config.clone(), weights);
        (config, model)
    }

    #[test]
    fn test_forward_with_test_model() {
        let (config, model) = load_test_model();
        let mut session = model.new_session();
        let logits = model.forward(&[1, 5, 10], &mut session);
        assert_eq!(logits.shape, vec![3, config.vocab_size]);
        assert!(logits.data.iter().all(|v| v.is_finite()), "non-finite logits detected");
    }

    #[test]
    fn test_decode_step_uses_kv_cache() {
        let (config, model) = load_test_model();
        let mut session = model.new_session();
        let _ = model.forward(&[1, 5, 10], &mut session);
        assert_eq!(session.kv_caches[0].len, 3);
        let logits = model.forward(&[20], &mut session);
        assert_eq!(logits.shape, vec![1, config.vocab_size]);
        assert_eq!(session.kv_caches[0].len, 4);
    }

    #[test]
    fn test_reset_clears_kv_cache() {
        let (_config, model) = load_test_model();
        let mut session = model.new_session();
        let _ = model.forward(&[1, 5, 10], &mut session);
        session.reset();
        assert_eq!(session.kv_caches[0].len, 0);
    }

    #[test]
    fn test_backend_name_is_cpu_by_default() {
        let (_config, model) = load_test_model();
        assert_eq!(model.backend_name(), "cpu");
    }

    #[test]
    fn test_with_explicit_backend() {
        let dir = tempfile::tempdir().unwrap();
        let config_gen = crate::test_utils::TestModelConfig::default();
        let path = crate::test_utils::generate_test_gguf(dir.path(), &config_gen);
        let header = crate::gguf::parse_header(&path).unwrap();
        let mapped = crate::gguf::MappedWeights::open(&path, header.data_offset).unwrap();
        let config = crate::model::config::LLMHostsModelConfig::from_gguf(&header).unwrap();
        let weights = LLMHostsModelWeights::from_gguf(&config, &header, &mapped).unwrap();

        let cpu = Box::new(crate::compute::cpu::CpuBackend::new());
        let model = LLMHostsModel::with_backend(config.clone(), weights, cpu);
        assert_eq!(model.backend_name(), "cpu");

        let mut session = model.new_session();
        let logits = model.forward(&[1, 5, 10], &mut session);
        assert_eq!(logits.shape, vec![3, config.vocab_size]);
        assert!(logits.data.iter().all(|v| v.is_finite()));
    }

    #[test]
    fn test_memory_usage_reports_nonzero() {
        let (_config, model) = load_test_model();
        let usage = model.memory_usage();
        assert!(usage.allocated > 0);
        assert!(usage.weights > 0);
        // KV cache no longer pre-allocated in model — lives in InferenceSession
        assert_eq!(usage.kv_cache, 0);
        assert_eq!(usage.allocated, usage.weights + usage.kv_cache);
        assert_eq!(usage.free, usage.total - usage.allocated);
    }

    #[test]
    fn test_inference_session_kv_cache_bytes() {
        let (_config, model) = load_test_model();
        let session = model.new_session();
        assert!(session.kv_cache_bytes() > 0, "session should pre-allocate KV cache buffers");
    }

    #[test]
    fn test_with_vram_budget_sufficient() {
        let dir = tempfile::tempdir().unwrap();
        let config_gen = crate::test_utils::TestModelConfig::default();
        let path = crate::test_utils::generate_test_gguf(dir.path(), &config_gen);
        let header = crate::gguf::parse_header(&path).unwrap();
        let mapped = crate::gguf::MappedWeights::open(&path, header.data_offset).unwrap();
        let config = crate::model::config::LLMHostsModelConfig::from_gguf(&header).unwrap();
        let weights = LLMHostsModelWeights::from_gguf(&config, &header, &mapped).unwrap();

        let cpu = Box::new(crate::compute::cpu::CpuBackend::new());
        let model = LLMHostsModel::with_vram_budget(config, weights, cpu, 100_000_000).unwrap();
        let usage = model.memory_usage();
        assert!(usage.free > 0);
    }

    #[test]
    fn test_with_vram_budget_insufficient() {
        let dir = tempfile::tempdir().unwrap();
        let config_gen = crate::test_utils::TestModelConfig::default();
        let path = crate::test_utils::generate_test_gguf(dir.path(), &config_gen);
        let header = crate::gguf::parse_header(&path).unwrap();
        let mapped = crate::gguf::MappedWeights::open(&path, header.data_offset).unwrap();
        let config = crate::model::config::LLMHostsModelConfig::from_gguf(&header).unwrap();
        let weights = LLMHostsModelWeights::from_gguf(&config, &header, &mapped).unwrap();

        let cpu = Box::new(crate::compute::cpu::CpuBackend::new());
        // Give a tiny budget that can't fit the model
        let result = LLMHostsModel::with_vram_budget(config, weights, cpu, 1024);
        match result {
            Err(err) => {
                let msg = format!("{err}");
                assert!(msg.contains("Out of memory"), "Expected OOM error, got: {msg}");
            }
            Ok(_) => panic!("Expected OOM error but model loaded successfully"),
        }
    }

    /// Verify that the forward pass produces deterministic output across resets.
    #[test]
    fn test_forward_deterministic_across_resets() {
        let (_config, model) = load_test_model();

        let mut session = model.new_session();
        let logits1 = model.forward(&[1, 5, 10], &mut session);
        session.reset();
        let logits2 = model.forward(&[1, 5, 10], &mut session);

        assert_eq!(logits1.shape, logits2.shape);
        for (a, b) in logits1.data.iter().zip(logits2.data.iter()) {
            assert!(
                (a - b).abs() < 1e-5,
                "Non-deterministic: {} vs {} (delta: {})",
                a, b, (a - b).abs()
            );
        }
    }

    /// Verify that two independent sessions do not interfere with each other.
    #[test]
    fn test_two_independent_sessions() {
        let (config, model) = load_test_model();

        let mut session_a = model.new_session();
        let mut session_b = model.new_session();

        // Run session A with one prompt
        let logits_a = model.forward(&[1, 5, 10], &mut session_a);
        assert_eq!(session_a.kv_caches[0].len, 3);
        assert_eq!(session_b.kv_caches[0].len, 0, "session B must be unaffected by session A");

        // Run session B with a different prompt
        let logits_b = model.forward(&[2, 7], &mut session_b);
        assert_eq!(session_b.kv_caches[0].len, 2);
        assert_eq!(session_a.kv_caches[0].len, 3, "session A must be unaffected by session B");

        assert_eq!(logits_a.shape, vec![3, config.vocab_size]);
        assert_eq!(logits_b.shape, vec![2, config.vocab_size]);
    }

    /// Concurrent inference: N threads run forward() simultaneously on the same model.
    #[test]
    fn test_concurrent_inference_sessions() {
        use std::sync::Arc;

        let (_config, model) = load_test_model();
        let model = Arc::new(model);

        let handles: Vec<_> = (0..4)
            .map(|i| {
                let model = Arc::clone(&model);
                std::thread::spawn(move || {
                    let mut session = model.new_session();
                    let token_ids: Vec<u32> = vec![1, 2, (i as u32) + 3];
                    let logits = model.forward(&token_ids, &mut session);
                    assert!(!logits.data.is_empty(), "thread {i} got empty logits");
                    assert!(logits.data.iter().all(|v| v.is_finite()), "thread {i} got non-finite logits");
                })
            })
            .collect();

        for h in handles {
            h.join().expect("concurrent inference thread panicked");
        }
    }
}
