//! InferenceEngine: loads a GGUF model via llama-sys FFI and manages the
//! lifetime of `llama_model` + `llama_context`.
//!
//! Design decisions:
//! - Single model, single context per engine instance (Phase 2 design).
//! - `LlamaLib` is loaded once and stored in an `Arc` so future session spawning
//!   can share it without re-loading the dynamic library.
//! - All raw pointers are wrapped in RAII guards that call the matching
//!   `llama_free` / `llama_model_free` on drop.
//! - The engine is `Send` but NOT `Sync` – callers must coordinate access
//!   (the scheduler in a later phase will own this exclusively).

use std::ffi::CString;
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{bail, Context, Result};
use tracing::{info, warn};

use llama_sys::{
    llama_context, llama_model, LlamaLib, LLAMA_FLASH_ATTN_TYPE_DISABLED,
    LLAMA_FLASH_ATTN_TYPE_ENABLED,
};

use crate::profile::RuntimeProfile;
use crate::sampler::SamplerConfig;
use crate::session::GenerationSession;

// ─── Configuration ─────────────────────────────────────────────────────────

/// Configuration for the inference engine.
#[derive(Debug, Clone)]
pub struct EngineConfig {
    /// Path to the GGUF model file.
    pub model_path: PathBuf,
    /// Runtime performance profile (used to fill in defaults).
    pub profile: RuntimeProfile,
    /// How many GPU layers to offload (999 = all).
    pub n_gpu_layers: i32,
    /// Context window size in tokens.
    pub n_ctx: u32,
    /// Physical batch size (prompt processing).
    pub n_batch: u32,
    /// Micro-batch size (generation step).
    pub n_ubatch: u32,
    /// Number of CPU threads for generation.
    pub n_threads: i32,
    /// Number of CPU threads for batch/prefill.
    pub n_threads_batch: i32,
    /// Enable flash attention.
    pub flash_attn: bool,
    /// Use mmap for model loading (false = load fully into RAM/VRAM).
    pub use_mmap: bool,
    /// RNG seed for sampling (0 = random).
    pub seed: u32,
}

impl Default for EngineConfig {
    fn default() -> Self {
        let prof = RuntimeProfile::Turbo4B;
        let d = prof.defaults();
        Self {
            model_path: PathBuf::new(),
            profile: prof,
            n_gpu_layers: d.n_gpu_layers,
            n_ctx: d.n_ctx,
            n_batch: d.n_batch,
            n_ubatch: d.n_ubatch,
            n_threads: d.n_threads,
            n_threads_batch: d.n_threads,
            flash_attn: d.flash_attn,
            use_mmap: false, // --no-mmap recommended for CUDA
            seed: 0,
        }
    }
}

impl EngineConfig {
    /// Build from a profile, applying its defaults.
    pub fn from_profile(profile: RuntimeProfile, model_path: PathBuf) -> Self {
        let d = profile.defaults();
        Self {
            model_path,
            profile,
            n_gpu_layers: d.n_gpu_layers,
            n_ctx: d.n_ctx,
            n_batch: d.n_batch,
            n_ubatch: d.n_ubatch,
            n_threads: d.n_threads,
            n_threads_batch: d.n_threads,
            flash_attn: d.flash_attn,
            use_mmap: false,
            seed: 0,
        }
    }
}

// ─── RAII wrappers ─────────────────────────────────────────────────────────

/// Owns a `*mut llama_model` and frees it on drop.
struct OwnedModel {
    ptr: *mut llama_model,
    lib: Arc<LlamaLib>,
}

impl Drop for OwnedModel {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            unsafe { (self.lib.model_free)(self.ptr) };
        }
    }
}

// SAFETY: llama_model is reference-counted internally; we are the only owner.
unsafe impl Send for OwnedModel {}

/// Owns a `*mut llama_context` and frees it on drop.
pub(crate) struct OwnedContext {
    pub ptr: *mut llama_context,
    pub lib: Arc<LlamaLib>,
}

impl Drop for OwnedContext {
    fn drop(&mut self) {
        if !self.ptr.is_null() {
            unsafe { (self.lib.free)(self.ptr) };
        }
    }
}

// SAFETY: we enforce single-threaded access through InferenceEngine.
unsafe impl Send for OwnedContext {}

// ─── Engine ────────────────────────────────────────────────────────────────

/// Core inference engine.
///
/// - Owns a loaded model and its associated context.
/// - Provides `new_session()` to start a streaming generation.
/// - NOT `Sync` – the scheduler must ensure single-threaded use.
pub struct InferenceEngine {
    pub config: EngineConfig,
    pub(crate) lib: Arc<LlamaLib>,
    _model: OwnedModel,
    pub(crate) ctx: OwnedContext,
}

impl InferenceEngine {
    /// Load a GGUF model and initialise a llama_context.
    ///
    /// Blocks until the model is fully loaded.
    pub fn load(config: EngineConfig) -> Result<Self> {
        if config.model_path.as_os_str().is_empty() {
            bail!("EngineConfig.model_path is empty");
        }
        if !config.model_path.exists() {
            bail!("Model file not found: {}", config.model_path.display());
        }

        let model_name = config
            .model_path
            .file_name()
            .map(|s| s.to_string_lossy().into_owned())
            .unwrap_or_else(|| "<unknown>".into());

        info!(
            model      = %model_name,
            profile    = ?config.profile,
            gpu_layers = config.n_gpu_layers,
            n_ctx      = config.n_ctx,
            n_batch    = config.n_batch,
            n_ubatch   = config.n_ubatch,
            flash_attn = config.flash_attn,
            use_mmap   = config.use_mmap,
            "engine_config"
        );

        let t_total = std::time::Instant::now();

        // 1. Load dynamic library
        let t0 = std::time::Instant::now();
        let lib = Arc::new(LlamaLib::load_auto().context("Failed to load llama.dll")?);

        // 2. Backend init
        unsafe {
            (lib.backend_load_all)();
            (lib.backend_init)();
        }
        info!(
            elapsed_ms = t0.elapsed().as_millis(),
            "llama backend initialised"
        );

        // 3. Model params
        let mut mparams = unsafe { (lib.model_default_params)() };
        mparams.n_gpu_layers = config.n_gpu_layers;
        mparams.use_mmap = config.use_mmap;

        // 4. Load model from file
        let model_path_c = CString::new(config.model_path.to_str().context("Non-UTF8 model path")?)
            .context("CString::new failed")?;

        let t_load = std::time::Instant::now();
        let model_ptr = unsafe { (lib.model_load_from_file)(model_path_c.as_ptr(), mparams) };
        if model_ptr.is_null() {
            bail!("llama_model_load_from_file returned NULL – check model path and VRAM");
        }
        let load_ms = t_load.elapsed().as_millis();

        // Log model info
        {
            let mut desc = vec![0u8; 256];
            unsafe {
                (lib.model_desc)(model_ptr, desc.as_mut_ptr() as *mut _, desc.len());
            }
            let desc_str = std::ffi::CStr::from_bytes_until_nul(&desc)
                .map(|s| s.to_string_lossy().into_owned())
                .unwrap_or_else(|_| "<unknown>".into());
            let n_params = unsafe { (lib.model_n_params)(model_ptr) };
            let is_hybrid = unsafe { (lib.model_is_hybrid)(model_ptr) };
            let is_recurrent = unsafe { (lib.model_is_recurrent)(model_ptr) };
            info!(
                desc         = %desc_str,
                params_b     = format!("{:.2}", n_params as f64 / 1e9),
                is_hybrid,
                is_recurrent,
                load_ms,
                "model_loaded"
            );
        }

        let model = OwnedModel {
            ptr: model_ptr,
            lib: Arc::clone(&lib),
        };

        // 5. Context params
        let mut cparams = unsafe { (lib.context_default_params)() };
        cparams.n_ctx = config.n_ctx;
        cparams.n_batch = config.n_batch;
        cparams.n_ubatch = config.n_ubatch;
        cparams.n_threads = config.n_threads;
        cparams.n_threads_batch = config.n_threads_batch;
        cparams.flash_attn_type = if config.flash_attn {
            LLAMA_FLASH_ATTN_TYPE_ENABLED
        } else {
            LLAMA_FLASH_ATTN_TYPE_DISABLED
        };
        cparams.offload_kqv = true;
        cparams.no_perf = true;

        // 6. Init context
        let ctx_ptr = unsafe { (lib.init_from_model)(model_ptr, cparams) };
        if ctx_ptr.is_null() {
            bail!("llama_init_from_model returned NULL – insufficient VRAM or bad params");
        }

        let actual_ctx = unsafe { (lib.n_ctx)(ctx_ptr) };
        info!(
            n_ctx = actual_ctx,
            elapsed_ms = t_total.elapsed().as_millis(),
            "context_ready"
        );

        let ctx = OwnedContext {
            ptr: ctx_ptr,
            lib: Arc::clone(&lib),
        };

        Ok(Self {
            config,
            lib,
            _model: model,
            ctx,
        })
    }

    /// Warm up: run a single dummy decode to pre-allocate CUDA buffers.
    ///
    /// Call once after `load()` before accepting real requests.
    pub fn warmup(&mut self) -> Result<()> {
        let t = std::time::Instant::now();
        info!("engine_warmup: starting dummy decode to pre-allocate CUDA buffers");

        unsafe {
            (self.lib.set_warmup)(self.ctx.ptr, true);
            // set_warmup may not exist in all builds; call decode with 1 token
            let dummy_token: llama_sys::llama_token = 1; // BOS-ish token
            let mut tokens = vec![dummy_token];
            let batch = (self.lib.batch_get_one)(tokens.as_mut_ptr(), 1);
            let ret = (self.lib.decode)(self.ctx.ptr, batch);
            (self.lib.set_warmup)(self.ctx.ptr, false);
            if ret != 0 {
                warn!(ret, "engine_warmup: decode returned non-zero (non-fatal)");
            }

            let memory = (self.lib.get_memory)(self.ctx.ptr);
            if memory.is_null() {
                warn!("engine_warmup: llama_get_memory returned NULL — KV cache not cleared");
            } else {
                (self.lib.memory_clear)(memory, false);
            }
        }

        info!(elapsed_ms = t.elapsed().as_millis(), "engine_ready");
        Ok(())
    }

    /// Create a new generation session for a single request.
    ///
    /// `prompt` is the fully-formatted string (system + user turns already joined).
    /// The engine must not be accessed concurrently while the session is active.
    pub fn new_session(
        &mut self,
        prompt: &str,
        sampler_cfg: SamplerConfig,
    ) -> Result<GenerationSession> {
        // Clear KV cache from any previous session so the new prefill starts at position 0.
        // `true` forces a full clear including M-RoPE position metadata; `false` only clears
        // cell occupancy flags and leaves position data intact, which causes the
        // "inconsistent sequence positions" error on models that use M-RoPE (e.g. Qwen).
        // Truncate prompt for logging — show first 120 chars to aid debugging without flooding logs.
        let prompt_preview = if prompt.len() > 120 {
            format!("{}…", &prompt[..120])
        } else {
            prompt.to_owned()
        };
        tracing::debug!(prompt_chars = prompt.len(), preview = %prompt_preview, "session.new");

        unsafe {
            let memory = (self.lib.get_memory)(self.ctx.ptr);
            if memory.is_null() {
                tracing::error!("llama_get_memory returned NULL — cannot clear KV cache, refusing to create session to avoid decode corruption");
                anyhow::bail!("KV cache clear failed: llama_get_memory returned NULL");
            } else {
                (self.lib.memory_clear)(memory, true);
                tracing::debug!("kv_cache_cleared");
            }
        }

        let vocab = unsafe { (self.lib.model_get_vocab)(self._model.ptr) };
        GenerationSession::with_vocab(
            Arc::clone(&self.lib),
            self.ctx.ptr,
            vocab,
            self.config.n_batch,
            prompt,
            sampler_cfg,
        )
    }
}

// SAFETY: we are the sole owner of the pointers; no concurrent access.
unsafe impl Send for InferenceEngine {}