//! Candle-based Embedding Engine
//!
//! Provides built-in text embedding using the Candle ML framework.
//! Supports various sentence transformer models from HuggingFace Hub.

use std::path::Path;

use candle_core::{Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::bert::{BertModel, Config as BertConfig, DTYPE};
use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};
use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer, TruncationParams};

use super::config::PerceptionConfig;
use super::error::PerceptionError;
use super::traits::{Modality, PerceptionEngine};

// ============================================================================
// Model Configuration
// ============================================================================

/// Known model configurations with their dimensions
#[derive(Debug, Clone)]
pub struct ModelInfo {
    pub model_id: String,
    pub dimension: usize,
    pub max_length: usize,
    pub normalize: bool,
}

impl ModelInfo {
    /// Get model info from a HuggingFace model ID
    pub fn from_model_id(model_id: &str) -> Self {
        match model_id {
            "sentence-transformers/all-MiniLM-L6-v2" => Self {
                model_id: model_id.to_string(),
                dimension: 384,
                max_length: 256,
                normalize: true,
            },
            "sentence-transformers/all-MiniLM-L12-v2" => Self {
                model_id: model_id.to_string(),
                dimension: 384,
                max_length: 256,
                normalize: true,
            },
            "BAAI/bge-small-en-v1.5" => Self {
                model_id: model_id.to_string(),
                dimension: 384,
                max_length: 512,
                normalize: true,
            },
            "BAAI/bge-base-en-v1.5" => Self {
                model_id: model_id.to_string(),
                dimension: 768,
                max_length: 512,
                normalize: true,
            },
            "intfloat/multilingual-e5-small" => Self {
                model_id: model_id.to_string(),
                dimension: 384,
                max_length: 512,
                normalize: true,
            },
            "intfloat/e5-small-v2" => Self {
                model_id: model_id.to_string(),
                dimension: 384,
                max_length: 512,
                normalize: true,
            },
            _ => Self {
                model_id: model_id.to_string(),
                dimension: 384, // Default assumption
                max_length: 512,
                normalize: true,
            },
        }
    }
}

// ============================================================================
// Candle Embedding Engine
// ============================================================================

/// Candle-based embedding engine using BERT/Sentence Transformers
pub struct CandleEmbeddingEngine {
    model: BertModel,
    tokenizer: Tokenizer,
    device: Device,
    model_info: ModelInfo,
    normalize: bool,
}

impl CandleEmbeddingEngine {
    /// Create engine from a PerceptionConfig
    pub fn from_config(config: &PerceptionConfig) -> Result<Self, PerceptionError> {
        let device = Self::get_device(&config.device)?;
        let _model_info = ModelInfo::from_model_id(&config.model_id);

        if config.auto_download {
            Self::from_pretrained(&config.model_id, device)
        } else if let Some(ref cache_dir) = config.cache_dir {
            Self::from_local(Path::new(cache_dir), device)
        } else {
            Err(PerceptionError::ConfigError(
                "Either auto_download must be true or cache_dir must be set".to_string(),
            ))
        }
    }

    /// Create engine from a pretrained HuggingFace model
    pub fn from_pretrained(model_id: &str, device: Device) -> Result<Self, PerceptionError> {
        tracing::info!("Loading model from HuggingFace Hub: {}", model_id);

        let model_info = ModelInfo::from_model_id(model_id);
        // Use ApiBuilder::from_env() to support HF_ENDPOINT environment variable for mirror sites
        let api = ApiBuilder::from_env()
            .with_progress(true)
            .build()
            .map_err(|e| PerceptionError::ModelDownloadError(e.to_string()))?;
        let repo = api.repo(Repo::new(model_id.to_string(), RepoType::Model));

        // Download model files
        let config_path = repo
            .get("config.json")
            .map_err(|e| PerceptionError::ModelDownloadError(e.to_string()))?;
        let tokenizer_path = repo
            .get("tokenizer.json")
            .map_err(|e| PerceptionError::ModelDownloadError(e.to_string()))?;
        let weights_path = repo
            .get("model.safetensors")
            .or_else(|_| repo.get("pytorch_model.bin"))
            .map_err(|e| PerceptionError::ModelDownloadError(e.to_string()))?;

        Self::load_from_files(&config_path, &tokenizer_path, &weights_path, device, model_info)
    }

    /// Create engine from local model files
    pub fn from_local(model_dir: &Path, device: Device) -> Result<Self, PerceptionError> {
        tracing::info!("Loading model from local path: {:?}", model_dir);

        let config_path = model_dir.join("config.json");
        let tokenizer_path = model_dir.join("tokenizer.json");
        let weights_path = if model_dir.join("model.safetensors").exists() {
            model_dir.join("model.safetensors")
        } else {
            model_dir.join("pytorch_model.bin")
        };

        // Try to infer model info from directory name
        let model_id = model_dir
            .file_name()
            .and_then(|s| s.to_str())
            .unwrap_or("unknown");
        let model_info = ModelInfo::from_model_id(model_id);

        Self::load_from_files(&config_path, &tokenizer_path, &weights_path, device, model_info)
    }

    /// Load model from specific file paths
    fn load_from_files(
        config_path: &Path,
        tokenizer_path: &Path,
        weights_path: &Path,
        device: Device,
        model_info: ModelInfo,
    ) -> Result<Self, PerceptionError> {
        // Load config
        let config_str = std::fs::read_to_string(config_path)
            .map_err(|e| PerceptionError::IoError(e))?;
        let config: BertConfig = serde_json::from_str(&config_str)
            .map_err(|e| PerceptionError::ConfigError(e.to_string()))?;

        // Load tokenizer
        let mut tokenizer = Tokenizer::from_file(tokenizer_path)
            .map_err(|e| PerceptionError::TokenizationError(e.to_string()))?;

        // Configure padding and truncation
        let padding = PaddingParams {
            strategy: PaddingStrategy::BatchLongest,
            ..Default::default()
        };
        let truncation = TruncationParams {
            max_length: model_info.max_length,
            ..Default::default()
        };
        tokenizer.with_padding(Some(padding));
        tokenizer.with_truncation(Some(truncation))
            .map_err(|e| PerceptionError::TokenizationError(e.to_string()))?;

        // Load weights
        let vb = if weights_path.extension().map_or(false, |e| e == "safetensors") {
            unsafe {
                VarBuilder::from_mmaped_safetensors(&[weights_path], DTYPE, &device)
                    .map_err(|e| PerceptionError::InferenceError(e.to_string()))?
            }
        } else {
            VarBuilder::from_pth(weights_path, DTYPE, &device)
                .map_err(|e| PerceptionError::InferenceError(e.to_string()))?
        };

        // Create model
        let model = BertModel::load(vb, &config)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        tracing::info!(
            "Model loaded successfully: {} (dim={})",
            model_info.model_id,
            model_info.dimension
        );

        Ok(Self {
            model,
            tokenizer,
            device,
            normalize: model_info.normalize,
            model_info,
        })
    }

    /// 从字符串获取计算设备 (Phase-62: 使用统一 DeviceSelector)
    ///
    /// 支持 `"auto"` 模式：自动检测 CUDA → Metal → CPU，
    /// GPU 不可用时安全降级到 CPU 并记录日志。
    ///
    /// | 输入 | 行为 |
    /// |---|---|
    /// | `"auto"` | CUDA → Metal → CPU 自动选择 |
    /// | `"cpu"` | 强制 CPU |
    /// | `"cuda"` / `"cuda:0"` | 尝试 CUDA:0，失败 fallback CPU |
    /// | `"cuda:N"` | 尝试 CUDA:N，失败 fallback CPU |
    /// | `"metal"` | 尝试 Metal，失败 fallback CPU |
    fn get_device(device_str: &str) -> Result<Device, PerceptionError> {
        use crate::device::{DevicePreference, DeviceSelector, GpuAccelConfig};

        let pref = DevicePreference::from_str_flexible(device_str);
        let config = GpuAccelConfig {
            device_preference: pref,
            ..Default::default()
        };
        let selector = DeviceSelector::new(config);
        let (device, is_gpu) = selector.candle_device_with_fallback("candle_embedding_inference");

        if is_gpu {
            tracing::info!(
                engine = "CandleEmbeddingEngine",
                device = ?device,
                config_device = device_str,
                "Using GPU for embedding inference"
            );
        } else {
            tracing::info!(
                engine = "CandleEmbeddingEngine",
                config_device = device_str,
                "Using CPU for embedding inference"
            );
        }

        Ok(device)
    }

    /// Tokenize a batch of texts
    fn tokenize(&self, texts: &[&str]) -> Result<(Tensor, Tensor), PerceptionError> {
        let encodings = self
            .tokenizer
            .encode_batch(texts.to_vec(), true)
            .map_err(|e| PerceptionError::TokenizationError(e.to_string()))?;

        let max_len = encodings.iter().map(|e| e.get_ids().len()).max().unwrap_or(0);

        let mut input_ids: Vec<u32> = Vec::with_capacity(texts.len() * max_len);
        let mut attention_mask: Vec<u32> = Vec::with_capacity(texts.len() * max_len);

        for encoding in &encodings {
            let ids = encoding.get_ids();
            let mask = encoding.get_attention_mask();
            input_ids.extend_from_slice(ids);
            attention_mask.extend_from_slice(mask);
        }

        let input_ids = Tensor::from_vec(input_ids, (texts.len(), max_len), &self.device)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;
        let attention_mask = Tensor::from_vec(attention_mask, (texts.len(), max_len), &self.device)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        Ok((input_ids, attention_mask))
    }

    /// Mean pooling over token embeddings
    fn mean_pooling(
        &self,
        token_embeddings: &Tensor,
        attention_mask: &Tensor,
    ) -> Result<Tensor, PerceptionError> {
        // Expand attention mask to embedding dimension
        let mask = attention_mask
            .unsqueeze(2)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?
            .to_dtype(token_embeddings.dtype())
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?
            .broadcast_as(token_embeddings.shape())
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        // Apply mask and sum
        let masked = token_embeddings
            .mul(&mask)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;
        let summed = masked
            .sum(1)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        // Divide by sum of mask (number of valid tokens)
        let mask_sum = mask
            .sum(1)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?
            .clamp(1e-9, f64::MAX)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        summed
            .div(&mask_sum)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))
    }

    /// L2 normalize embeddings
    fn l2_normalize(&self, embeddings: &Tensor) -> Result<Tensor, PerceptionError> {
        let norm = embeddings
            .sqr()
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?
            .sum_keepdim(1)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?
            .sqrt()
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?
            .clamp(1e-9, f64::MAX)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        embeddings
            .broadcast_div(&norm)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))
    }

    /// Run inference on a batch of texts (Phase-62: GPU fallback 支持)
    ///
    /// 如果 GPU 推理失败（OOM、驱动错误等），自动将输入张量
    /// 移动到 CPU 上重试，确保服务不中断。
    fn encode_batch_internal(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, PerceptionError> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }

        // Tokenize
        let (input_ids, attention_mask) = self.tokenize(texts)?;

        // Get token type IDs (all zeros for single sentence)
        let token_type_ids = input_ids
            .zeros_like()
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        // Run model inference (with GPU fallback)
        let output = match self
            .model
            .forward(&input_ids, &token_type_ids, Some(&attention_mask))
        {
            Ok(output) => output,
            Err(e) if !matches!(self.device, Device::Cpu) => {
                // GPU 推理失败 → 尝试在 CPU 上重试
                tracing::warn!(
                    engine = "CandleEmbeddingEngine",
                    device = ?self.device,
                    error = %e,
                    batch_size = texts.len(),
                    "GPU inference failed, attempting CPU fallback for tokenized tensors"
                );

                // 将输入张量移到 CPU
                let cpu_input_ids = input_ids
                    .to_device(&Device::Cpu)
                    .map_err(|e| PerceptionError::InferenceError(
                        format!("Failed to move input_ids to CPU: {e}")
                    ))?;
                let cpu_attention_mask = attention_mask
                    .to_device(&Device::Cpu)
                    .map_err(|e| PerceptionError::InferenceError(
                        format!("Failed to move attention_mask to CPU: {e}")
                    ))?;
                let cpu_token_type_ids = cpu_input_ids
                    .zeros_like()
                    .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

                // 注意：模型权重仍在 GPU 上，CPU forward 可能也会失败
                // 这种情况下返回原始错误 + CPU 重试失败的信息
                self.model
                    .forward(&cpu_input_ids, &cpu_token_type_ids, Some(&cpu_attention_mask))
                    .map_err(|cpu_err| {
                        crate::device::GpuFallbackEvent {
                            reason: crate::device::FallbackReason::RuntimeError,
                            detail: format!(
                                "GPU error: {e}; CPU retry also failed: {cpu_err}. \
                                 Consider using device: \"cpu\" in config to load model on CPU."
                            ),
                            operation: "candle_embedding_forward".into(),
                            requested_device: crate::device::DevicePreference::Auto,
                        }.log();
                        PerceptionError::InferenceError(
                            format!("GPU inference failed ({e}) and CPU fallback failed ({cpu_err})")
                        )
                    })?
            }
            Err(e) => {
                return Err(PerceptionError::InferenceError(e.to_string()));
            }
        };

        // Mean pooling
        let pooled = self.mean_pooling(&output, &attention_mask)?;

        // Normalize if configured
        let final_embeddings = if self.normalize {
            self.l2_normalize(&pooled)?
        } else {
            pooled
        };

        // Convert to Vec<Vec<f32>>
        let embeddings_vec: Vec<Vec<f32>> = final_embeddings
            .to_vec2()
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        Ok(embeddings_vec)
    }
}

impl PerceptionEngine for CandleEmbeddingEngine {
    fn conceptualize(&self, input: &str) -> Result<Vec<f32>, PerceptionError> {
        let results = self.encode_batch_internal(&[input])?;
        results
            .into_iter()
            .next()
            .ok_or_else(|| PerceptionError::InferenceError("No embedding generated".to_string()))
    }

    fn conceptualize_batch(&self, inputs: &[&str]) -> Result<Vec<Vec<f32>>, PerceptionError> {
        self.encode_batch_internal(inputs)
    }

    fn dimension(&self) -> usize {
        self.model_info.dimension
    }

    fn supported_modalities(&self) -> Vec<Modality> {
        vec![Modality::Text, Modality::Code]
    }

    fn model_name(&self) -> &str {
        &self.model_info.model_id
    }

    fn max_input_length(&self) -> Option<usize> {
        Some(self.model_info.max_length)
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    // --- Phase-62: DeviceSelector 集成测试 ---

    #[test]
    fn test_get_device_cpu() {
        let device = CandleEmbeddingEngine::get_device("cpu").unwrap();
        assert!(matches!(device, Device::Cpu));
    }

    #[test]
    fn test_get_device_auto_never_panics() {
        // auto 模式永远不 panic — 最差返回 CPU
        let device = CandleEmbeddingEngine::get_device("auto").unwrap();
        let _ = device; // 不 panic 就是成功
    }

    #[test]
    fn test_get_device_invalid_cuda_fallback() {
        // 不存在的 CUDA 设备 → DeviceSelector 自动 fallback
        let device = CandleEmbeddingEngine::get_device("cuda:99").unwrap();
        let _ = device; // 不 panic 就是成功
    }

    #[test]
    fn test_get_device_unknown_string_fallback() {
        // 未知字符串 → DevicePreference::Auto → 不 panic
        let device = CandleEmbeddingEngine::get_device("tpu").unwrap();
        let _ = device;
    }

    // Note: These tests require model download, so they're ignored by default
    // Run with: cargo test --features candle -- --ignored

    #[test]
    #[ignore]
    fn test_model_info() {
        let info = ModelInfo::from_model_id("sentence-transformers/all-MiniLM-L6-v2");
        assert_eq!(info.dimension, 384);
        assert_eq!(info.max_length, 256);
    }

    #[test]
    #[ignore]
    fn test_load_model() {
        let engine =
            CandleEmbeddingEngine::from_pretrained("sentence-transformers/all-MiniLM-L6-v2", Device::Cpu);
        assert!(engine.is_ok());
    }

    #[test]
    #[ignore]
    fn test_embedding() {
        let engine = CandleEmbeddingEngine::from_pretrained(
            "sentence-transformers/all-MiniLM-L6-v2",
            Device::Cpu,
        )
        .unwrap();

        let embedding = engine.conceptualize("Hello, world!").unwrap();
        assert_eq!(embedding.len(), 384);

        // Check normalization
        let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 1e-4);
    }

    #[test]
    #[ignore]
    fn test_batch_embedding() {
        let engine = CandleEmbeddingEngine::from_pretrained(
            "sentence-transformers/all-MiniLM-L6-v2",
            Device::Cpu,
        )
        .unwrap();

        let texts = vec!["Hello, world!", "How are you?", "Test sentence"];
        let embeddings = engine.conceptualize_batch(&texts).unwrap();

        assert_eq!(embeddings.len(), 3);
        for emb in &embeddings {
            assert_eq!(emb.len(), 384);
        }
    }
}
