import logging
from typing import Callable, List, Optional

_logger = logging.getLogger(__name__)

# Supported sparse embedding models with their configurations
SUPPORTED_SPARSE_MODELS = {
    # SPLADE (FastEmbed-supported) - Standard recommendation
    "splade_pp": {
        "model_name": "prithivida/Splade_PP_en_v1",
        "vocab_size": 30522,
        "requires_idf": False,
        "description": "SPLADE++ Model for English - Best for semantic + keyword hybrid search",
    },
}


def get_model_vocab_size(model_key: str) -> int:
    """
    Get the vocabulary size for a given model key.

    Args:
        model_key: Short model identifier (e.g., 'splade_pp')

    Returns:
        Vocabulary size (sparse_dim) for the model
    """
    if model_key in SUPPORTED_SPARSE_MODELS:
        return SUPPORTED_SPARSE_MODELS[model_key]["vocab_size"]

    # Default to SPLADE vocab size for unknown models
    _logger.warning(f"Unknown model key '{model_key}', defaulting to vocab_size=30522")
    return 30522


def _initialize_sparse_encoder_fastembed(
    model_name: str,
    batch_size: int = 256,
    cache_dir: Optional[str] = None,
    threads: Optional[int] = None,
) -> Callable:
    """
    Initialize a sparse encoder using FastEmbed (recommended for SPLADE models).

    Args:
        model_name: Model key from SUPPORTED_SPARSE_MODELS or full model path
        batch_size: Batch size for encoding
        cache_dir: Cache directory for model files
        threads: Number of threads for CPU inference

    Returns:
        Callable that computes sparse vectors from texts
    """
    try:
        from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
    except ImportError as e:
        raise ImportError(
            "FastEmbed is required for hybrid search but not installed.\n"
            "Install with: pip install fastembed\n"
            "For dense-only search, create vector store without sparse_dim/model_name."
        ) from e

    # Resolve model name from key or use directly
    if model_name in SUPPORTED_SPARSE_MODELS:
        resolved_model_name = SUPPORTED_SPARSE_MODELS[model_name]["model_name"]
        _logger.info(f"Using model alias '{model_name}' -> '{resolved_model_name}'")
    else:
        resolved_model_name = model_name
        _logger.info(f"Using custom model: '{resolved_model_name}'")

    # Initialize model with GPU support if available
    try:
        model = SparseTextEmbedding(
            resolved_model_name,
            cache_dir=cache_dir,
            threads=threads,
            providers=["CUDAExecutionProvider"],
        )
        _logger.info(f"✓ Initialized sparse encoder '{resolved_model_name}' on GPU")
    except Exception as gpu_error:
        _logger.debug(f"GPU initialization failed: {gpu_error}")
        try:
            model = SparseTextEmbedding(
                resolved_model_name, cache_dir=cache_dir, threads=threads
            )
            _logger.info(f"✓ Initialized sparse encoder '{resolved_model_name}' on CPU")
        except Exception as cpu_error:
            _logger.error(
                f"Failed to initialize model '{resolved_model_name}': {cpu_error}"
            )
            raise

    def compute_vectors(texts: List[str]) -> tuple:
        """
        Compute sparse vectors (indices, values) for a list of texts.

        Args:
            texts: List of text strings to encode

        Returns:
            Tuple of (indices, values) where each is a list of lists
        """
        try:
            embeddings = model.embed(texts, batch_size=batch_size)
            indices = []
            values = []

            for embedding in embeddings:
                indices.append(embedding.indices.tolist())
                values.append(embedding.values.tolist())

            return indices, values
        except Exception as e:
            _logger.error(f"Error computing sparse vectors: {e}")
            raise

    return compute_vectors


def _initialize_sparse_encoder_transformers(model_name: str) -> Callable:
    """
    Initialize a sparse encoder using Transformers library.
    This is an alternative to FastEmbed for SPLADE models.

    Args:
        model_name: Model key from SUPPORTED_SPARSE_MODELS or full model path

    Returns:
        Callable that computes sparse vectors from texts
    """
    try:
        import torch
        from transformers import AutoModelForMaskedLM, AutoTokenizer
    except ImportError as e:
        raise ImportError(
            "Transformers is required for this encoder but not installed.\n"
            'Install with: pip install "transformers[torch]"'
        ) from e

    # Resolve model name from key or use directly
    if model_name in SUPPORTED_SPARSE_MODELS:
        resolved_model_name = SUPPORTED_SPARSE_MODELS[model_name]["model_name"]
        _logger.info(f"Using model alias '{model_name}' -> '{resolved_model_name}'")
    else:
        resolved_model_name = model_name
        _logger.info(f"Using custom model: '{resolved_model_name}'")

    tokenizer = AutoTokenizer.from_pretrained(resolved_model_name)
    model = AutoModelForMaskedLM.from_pretrained(resolved_model_name)

    if torch.cuda.is_available():
        model = model.to("cuda")
        _logger.info(f"✓ Initialized sparse encoder '{resolved_model_name}' on GPU")
    else:
        _logger.info(f"✓ Initialized sparse encoder '{resolved_model_name}' on CPU")

    def compute_vectors(texts: List[str]) -> tuple:
        """
        Compute sparse vectors from logits using SPLADE methodology.

        Args:
            texts: List of text strings to encode

        Returns:
            Tuple of (indices, values) where each is a list of lists
        """
        tokens = tokenizer(
            texts, truncation=True, padding=True, max_length=512, return_tensors="pt"
        )

        if torch.cuda.is_available():
            tokens = tokens.to("cuda")

        with torch.no_grad():
            output = model(**tokens)
            logits, attention_mask = output.logits, tokens.attention_mask

            # SPLADE: relu(log(1 + relu(logits))) weighted by attention
            relu_log = torch.log(1 + torch.relu(logits))
            weighted_log = relu_log * attention_mask.unsqueeze(-1)

            # Max pooling across sequence dimension
            tvecs, _ = torch.max(weighted_log, dim=1)

        indices = []
        values = []
        for batch in tvecs:
            # Get non-zero indices and their values
            nz_indices = batch.nonzero(as_tuple=True)[0].tolist()
            indices.append(nz_indices)
            values.append(batch[nz_indices].tolist())

        return indices, values

    return compute_vectors


def get_sparse_encoder(
    model_name: Optional[str] = None,
    use_fastembed: bool = True,
    batch_size: int = 256,
    cache_dir: Optional[str] = None,
    threads: Optional[int] = None,
) -> Optional[Callable]:
    """
    Get a sparse encoder function for the specified model.

    Args:
        model_name: Model key (e.g., 'splade_pp') or full model path
        use_fastembed: Whether to use FastEmbed (recommended) vs Transformers
        batch_size: Batch size for encoding (FastEmbed only)
        cache_dir: Cache directory for model files (FastEmbed only)
        threads: Number of threads for CPU inference (FastEmbed only)

    Returns:
        Callable that takes List[str] and returns (indices, values) tuple,
        or None if model_name is None

    Example:
        >>> encoder = get_sparse_encoder('splade_pp')
        >>> indices, values = encoder(['Hello world', 'Test document'])
    """
    if model_name is None:
        return None

    if use_fastembed:
        return _initialize_sparse_encoder_fastembed(
            model_name=model_name,
            batch_size=batch_size,
            cache_dir=cache_dir,
            threads=threads,
        )
    else:
        return _initialize_sparse_encoder_transformers(model_name=model_name)


def list_supported_models() -> dict:
    """
    List all supported sparse models with their configurations.

    Returns:
        Dictionary of model configurations
    """
    return SUPPORTED_SPARSE_MODELS.copy()
