"""
Embedding Manager Module

Handles creation of embeddings for text chunks using sentence transformers.
"""

import asyncio
import numpy as np
from typing import List
from sentence_transformers import SentenceTransformer
from config.config import EMBEDDING_MODEL, BATCH_SIZE


class EmbeddingManager:
    """Handles embedding creation for text chunks."""
    
    def __init__(self):
        """Initialize the embedding manager."""
        self.embedding_model = None
        self._init_embedding_model()
    
    def _init_embedding_model(self):
        """Initialize the embedding model."""
        print(f"🔄 Loading embedding model: {EMBEDDING_MODEL}")
        self.embedding_model = SentenceTransformer(EMBEDDING_MODEL)
        print(f"✅ Embedding model loaded successfully")
    
    async def create_embeddings(self, chunks: List[str]) -> np.ndarray:
        """
        Create embeddings for text chunks.
        
        Args:
            chunks: List of text chunks to embed
            
        Returns:
            np.ndarray: Array of embeddings with shape (num_chunks, embedding_dim)
        """
        print(f"🧠 Creating embeddings for {len(chunks)} chunks")
        
        if not chunks:
            raise ValueError("No chunks provided for embedding creation")
        
        def create_embeddings_sync():
            """Synchronous embedding creation to run in thread pool."""
            embeddings = self.embedding_model.encode(
                chunks,
                batch_size=BATCH_SIZE,
                show_progress_bar=True,
                normalize_embeddings=True
            )
            return np.array(embeddings).astype("float32")
        
        # Run in thread pool to avoid blocking the event loop
        loop = asyncio.get_event_loop()
        embeddings = await loop.run_in_executor(None, create_embeddings_sync)
        
        print(f"✅ Created embeddings with shape: {embeddings.shape}")
        return embeddings
    
    def get_embedding_dimension(self) -> int:
        """
        Get the dimension of embeddings produced by the model.
        
        Returns:
            int: Embedding dimension
        """
        if self.embedding_model is None:
            raise RuntimeError("Embedding model not initialized")
        
        # Get dimension from model
        return self.embedding_model.get_sentence_embedding_dimension()
    
    def validate_embeddings(self, embeddings: np.ndarray, expected_count: int) -> bool:
        """
        Validate that embeddings have the expected shape and properties.
        
        Args:
            embeddings: The embeddings array to validate
            expected_count: Expected number of embeddings
            
        Returns:
            bool: True if embeddings are valid, False otherwise
        """
        if embeddings is None:
            return False
        
        if embeddings.shape[0] != expected_count:
            print(f"❌ Embedding count mismatch: expected {expected_count}, got {embeddings.shape[0]}")
            return False
        
        if embeddings.dtype != np.float32:
            print(f"❌ Embedding dtype mismatch: expected float32, got {embeddings.dtype}")
            return False
        
        # Check for NaN or infinite values
        if np.any(np.isnan(embeddings)) or np.any(np.isinf(embeddings)):
            print("❌ Embeddings contain NaN or infinite values")
            return False
        
        print(f"✅ Embeddings validation passed: {embeddings.shape}")
        return True
    
    def get_model_info(self) -> dict:
        """
        Get information about the embedding model.
        
        Returns:
            dict: Model information
        """
        if self.embedding_model is None:
            return {"model_name": EMBEDDING_MODEL, "status": "not_loaded"}
        
        return {
            "model_name": EMBEDDING_MODEL,
            "embedding_dimension": self.get_embedding_dimension(),
            "max_sequence_length": getattr(self.embedding_model, 'max_seq_length', 'unknown'),
            "status": "loaded"
        }