"""
Embedding Generation Module

This module provides comprehensive embedding generation capabilities for text,
images, audio, and multi-modal content, with support for multiple embedding
models and optimization strategies.

Key Features:
    - Text embedding generation (multiple models: sentence-transformers, OpenAI, BGE, etc.)
    - Image embedding generation
    - Audio embedding generation
    - Multi-modal embedding support
    - Batch processing for efficiency
    - Embedding optimization and compression
    - Similarity comparison utilities

Example Usage:
    >>> from semantica.embeddings import EmbeddingGenerator
    >>> generator = EmbeddingGenerator()
    >>> embeddings = generator.generate_embeddings("Hello world", data_type="text")
    >>> similarity = generator.compare_embeddings(emb1, emb2, method="cosine")

Author: Semantica Contributors
License: MIT
"""

from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import numpy as np

from ..utils.exceptions import ProcessingError
from ..utils.logging import get_logger
from .text_embedder import TextEmbedder
from .image_embedder import ImageEmbedder
from .audio_embedder import AudioEmbedder
from .multimodal_embedder import MultimodalEmbedder
from .embedding_optimizer import EmbeddingOptimizer


class EmbeddingGenerator:
    """
    Main embedding generation handler.
    
    This class provides a unified interface for generating embeddings from
    various data types (text, images, audio) using multiple embedding models.
    It handles batch processing, optimization, and similarity calculations.
    
    Supported Models:
        - sentence-transformers: High-quality sentence embeddings
        - openai: OpenAI text-embedding models
        - bge: BAAI General Embedding models
        - clip: CLIP for image-text embeddings
    
    Example Usage:
        >>> generator = EmbeddingGenerator()
        >>> # Single text embedding
        >>> emb = generator.generate_embeddings("Hello world", data_type="text")
        >>> # Batch processing
        >>> embs = generator.process_batch(["text1", "text2", "text3"])
    """
    
    def __init__(self, config: Optional[Dict[str, Any]] = None, **kwargs):
        """
        Initialize embedding generator.
        
        Sets up embedders for different data types and the embedding optimizer.
        Configuration can be provided via config dict or keyword arguments.
        
        Args:
            config: Configuration dictionary with keys:
                - text: Text embedder configuration
                - image: Image embedder configuration
                - audio: Audio embedder configuration
                - multimodal: Multi-modal embedder configuration
                - optimizer: Embedding optimizer configuration
            **kwargs: Additional configuration (merged into config)
        """
        self.logger = get_logger("embedding_generator")
        
        # Merge configuration
        self.config = config or {}
        self.config.update(kwargs)
        
        # Initialize embedders for different data types
        # These are lazy-loaded and only initialized when needed
        text_config = self.config.get("text", {})
        image_config = self.config.get("image", {})
        audio_config = self.config.get("audio", {})
        multimodal_config = self.config.get("multimodal", {})
        optimizer_config = self.config.get("optimizer", {})
        
        self.text_embedder = TextEmbedder(**text_config)
        self.image_embedder = ImageEmbedder(**image_config)
        self.audio_embedder = AudioEmbedder(**audio_config)
        self.multimodal_embedder = MultimodalEmbedder(**multimodal_config)
        self.embedding_optimizer = EmbeddingOptimizer(**optimizer_config)
        
        # List of supported embedding models
        self.supported_models = ["sentence-transformers", "openai", "bge", "clip"]
        
        # Initialize progress tracker
        from ..utils.progress_tracker import get_progress_tracker
        self.progress_tracker = get_progress_tracker()
        
        self.logger.info("Embedding generator initialized")
    
    def generate_embeddings(
        self, 
        data: Union[str, Path, List[Union[str, Path]]], 
        data_type: Optional[str] = None, 
        **options
    ) -> np.ndarray:
        """
        Generate embeddings for input data.
        
        This method automatically detects the data type if not specified and
        routes to the appropriate embedder. Supports both single items and batches.
        
        Args:
            data: Input data to embed:
                - str: Text string or file path
                - Path: File path object
                - List: Batch of texts or file paths
            data_type: Explicit data type ("text", "image", "audio")
                      If None, auto-detects from input
            **options: Additional generation options passed to embedder
            
        Returns:
            np.ndarray: Generated embeddings
                - For single input: 1D array
                - For batch input: 2D array (batch_size, embedding_dim)
                
        Raises:
            ProcessingError: If data type is unsupported or embedding fails
            
        Example:
            >>> # Single text
            >>> emb = generator.generate_embeddings("Hello world")
            >>> # Batch of texts
            >>> embs = generator.generate_embeddings(["text1", "text2"])
            >>> # Image file
            >>> emb = generator.generate_embeddings("image.jpg", data_type="image")
        """
        # Auto-detect data type if not provided
        if data_type is None:
            data_type = self._detect_data_type(data)
            self.logger.debug(f"Auto-detected data type: {data_type}")
        
        # Route to appropriate embedder based on data type
        if data_type == "text":
            if isinstance(data, str):
                return self.text_embedder.embed_text(data, **options)
            elif isinstance(data, list):
                return self.text_embedder.embed_batch(data, **options)
        elif data_type == "image":
            if isinstance(data, (str, Path)):
                return self.image_embedder.embed_image(data, **options)
            elif isinstance(data, list):
                return self.image_embedder.embed_batch(data, **options)
        elif data_type == "audio":
            if isinstance(data, (str, Path)):
                return self.audio_embedder.embed_audio(data, **options)
            elif isinstance(data, list):
                return self.audio_embedder.embed_batch(data, **options)
        else:
            error_msg = (
                f"Unsupported data type: {data_type}. "
                f"Supported types: text, image, audio"
            )
            raise ProcessingError(error_msg)
    
    def _detect_data_type(self, data: Union[str, Path, List]) -> str:
        """
        Detect data type from input automatically.
        
        This method analyzes the input to determine whether it's text, image,
        or audio data. For file paths, it uses file extensions. For strings,
        it defaults to text.
        
        Args:
            data: Input data to analyze
            
        Returns:
            str: Detected data type ("text", "image", or "audio")
        """
        # Handle list inputs - check first item
        if isinstance(data, list):
            if not data:
                # Empty list defaults to text
                return "text"
            # Recursively check first item
            return self._detect_data_type(data[0])
        
        # Check if input is a file path
        file_path = None
        if isinstance(data, Path):
            file_path = data
        elif isinstance(data, str):
            # Check if string is a valid file path
            potential_path = Path(data)
            if potential_path.exists():
                file_path = potential_path
        
        # If it's a file path, detect type by extension
        if file_path:
            suffix = file_path.suffix.lower()
            
            # Common image file extensions
            image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp', '.svg']
            # Common audio file extensions
            audio_extensions = ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.m4a', '.wma']
            
            if suffix in image_extensions:
                return "image"
            elif suffix in audio_extensions:
                return "audio"
        
        # Default to text for strings or unknown file types
        return "text"
    
    def optimize_embeddings(self, embeddings: np.ndarray, **options) -> np.ndarray:
        """
        Optimize embedding quality and performance.
        
        Args:
            embeddings: Input embeddings
            **options: Optimization options
            
        Returns:
            np.ndarray: Optimized embeddings
        """
        return self.embedding_optimizer.compress(embeddings, **options)
    
    def compare_embeddings(self, embedding1: np.ndarray, embedding2: np.ndarray, **options) -> float:
        """
        Compare embeddings for similarity.
        
        Args:
            embedding1: First embedding
            embedding2: Second embedding
            **options: Comparison options:
                - method: Similarity method ("cosine", "euclidean")
                
        Returns:
            float: Similarity score (0-1)
        """
        method = options.get("method", "cosine")
        
        if method == "cosine":
            return self._cosine_similarity(embedding1, embedding2)
        elif method == "euclidean":
            return self._euclidean_similarity(embedding1, embedding2)
        else:
            return self._cosine_similarity(embedding1, embedding2)
    
    def _cosine_similarity(self, emb1: np.ndarray, emb2: np.ndarray) -> float:
        """Compute cosine similarity."""
        dot_product = np.dot(emb1, emb2)
        norm1 = np.linalg.norm(emb1)
        norm2 = np.linalg.norm(emb2)
        
        if norm1 == 0 or norm2 == 0:
            return 0.0
        
        return float(dot_product / (norm1 * norm2))
    
    def _euclidean_similarity(self, emb1: np.ndarray, emb2: np.ndarray) -> float:
        """Compute Euclidean similarity (converted to 0-1 scale)."""
        distance = np.linalg.norm(emb1 - emb2)
        # Normalize to 0-1 (simple approach)
        max_distance = np.linalg.norm(emb1) + np.linalg.norm(emb2)
        similarity = 1.0 - (distance / max_distance) if max_distance > 0 else 0.0
        return max(0.0, min(1.0, similarity))
    
    def process_batch(self, data_items: List[Union[str, Path]], **options) -> Dict[str, Any]:
        """
        Process multiple data items for embedding generation.
        
        Args:
            data_items: List of data items
            **options: Processing options
            
        Returns:
            dict: Batch processing results
        """
        # Track batch processing
        tracking_id = self.progress_tracker.start_tracking(
            module="embeddings",
            submodule="EmbeddingGenerator",
            message=f"Batch of {len(data_items)} items"
        )
        
        try:
            results = {
                "embeddings": [],
                "successful": [],
                "failed": []
            }
            
            for idx, item in enumerate(data_items, 1):
                try:
                    self.progress_tracker.update_tracking(tracking_id, 
                                                         message=f"Processing item {idx}/{len(data_items)}")
                    embedding = self.generate_embeddings(item, **options)
                    results["embeddings"].append(embedding)
                    results["successful"].append(str(item))
                except Exception as e:
                    results["failed"].append({
                        "item": str(item),
                        "error": str(e)
                    })
            
            results["total"] = len(data_items)
            results["success_count"] = len(results["successful"])
            results["failure_count"] = len(results["failed"])
            
            self.progress_tracker.stop_tracking(tracking_id, status="completed",
                                               message=f"Processed {results['success_count']}/{len(data_items)} items successfully")
            return results
            
        except Exception as e:
            self.progress_tracker.stop_tracking(tracking_id, status="failed", message=str(e))
            raise


# Re-export classes from other modules for convenience
from .text_embedder import TextEmbedder as TextEmbedderImpl
from .multimodal_embedder import MultimodalEmbedder as MultimodalEmbedderImpl
from .embedding_optimizer import EmbeddingOptimizer as EmbeddingOptimizerImpl

# Make classes available with same names
TextEmbedder = TextEmbedderImpl
MultiModalEmbedder = MultimodalEmbedderImpl
EmbeddingOptimizer = EmbeddingOptimizerImpl
