#![cfg(feature = "inference")]

//! Sleipnir — Odin's 8-legged horse; here, the continuous-time manifold driver.
//! Implements Phase 10 [S1] Liquid Cognitive Core.
//!
//! Uses the LFM2.5-VL-450M model (sleipnir.gguf) for semantic embedding.
//! Embedding dimension = 1024 (LFM2 hidden_size).

use mimirs_core::{DensityMemory, MimirError};
use mimirs_wyrd::{Device, Tensor};
use std::path::Path;
use std::sync::Mutex;
use tokenizers::Tokenizer;

/// SleipnirEngine - The Liquid Foundation Model cognitive driver.
///
/// The underlying Lfm2Model is loaded once at construction and shared
/// across all encode() calls via an Mutex, eliminating repeated disk I/O.
/// Multiple encode() calls can run concurrently since the model is read-only
/// during inference (forward pass only, no mutation).
pub struct SleipnirEngine {
    tokenizer: Tokenizer,
    device: Device,
    /// Cached model — multiple concurrent encode() calls share via Mutex.
    model: Mutex<mimirs_judge::Lfm2Model>,
}

impl SleipnirEngine {
    /// LFM2.5-VL-450M embedding dimension (hidden_size).
    pub const EMBED_DIM: usize = 1024;

    /// Loads the Sleipnir engine (Liquid LFM 2.5) from a GGUF file.
    ///
    /// CONTRACT: The model is loaded exactly once here and cached for the
    /// lifetime of the engine. Subsequent `encode()` calls share this instance.
    pub fn load(model_path: &Path, device: Device) -> Result<Self, MimirError> {
        let mut file = std::fs::File::open(model_path)
            .map_err(|e| MimirError::Internal(format!("Failed to open model file: {}", e)))?;

        let content = mimirs_wyrd::quantized::gguf_file::Content::read(&mut file)
            .map_err(|e| MimirError::Internal(format!("Failed to read GGUF: {}", e)))?;

        let tokenizer = Self::load_tokenizer(model_path, &content, &mut file)?;

        let model = mimirs_judge::Lfm2Model::from_gguf(model_path, &device)
            .map_err(|e| MimirError::Internal(format!("Failed to load LFM2 model: {}", e)))?;

        Ok(Self {
            tokenizer,
            device,
            model: Mutex::new(model),
        })
    }

    /// Load tokenizer from GGUF metadata or fall back to sibling tokenizer.json.
    fn load_tokenizer(
        model_path: &Path,
        content: &mimirs_wyrd::quantized::gguf_file::Content,
        file: &mut std::fs::File,
    ) -> Result<Tokenizer, MimirError> {
        // Try sibling tokenizer.json first (most reliable)
        if let Some(parent) = model_path.parent() {
            let tokenizer_json = parent.join("tokenizer.json");
            if tokenizer_json.exists() {
                return Tokenizer::from_file(&tokenizer_json).map_err(|e| {
                    MimirError::Internal(format!(
                        "Failed to load tokenizer from {:?}: {}",
                        tokenizer_json, e
                    ))
                });
            }
        }

        // Try to extract tokenizer from GGUF metadata
        if let Some(tokens_value) = content.metadata.get("tokenizer.ggml.tokens")
            && let mimirs_wyrd::quantized::gguf_file::Value::Array(tokens) = tokens_value
        {
            return Self::build_tokenizer_from_gguf(tokens, content, file);
        }

        // Last resort: try tokenizer.json / tokenizer.model in model directory
        let model_dir = model_path.parent().unwrap_or(Path::new("."));
        for candidate in &["tokenizer.json", "tokenizer.model"] {
            let path = model_dir.join(candidate);
            if path.exists() {
                return Tokenizer::from_file(&path).map_err(|e| {
                    MimirError::Internal(format!("Failed to load tokenizer from {:?}: {}", path, e))
                });
            }
        }

        Err(MimirError::Internal(
            "No tokenizer found: checked sibling tokenizer.json, tokenizer.model, and GGUF metadata".to_string(),
        ))
    }

    /// Build a Tokenizer from GGUF vocabulary metadata.
    fn build_tokenizer_from_gguf(
        tokens: &[mimirs_wyrd::quantized::gguf_file::Value],
        _content: &mimirs_wyrd::quantized::gguf_file::Content,
        _file: &mut std::fs::File,
    ) -> Result<Tokenizer, MimirError> {
        let token_strings: Vec<String> = tokens
            .iter()
            .filter_map(|v| {
                if let mimirs_wyrd::quantized::gguf_file::Value::String(s) = v {
                    Some(s.clone())
                } else {
                    None
                }
            })
            .collect();

        if token_strings.is_empty() {
            return Err(MimirError::Internal(
                "GGUF tokenizer.ggml.tokens array is empty".to_string(),
            ));
        }

        let mut vocab = serde_json::Map::new();
        for (i, token) in token_strings.iter().enumerate() {
            vocab.insert(token.clone(), serde_json::Value::Number(i.into()));
        }

        let tokenizer_json = serde_json::json!({
            "version": "1.0",
            "model": {
                "type": "BPE",
                "vocab": vocab,
                "merges": []
            }
        });

        let json_bytes = serde_json::to_vec(&tokenizer_json).map_err(|e| {
            MimirError::Internal(format!("Failed to serialize tokenizer JSON: {}", e))
        })?;

        Tokenizer::from_bytes(&json_bytes).map_err(|e| {
            MimirError::Internal(format!("Failed to build tokenizer from GGUF vocab: {}", e))
        })
    }

    /// Encodes text into a semantic vector using the Sleipnir model.
    ///
    /// Runs the text through the cached LFM2 model and extracts the last hidden
    /// state as a 1024-dim semantic embedding vector. The model is locked during
    pub fn encode(&self, text: &str) -> Result<Vec<f32>, MimirError> {
        let mut model = self
            .model
            .lock()
            .map_err(|_| MimirError::Internal("Sleipnir model Mutex poisoned".to_string()))?;

        // Tokenize
        let encoding = self
            .tokenizer
            .encode(text, true)
            .map_err(|e| MimirError::Validation(format!("Tokenization failed: {}", e)))?;

        let token_ids: Vec<u32> = encoding.get_ids().to_vec();
        if token_ids.is_empty() {
            return Err(MimirError::Validation(
                "Empty tokenization result".to_string(),
            ));
        }

        // Convert token_ids to a Tensor
        let input_tensor = Tensor::new(token_ids.as_slice(), &self.device)
            .map_err(|e| MimirError::Internal(format!("Tensor creation failed: {}", e)))?
            .unsqueeze(0)
            .map_err(|e| MimirError::Internal(format!("Unsqueeze failed: {}", e)))?;

        // Run forward pass to get hidden states
        let hidden_states = model
            .forward(&input_tensor, 0)
            .map_err(|e| MimirError::Internal(format!("LFM2 forward pass failed: {}", e)))?;

        // Mean pooling over sequence dimension to get a single 1024-dim vector
        // hidden_states shape: [1, seq_len, hidden_size]
        let seq_len = hidden_states
            .dim(1)
            .map_err(|e| MimirError::Internal(format!("Failed to get seq_len: {}", e)))?;
        let hidden_size = hidden_states
            .dim(2)
            .map_err(|e| MimirError::Internal(format!("Failed to get hidden_size: {}", e)))?;

        let pooled = hidden_states
            .sum(1)
            .map_err(|e| MimirError::Internal(format!("Sum pooling failed: {}", e)))?
            .reshape(&[hidden_size])
            .map_err(|e| MimirError::Internal(format!("Reshape failed: {}", e)))?
            .affine(1.0 / (seq_len as f64), 0.0)
            .map_err(|e| MimirError::Internal(format!("Scale failed: {}", e)))?;

        // L2 normalize
        let l2_norm = pooled
            .sqr()
            .map_err(|e| MimirError::Internal(format!("Sqr failed: {}", e)))?
            .sum_all()
            .map_err(|e| MimirError::Internal(format!("Sum all failed: {}", e)))?
            .to_scalar::<f32>()
            .map_err(|e| MimirError::Internal(format!("To scalar failed: {}", e)))?
            .sqrt();

        let pooled = if l2_norm > 1e-6 {
            pooled
                .affine(1.0 / (l2_norm as f64), 0.0)
                .map_err(|e| MimirError::Internal(format!("Normalize failed: {}", e)))?
        } else {
            pooled
        };

        let vector: Vec<f32> = pooled
            .to_vec1()
            .map_err(|e| MimirError::Internal(format!("Tensor to vec failed: {}", e)))?;

        Ok(vector)
    }

    /// Encodes text into a DensityMemory state using continuous-time hidden dynamics.
    pub fn encode_with_dynamics(&self, text: &str, _t: f32) -> Result<DensityMemory, MimirError> {
        let embedding = self.encode(text)?;
        DensityMemory::from_pure(&embedding)
    }

    /// Predicts temporal complexity to tune reservoir reset rate ε.
    pub fn predict_complexity(&self, text: &str) -> f32 {
        match self.tokenizer.encode(text, true) {
            Ok(tokens) => {
                let ids = tokens.get_ids();
                if ids.is_empty() {
                    return 0.1;
                }
                let unique_tokens = ids.iter().collect::<std::collections::HashSet<_>>().len();
                let complexity = (unique_tokens as f32 / ids.len() as f32).clamp(0.0, 1.0);
                0.1 + (complexity * 0.8)
            }
            Err(_) => 0.5,
        }
    }

    /// Extracts unstructured PHI/PII entities from text.
    ///
    /// WHY: Uses the local LFM2.5 model to identify name/location/org entities that regexes missed.
    ///
    /// CONTRACT: Returns a list of (entity, placeholder) pairs.
    pub fn extract_phi_entities(&self, text: &str) -> Result<Vec<(String, String)>, MimirError> {
        let mut model = self
            .model
            .lock()
            .map_err(|_| MimirError::Internal("Sleipnir model Mutex poisoned".to_string()))?;

        let prompt = format!(
            "Find any unredacted patient names, doctor names, hospital/clinic names, or city/street/state locations in the text.\n\
             Format the response as exactly:\n\
             Name: <value>\n\
             Location: <value>\n\
             Organization: <value>\n\n\
             Text: \"{}\"\n\
             Response:\n",
            text
        );

        // Tokenize
        let encoding = self
            .tokenizer
            .encode(prompt.as_str(), true)
            .map_err(|e| MimirError::Validation(format!("Tokenization failed: {}", e)))?;
        let token_ids = encoding.get_ids();

        // Generate response (limit to 100 tokens)
        let generated_ids = model
            .generate(token_ids, 100)
            .map_err(|e| MimirError::Internal(format!("Generation failed: {}", e)))?;

        // Decode
        let decoded = self
            .tokenizer
            .decode(&generated_ids, true)
            .map_err(|e| MimirError::Internal(format!("Decode failed: {}", e)))?;

        // Parse lines
        let mut entities = Vec::new();
        for line in decoded.lines() {
            let line = line.trim();
            #[allow(clippy::collapsible_if)]
            if let Some(val) = line.strip_prefix("Name:").map(|s| s.trim().to_string()) {
                if !val.is_empty() && val != "<value>" {
                    entities.push((val, "[REDACTED_NAME]".to_string()));
                }
            } else if let Some(val) = line.strip_prefix("Location:").map(|s| s.trim().to_string()) {
                if !val.is_empty() && val != "<value>" {
                    entities.push((val, "[REDACTED_LOCATION]".to_string()));
                }
            } else if let Some(val) = line
                .strip_prefix("Organization:")
                .map(|s| s.trim().to_string())
            {
                if !val.is_empty() && val != "<value>" {
                    entities.push((val, "[REDACTED_ORG]".to_string()));
                }
            }
        }

        Ok(entities)
    }
}
