//! EmbeddingGemma 推理引擎
//!
//! 实现了 `PerceptionEngine` trait，提供完整的 embedding 生成能力。
//!
//! # 使用示例
//!
//! ```ignore
//! use aifastdb_core::perception::embedding_gemma::{
//!     EmbeddingGemmaEngine, TaskType, MatryoshkaDim
//! };
//!
//! // 从 HuggingFace Hub 加载
//! let engine = EmbeddingGemmaEngine::from_pretrained(
//!     "google/embeddinggemma-300m",
//!     Device::Cpu,
//! )?;
//!
//! // 生成 embedding
//! let embedding = engine.conceptualize("Hello, world!")?;
//!
//! // 使用特定任务类型
//! let engine = engine.with_task(TaskType::SearchDocument);
//! let doc_embedding = engine.conceptualize("Document content...")?;
//! ```

use std::path::Path;

use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};
use tokenizers::{PaddingParams, PaddingStrategy, Tokenizer, TruncationParams};

use super::config::{EmbeddingGemmaConfig, MatryoshkaDim, TaskType};
use super::model::EmbeddingGemmaModel;
use crate::perception::error::PerceptionError;
use crate::perception::traits::{Modality, PerceptionEngine};
use crate::perception::PerceptionConfig;

/// EmbeddingGemma 推理引擎
///
/// 提供完整的文本到 embedding 的转换能力，包括：
/// - 自动任务前缀处理
/// - Matryoshka 维度截断
/// - 批量处理支持
pub struct EmbeddingGemmaEngine {
    /// 底层模型
    model: EmbeddingGemmaModel,
    /// Tokenizer
    tokenizer: Tokenizer,
    /// 计算设备
    device: Device,
    /// 模型配置
    config: EmbeddingGemmaConfig,
    /// Matryoshka 截断维度
    matryoshka_dim: Option<usize>,
    /// 当前任务类型
    task_type: TaskType,
    /// 模型 ID
    model_id: String,
}

impl EmbeddingGemmaEngine {
    /// 从 PerceptionConfig 创建引擎
    pub fn from_config(config: &PerceptionConfig) -> Result<Self, PerceptionError> {
        let device = Self::get_device(&config.device)?;

        if config.auto_download {
            let mut engine = Self::from_pretrained(&config.model_id, device)?;
            
            // 应用 Matryoshka 维度配置
            if let Some(dim) = config.dimension {
                if dim < 768 {
                    engine = engine.with_matryoshka_dim(MatryoshkaDim::from_dim(dim)
                        .ok_or_else(|| PerceptionError::ConfigError(
                            format!("Invalid Matryoshka dimension: {}. Valid values: 768, 512, 256, 128", dim)
                        ))?);
                }
            }
            
            Ok(engine)
        } 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(),
            ))
        }
    }

    /// 从 HuggingFace Hub 下载并加载模型
    ///
    /// # Arguments
    /// * `model_id` - HuggingFace 模型 ID (如 "google/embeddinggemma-300m")
    /// * `device` - 计算设备
    pub fn from_pretrained(model_id: &str, device: Device) -> Result<Self, PerceptionError> {
        tracing::info!("Loading EmbeddingGemma model from HuggingFace Hub: {}", 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));

        // 下载配置文件
        let config_path = repo
            .get("config.json")
            .map_err(|e| PerceptionError::ModelDownloadError(e.to_string()))?;

        // 下载 tokenizer
        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_id.to_string(),
        )
    }

    /// 从本地文件加载模型
    ///
    /// # Arguments
    /// * `model_dir` - 模型目录路径
    /// * `device` - 计算设备
    pub fn from_local(model_dir: &Path, device: Device) -> Result<Self, PerceptionError> {
        tracing::info!("Loading EmbeddingGemma 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")
        };

        let model_id = model_dir
            .file_name()
            .and_then(|s| s.to_str())
            .unwrap_or("local-model")
            .to_string();

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

    /// 从指定文件加载模型
    fn load_from_files(
        config_path: &Path,
        tokenizer_path: &Path,
        weights_path: &Path,
        device: Device,
        model_id: String,
    ) -> Result<Self, PerceptionError> {
        // 加载配置
        let config_str = std::fs::read_to_string(config_path)
            .map_err(|e| PerceptionError::IoError(e))?;
        
        // 尝试解析为 EmbeddingGemmaConfig，如果失败则使用默认配置
        let config: EmbeddingGemmaConfig = serde_json::from_str(&config_str)
            .unwrap_or_else(|_| {
                tracing::warn!("Failed to parse config.json, using default EmbeddingGemma config");
                EmbeddingGemmaConfig::default()
            });

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

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

        // 加载模型权重
        let dtype = DType::F32; // 使用 F32 以获得最佳精度
        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()))?
        };

        // 创建模型
        let model = EmbeddingGemmaModel::new(config.clone(), vb)
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        tracing::info!(
            "EmbeddingGemma model loaded successfully: {} (dim={}, layers={})",
            model_id,
            config.hidden_size,
            config.num_hidden_layers
        );

        Ok(Self {
            model,
            tokenizer,
            device,
            config,
            matryoshka_dim: None,
            task_type: TaskType::default(),
            model_id,
        })
    }

    /// 获取计算设备 (Phase-62: 使用统一 DeviceSelector)
    ///
    /// 支持 `"auto"` 模式：自动检测 CUDA → Metal → CPU，
    /// GPU 不可用时安全降级到 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("embedding_gemma_inference");

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

        Ok(device)
    }

    /// 设置 Matryoshka 截断维度
    ///
    /// # Arguments
    /// * `dim` - 目标维度 (768, 512, 256, 128)
    pub fn with_matryoshka_dim(mut self, dim: MatryoshkaDim) -> Self {
        self.matryoshka_dim = Some(dim.dim());
        self
    }

    /// 设置任务类型
    ///
    /// 不同任务类型会使用不同的前缀，以获得最佳效果。
    pub fn with_task(mut self, task: TaskType) -> Self {
        self.task_type = task;
        self
    }

    /// 设置任务类型（可变引用版本）
    pub fn set_task(&mut self, task: TaskType) {
        self.task_type = task;
    }

    /// 获取当前任务类型
    pub fn task_type(&self) -> TaskType {
        self.task_type
    }

    /// 获取当前 Matryoshka 维度
    pub fn matryoshka_dim(&self) -> Option<usize> {
        self.matryoshka_dim
    }

    /// 对输入文本添加任务前缀
    fn add_task_prefix(&self, text: &str) -> String {
        format!("{}{}", self.task_type.prefix(), text)
    }

    /// Tokenize 输入文本
    fn tokenize(&self, texts: &[&str]) -> Result<(Tensor, Tensor), PerceptionError> {
        // 添加任务前缀
        let prefixed_texts: Vec<String> = texts.iter().map(|t| self.add_task_prefix(t)).collect();
        let text_refs: Vec<&str> = prefixed_texts.iter().map(|s| s.as_str()).collect();

        // Tokenize
        let encodings = self
            .tokenizer
            .encode_batch(text_refs, true)
            .map_err(|e| PerceptionError::TokenizationError(e.to_string()))?;

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

        // 构建 input_ids 和 attention_mask
        let mut input_ids: Vec<u32> = Vec::with_capacity(batch_size * max_len);
        let mut attention_mask: Vec<f32> = Vec::with_capacity(batch_size * max_len);

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

            // 添加 token IDs
            input_ids.extend_from_slice(ids);
            // 填充到 max_len
            input_ids.extend(std::iter::repeat(0u32).take(max_len - ids.len()));

            // 添加 attention mask
            attention_mask.extend(mask.iter().map(|&m| m as f32));
            // 填充到 max_len
            attention_mask.extend(std::iter::repeat(0.0f32).take(max_len - mask.len()));
        }

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

        Ok((input_ids, attention_mask))
    }

    /// 运行推理 (Phase-62: GPU fallback 支持)
    ///
    /// GPU 推理失败时尝试将输入移至 CPU 重试。
    fn encode_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)?;

        // 运行模型（带 GPU fallback）
        let embeddings = match self
            .model
            .encode(&input_ids, &attention_mask, self.matryoshka_dim)
        {
            Ok(emb) => emb,
            Err(e) if !matches!(self.device, Device::Cpu) => {
                // GPU 推理失败 → 尝试 CPU fallback
                tracing::warn!(
                    engine = "EmbeddingGemmaEngine",
                    device = ?self.device,
                    error = %e,
                    batch_size = texts.len(),
                    "GPU inference failed, attempting CPU fallback"
                );

                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}")
                    ))?;

                self.model
                    .encode(&cpu_input_ids, &cpu_attention_mask, self.matryoshka_dim)
                    .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."
                            ),
                            operation: "embedding_gemma_encode".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()));
            }
        };

        // 转换为 Vec<Vec<f32>>
        let embeddings_vec: Vec<Vec<f32>> = embeddings
            .to_vec2()
            .map_err(|e| PerceptionError::InferenceError(e.to_string()))?;

        Ok(embeddings_vec)
    }

    /// 获取模型配置
    pub fn config(&self) -> &EmbeddingGemmaConfig {
        &self.config
    }

    /// 获取计算设备
    pub fn device(&self) -> &Device {
        &self.device
    }
}

impl PerceptionEngine for EmbeddingGemmaEngine {
    fn conceptualize(&self, input: &str) -> Result<Vec<f32>, PerceptionError> {
        let results = self.encode_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_internal(inputs)
    }

    fn dimension(&self) -> usize {
        self.matryoshka_dim.unwrap_or(self.config.hidden_size)
    }

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

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

    fn max_input_length(&self) -> Option<usize> {
        Some(self.config.max_position_embeddings)
    }
}

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

    #[test]
    fn test_task_prefix() {
        // 创建一个模拟的引擎状态来测试前缀
        let task = TaskType::SearchQuery;
        assert_eq!(task.prefix(), "task: search result | query: ");

        let task = TaskType::SearchDocument;
        assert_eq!(task.prefix(), "title: none | text: ");
    }

    #[test]
    fn test_matryoshka_dim() {
        assert_eq!(MatryoshkaDim::Full.dim(), 768);
        assert_eq!(MatryoshkaDim::Dim512.dim(), 512);
        assert_eq!(MatryoshkaDim::Dim256.dim(), 256);
        assert_eq!(MatryoshkaDim::Dim128.dim(), 128);
    }

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

    #[test]
    fn test_get_device_auto_never_panics() {
        // auto 模式永远不 panic，最差返回 CPU
        let device = EmbeddingGemmaEngine::get_device("auto");
        assert!(device.is_ok());
    }

    #[test]
    fn test_get_device_invalid_cuda_fallback() {
        // 不存在的 CUDA 设备 → 应 fallback 到某个可用设备（不 panic）
        let device = EmbeddingGemmaEngine::get_device("cuda:99");
        assert!(device.is_ok());
    }

    #[test]
    fn test_task_type_from_str() {
        assert_eq!(TaskType::from_str("query"), Some(TaskType::SearchQuery));
        assert_eq!(TaskType::from_str("document"), Some(TaskType::SearchDocument));
        assert_eq!(TaskType::from_str("clustering"), Some(TaskType::Clustering));
        assert_eq!(TaskType::from_str("unknown"), None);
    }
}
