//! Embedding model implementations.

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

use candle_core::{Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::bert::{BertModel, Config as BertConfig};
use hf_hub::{api::sync::Api, Repo, RepoType};
use parking_lot::Mutex;
use tokenizers::Tokenizer;

use crate::config::EmbeddingModel as EmbeddingModelConfig;
use crate::embeddings::{Embedder, PoolingStrategy};
use crate::error::{Error, Result};

/// Local embedding model using Candle.
pub struct LocalEmbedder {
    model: Arc<Mutex<BertModel>>,
    tokenizer: Arc<Tokenizer>,
    device: Device,
    dimension: usize,
    pooling: PoolingStrategy,
    max_length: usize,
}

impl LocalEmbedder {
    /// Create a new local embedder with the specified model.
    pub fn new(config: &EmbeddingModelConfig) -> Result<Self> {
        let device = Device::Cpu; // Use CPU for maximum compatibility

        let (model, tokenizer, dimension) = match config {
            EmbeddingModelConfig::AllMiniLmL6V2 => {
                Self::load_model("sentence-transformers/all-MiniLM-L6-v2", &device)?
            }
            EmbeddingModelConfig::AllMiniLmL12V2 => {
                Self::load_model("sentence-transformers/all-MiniLM-L12-v2", &device)?
            }
            EmbeddingModelConfig::CodeBert => {
                Self::load_model("microsoft/codebert-base", &device)?
            }
            EmbeddingModelConfig::Custom {
                model_id,
                dimension,
                ..
            } => {
                let (model, tokenizer, _) = Self::load_model(model_id, &device)?;
                (model, tokenizer, *dimension)
            }
        };

        Ok(Self {
            model: Arc::new(Mutex::new(model)),
            tokenizer: Arc::new(tokenizer),
            device,
            dimension,
            pooling: PoolingStrategy::Mean,
            max_length: 512,
        })
    }

    /// Load model from HuggingFace Hub.
    fn load_model(model_id: &str, device: &Device) -> Result<(BertModel, Tokenizer, usize)> {
        tracing::info!("Loading embedding model: {}", model_id);

        let api = Api::new().map_err(|e| Error::model_load(e.to_string()))?;
        let repo = api.repo(Repo::new(model_id.to_string(), RepoType::Model));

        // Download config
        let config_path = repo
            .get("config.json")
            .map_err(|e| Error::model_load(format!("Failed to download config: {}", e)))?;

        let config_str = std::fs::read_to_string(&config_path)
            .map_err(|e| Error::model_load(format!("Failed to read config: {}", e)))?;

        let config: BertConfig = serde_json::from_str(&config_str)
            .map_err(|e| Error::model_load(format!("Failed to parse config: {}", e)))?;

        let dimension = config.hidden_size;

        // Download tokenizer
        let tokenizer_path = repo
            .get("tokenizer.json")
            .map_err(|e| Error::model_load(format!("Failed to download tokenizer: {}", e)))?;

        let tokenizer = Tokenizer::from_file(&tokenizer_path)
            .map_err(|e| Error::model_load(format!("Failed to load tokenizer: {}", e)))?;

        // Download model weights
        let weights_path = repo
            .get("model.safetensors")
            .or_else(|_| repo.get("pytorch_model.bin"))
            .map_err(|e| Error::model_load(format!("Failed to download weights: {}", e)))?;

        // Load weights
        let vb = if weights_path.extension().map_or(false, |ext| ext == "safetensors") {
            unsafe {
                VarBuilder::from_mmaped_safetensors(&[weights_path], candle_core::DType::F32, device)
                    .map_err(|e| Error::model_load(format!("Failed to load safetensors: {}", e)))?
            }
        } else {
            return Err(Error::model_load("Pytorch weights not supported, need safetensors"));
        };

        // Create model
        let model = BertModel::load(vb, &config)
            .map_err(|e| Error::model_load(format!("Failed to create model: {}", e)))?;

        tracing::info!("Model loaded successfully, dimension: {}", dimension);

        Ok((model, tokenizer, dimension))
    }

    /// Create a simple embedder that doesn't require downloading models.
    /// Uses a random projection for testing purposes.
    pub fn simple(dimension: usize) -> Self {
        Self {
            model: Arc::new(Mutex::new(Self::create_dummy_model())),
            tokenizer: Arc::new(Self::create_dummy_tokenizer()),
            device: Device::Cpu,
            dimension,
            pooling: PoolingStrategy::Mean,
            max_length: 512,
        }
    }

    fn create_dummy_model() -> BertModel {
        // This would panic, so we need a different approach for testing
        unimplemented!("Use LocalEmbedder::new() with a real model")
    }

    fn create_dummy_tokenizer() -> Tokenizer {
        // This would also need a real tokenizer
        unimplemented!("Use LocalEmbedder::new() with a real model")
    }

    /// Tokenize text.
    fn tokenize(&self, text: &str) -> Result<tokenizers::Encoding> {
        let encoding = self
            .tokenizer
            .encode(text, true)
            .map_err(|e| Error::embedding(format!("Tokenization failed: {}", e)))?;

        Ok(encoding)
    }

    /// Apply pooling to get a single vector from token embeddings.
    fn pool(&self, embeddings: &Tensor, attention_mask: &Tensor) -> Result<Tensor> {
        match self.pooling {
            PoolingStrategy::Mean => {
                // Mean pooling: average over non-masked tokens
                let mask = attention_mask
                    .unsqueeze(2)?
                    .broadcast_as(embeddings.shape())?
                    .to_dtype(embeddings.dtype())?;

                let masked = (embeddings * &mask)?;
                let sum = masked.sum(1)?;
                let count = mask.sum(1)?;
                let mean = (sum / count)?;
                Ok(mean)
            }
            PoolingStrategy::Cls => {
                // CLS token (first token)
                Ok(embeddings.narrow(1, 0, 1)?.squeeze(1)?)
            }
            PoolingStrategy::Max => {
                // Max pooling over sequence
                Ok(embeddings.max(1)?)
            }
        }
    }

    /// Normalize a vector to unit length.
    fn normalize(&self, tensor: &Tensor) -> Result<Tensor> {
        let norm = tensor.sqr()?.sum_keepdim(1)?.sqrt()?;
        Ok((tensor / norm)?)
    }
}

impl Embedder for LocalEmbedder {
    fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let batch = self.embed_batch(&[text])?;
        Ok(batch.into_iter().next().unwrap())
    }

    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }

        // Tokenize all texts
        let encodings: Vec<_> = texts
            .iter()
            .map(|t| self.tokenize(t))
            .collect::<Result<Vec<_>>>()?;

        // Determine max length for padding
        let max_len = encodings
            .iter()
            .map(|e| e.get_ids().len())
            .max()
            .unwrap_or(0)
            .min(self.max_length);

        // Prepare input tensors
        let mut input_ids = Vec::new();
        let mut attention_masks = Vec::new();

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

            let len = ids.len().min(max_len);

            // Pad or truncate
            let mut padded_ids = vec![0u32; max_len];
            let mut padded_mask = vec![0u32; max_len];

            padded_ids[..len].copy_from_slice(&ids[..len]);
            padded_mask[..len].copy_from_slice(&mask[..len]);

            input_ids.extend(padded_ids);
            attention_masks.extend(padded_mask);
        }

        let batch_size = texts.len();

        let input_ids = Tensor::from_vec(input_ids, (batch_size, max_len), &self.device)?;
        let attention_mask = Tensor::from_vec(attention_masks, (batch_size, max_len), &self.device)?;
        let token_type_ids = Tensor::zeros((batch_size, max_len), candle_core::DType::U32, &self.device)?;

        // Run model
        let model = self.model.lock();
        let embeddings = model.forward(&input_ids, &token_type_ids, Some(&attention_mask))?;

        // Pool and normalize
        let pooled = self.pool(&embeddings, &attention_mask)?;
        let normalized = self.normalize(&pooled)?;

        // Convert to Vec<Vec<f32>>
        let flat: Vec<f32> = normalized.flatten_all()?.to_vec1()?;
        let results: Vec<Vec<f32>> = flat
            .chunks(self.dimension)
            .map(|chunk| chunk.to_vec())
            .collect();

        Ok(results)
    }

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

    fn name(&self) -> &'static str {
        "local"
    }
}

/// Embedding model configuration and loading utilities.
pub struct EmbeddingModel;

impl EmbeddingModel {
    /// Get default model path for caching.
    pub fn default_cache_dir() -> PathBuf {
        dirs::cache_dir()
            .unwrap_or_else(|| PathBuf::from("."))
            .join("aurora-semantic")
            .join("models")
    }

    /// Check if a model is cached locally.
    pub fn is_cached(model_id: &str) -> bool {
        let cache_dir = Self::default_cache_dir();
        let model_dir = cache_dir.join(model_id.replace('/', "--"));
        model_dir.exists()
    }
}

/// Simple embedder for testing that uses random projections.
pub struct RandomEmbedder {
    dimension: usize,
}

impl RandomEmbedder {
    pub fn new(dimension: usize) -> Self {
        Self { dimension }
    }

    fn hash_to_embedding(&self, text: &str) -> Vec<f32> {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut result = vec![0.0f32; self.dimension];
        let hash = {
            let mut hasher = DefaultHasher::new();
            text.hash(&mut hasher);
            hasher.finish()
        };

        // Generate pseudo-random values from hash
        let mut seed = hash;
        for i in 0..self.dimension {
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            result[i] = ((seed >> 32) as f32 / u32::MAX as f32) * 2.0 - 1.0;
        }

        // Normalize
        let norm: f32 = result.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            for x in &mut result {
                *x /= norm;
            }
        }

        result
    }
}

impl Embedder for RandomEmbedder {
    fn embed(&self, text: &str) -> Result<Vec<f32>> {
        Ok(self.hash_to_embedding(text))
    }

    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        Ok(texts.iter().map(|t| self.hash_to_embedding(t)).collect())
    }

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

    fn name(&self) -> &'static str {
        "random"
    }
}

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

    #[test]
    fn test_random_embedder() {
        let embedder = RandomEmbedder::new(384);

        let embedding = embedder.embed("test code").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() < 0.01);

        // Same text should produce same embedding
        let embedding2 = embedder.embed("test code").unwrap();
        assert_eq!(embedding, embedding2);

        // Different text should produce different embedding
        let embedding3 = embedder.embed("other code").unwrap();
        assert_ne!(embedding, embedding3);
    }

    #[test]
    fn test_batch_embedding() {
        let embedder = RandomEmbedder::new(128);

        let texts = vec!["hello", "world", "test"];
        let embeddings = embedder.embed_batch(&texts).unwrap();

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