"""
Embedding Model Management Mixin

Provides FastEmbed and ColBERT embedding initialization and management
for the ModuleWrapper system.
"""

import logging
import shutil
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional

from adapters.module_wrapper.types import (
    COLBERT_DIM,
    MINILM_DIM,
    EmbeddingConfig,
    EmbeddingDimension,
    EmbeddingVector,
    MultiVector,
)

logger = logging.getLogger(__name__)


# =============================================================================
# LAZY-LOADED IMPORTS
# =============================================================================

_fastembed = None
_colbert_embed = None


def _get_fastembed():
    """Lazy load FastEmbed when first needed."""
    global _fastembed
    if _fastembed is None:
        logger.info("Loading FastEmbed (first use)...")
        from fastembed import TextEmbedding

        _fastembed = TextEmbedding
        logger.info("FastEmbed loaded")
    return _fastembed


def _get_colbert_embed():
    """Lazy load ColBERT LateInteractionTextEmbedding when first needed."""
    global _colbert_embed
    if _colbert_embed is None:
        logger.info("Loading ColBERT LateInteractionTextEmbedding (first use)...")
        from fastembed import LateInteractionTextEmbedding

        _colbert_embed = LateInteractionTextEmbedding
        logger.info("ColBERT LateInteractionTextEmbedding loaded")
    return _colbert_embed


# =============================================================================
# EMBEDDING MIXIN
# =============================================================================


class EmbeddingMixin:
    """
    Mixin providing embedding model initialization and management.

    Expects the following attributes on self:
    - embedding_model_name: str
    - enable_colbert: bool
    - colbert_model_name: str
    """

    # --- Mixin dependency contract ---
    _MIXIN_PROVIDES = frozenset(
        {
            "embedder",
            "embedding_dim",
            "colbert_embedder",
            "colbert_embedding_dim",
            "_colbert_initialized",
            "_initialize_embedder",
            "_initialize_colbert_embedder",
            "embedding_metadata",
        }
    )
    _MIXIN_REQUIRES = frozenset(
        {
            "embedding_model_name",
            "enable_colbert",
            "colbert_model_name",
        }
    )
    _MIXIN_INIT_ORDER = 20

    # Attributes to be set by __init__ in the composed class
    embedder: Any = None
    embedding_dim: Optional[EmbeddingDimension] = None
    colbert_embedder: Any = None
    colbert_embedding_dim: EmbeddingDimension = COLBERT_DIM
    _colbert_initialized: bool = False

    def _clear_fastembed_cache(self, model_name: Optional[str] = None) -> bool:
        """
        Clear corrupted FastEmbed cache to allow re-download.

        Args:
            model_name: Optional specific model name to clear (clears all if None)

        Returns:
            bool: True if cache was cleared successfully
        """
        import os

        cleared = False
        cache_locations = [
            Path(tempfile.gettempdir()) / "fastembed_cache",
            Path.home() / ".cache" / "fastembed",
        ]

        # Also check macOS-specific temp locations
        try:
            # Get the actual temp directory which may be in /var/folders on macOS
            actual_temp = Path(os.path.realpath(tempfile.gettempdir()))
            if actual_temp not in cache_locations:
                cache_locations.append(actual_temp / "fastembed_cache")
        except Exception:
            pass

        for cache_dir in cache_locations:
            if cache_dir.exists():
                try:
                    if model_name:
                        # Clear specific model cache
                        model_short = model_name.split("/")[-1]
                        for subdir in cache_dir.iterdir():
                            if model_short in subdir.name or "MiniLM" in subdir.name:
                                logger.info(f"Clearing corrupted cache: {subdir}")
                                shutil.rmtree(subdir)
                                cleared = True
                    else:
                        # Clear entire cache
                        logger.info(f"Clearing entire FastEmbed cache: {cache_dir}")
                        shutil.rmtree(cache_dir)
                        cleared = True
                except Exception as clear_error:
                    logger.warning(f"Could not clear cache {cache_dir}: {clear_error}")

        return cleared

    def _initialize_embedder(self):
        """Initialize the embedding model with retry logic for corrupted cache."""
        global _fastembed

        max_retries = 2
        last_error = None

        for attempt in range(max_retries + 1):
            try:
                TextEmbedding = _get_fastembed()
                self.embedder = TextEmbedding(model_name=self.embedding_model_name)

                # Try to get embedding dimension dynamically
                try:
                    # Generate a test embedding to get the dimension
                    test_embedding = list(self.embedder.embed(["test"]))[0]
                    self.embedding_dim = (
                        len(test_embedding)
                        if hasattr(test_embedding, "__len__")
                        else 384
                    )
                except Exception:
                    # Fallback to known dimensions for common models
                    model_dims = {
                        "sentence-transformers/all-MiniLM-L6-v2": 384,
                        "sentence-transformers/all-mpnet-base-v2": 768,
                    }
                    self.embedding_dim = model_dims.get(self.embedding_model_name, 384)

                logger.info(
                    f"Embedding model loaded: {self.embedding_model_name} (dim: {self.embedding_dim})"
                )
                return  # Success!

            except Exception as e:
                last_error = e
                error_str = str(e).lower()

                # Check if this is a cache/file corruption error that we can recover from
                is_recoverable = any(
                    keyword in error_str
                    for keyword in [
                        "no_suchfile",
                        "file doesn't exist",
                        "corrupted",
                        "model.onnx",
                        "failed to load",
                        "invalid model",
                    ]
                )

                if is_recoverable and attempt < max_retries:
                    logger.warning(
                        f"Embedding model load failed (attempt {attempt + 1}/{max_retries + 1}): {e}"
                    )
                    logger.info(
                        "Attempting to clear corrupted cache and re-download model..."
                    )

                    # Clear the cache and retry
                    if self._clear_fastembed_cache(self.embedding_model_name):
                        logger.info("Cache cleared, retrying model download...")
                        # Reset the global fastembed reference to force reload
                        _fastembed = None
                        continue
                    else:
                        logger.warning("Could not clear cache, retrying anyway...")
                        continue
                else:
                    # Non-recoverable error or out of retries
                    break

        # All retries exhausted
        logger.error(
            f"Failed to initialize embedding model after {max_retries + 1} attempts: {last_error}"
        )
        raise last_error

    def _initialize_colbert_embedder(self):
        """Initialize the ColBERT late interaction embedding model."""
        if not self.enable_colbert:
            return

        try:
            LateInteractionTextEmbedding = _get_colbert_embed()
            logger.info(f"Initializing ColBERT model: {self.colbert_model_name}")
            self.colbert_embedder = LateInteractionTextEmbedding(
                model_name=self.colbert_model_name
            )

            # Get embedding dimension from model
            colbert_dims = {
                "colbert-ir/colbertv2.0": 128,
                "answerdotai/answerai-colbert-small-v1": 96,
            }
            self.colbert_embedding_dim = colbert_dims.get(self.colbert_model_name, 128)

            logger.info(
                f"ColBERT model loaded: {self.colbert_model_name} (dim: {self.colbert_embedding_dim})"
            )
            self._colbert_initialized = True

        except Exception as e:
            logger.error(f"Failed to initialize ColBERT embedder: {e}")
            self.enable_colbert = False
            raise

    @property
    def embedding_metadata(self) -> Dict[str, Any]:
        """Get metadata about the embedding configuration."""
        return {
            "embedding_model": getattr(self, "embedding_model_name", None),
            "embedding_dim": self.embedding_dim,
            "colbert_enabled": getattr(self, "enable_colbert", False),
            "colbert_model": (
                getattr(self, "colbert_model_name", None)
                if getattr(self, "enable_colbert", False)
                else None
            ),
            "colbert_dim": (
                self.colbert_embedding_dim
                if getattr(self, "_colbert_initialized", False)
                else None
            ),
        }


# Export for convenience
__all__ = [
    "EmbeddingMixin",
    "_get_fastembed",
    "_get_colbert_embed",
]
