from functools import lru_cache
from typing import Any

from . import constants as cs
from . import exceptions as ex
from .config import settings
from .utils.dependencies import has_fastembed, has_torch, has_transformers


def _resolve_provider(provider: str | None = None) -> str:
    resolved = provider or settings.EMBEDDING_PROVIDER
    return resolved.lower().strip()


def _resolve_model_name(provider: str, model_name: str | None = None) -> str:
    if model_name is not None and model_name.strip():
        return model_name.strip()
    if settings.EMBEDDING_MODEL.strip():
        return settings.EMBEDDING_MODEL.strip()
    if provider == cs.EMBEDDING_PROVIDER_FASTEMBED:
        return cs.FASTEMBED_MODEL
    return cs.UNIXCODER_MODEL


def _ensure_backend_dependencies(provider: str) -> None:
    if provider == cs.EMBEDDING_PROVIDER_UNIXCODER:
        if has_torch() and has_transformers():
            return
        raise RuntimeError(ex.SEMANTIC_EXTRA)
    if provider == cs.EMBEDDING_PROVIDER_FASTEMBED:
        if has_fastembed():
            return
        raise RuntimeError(ex.SEMANTIC_EXTRA)
    raise ValueError(f"Unsupported embedding provider: {provider}")


def _vector_to_list(vector: Any) -> list[float]:
    if hasattr(vector, "tolist"):
        result = vector.tolist()
        return [float(value) for value in result]
    return [float(value) for value in vector]


@lru_cache(maxsize=4)
def get_model(provider: str | None = None, model_name: str | None = None) -> Any:
    resolved_provider = _resolve_provider(provider)
    resolved_model = _resolve_model_name(resolved_provider, model_name)
    _ensure_backend_dependencies(resolved_provider)

    if resolved_provider == cs.EMBEDDING_PROVIDER_UNIXCODER:
        import torch

        from .unixcoder import UniXcoder

        model = UniXcoder(resolved_model)
        model.eval()
        if torch.cuda.is_available():
            model = model.cuda()
        return model

    from fastembed import TextEmbedding

    return TextEmbedding(model_name=resolved_model)


def embed_code(code: str, max_length: int | None = None) -> list[float]:
    provider = _resolve_provider()
    if provider == cs.EMBEDDING_PROVIDER_FASTEMBED:
        model = get_model(provider)
        embedding = next(iter(model.embed([code])))
        return _vector_to_list(embedding)

    if max_length is None:
        max_length = settings.EMBEDDING_MAX_LENGTH
    model = get_model(provider)
    import torch

    device = next(model.parameters()).device
    tokens = model.tokenize([code], max_length=max_length)
    tokens_tensor = torch.tensor(tokens).to(device)
    with torch.no_grad():
        _, sentence_embeddings = model(tokens_tensor)
    return _vector_to_list(sentence_embeddings[0].cpu().numpy())


def embed_code_batch(
    code_batch: list[str], max_length: int | None = None
) -> list[list[float]]:
    if not code_batch:
        return []

    provider = _resolve_provider()
    if provider == cs.EMBEDDING_PROVIDER_FASTEMBED:
        model = get_model(provider)
        return [_vector_to_list(embedding) for embedding in model.embed(code_batch)]

    if max_length is None:
        max_length = settings.EMBEDDING_MAX_LENGTH
    model = get_model(provider)
    import torch

    device = next(model.parameters()).device
    tokens = model.tokenize(code_batch, max_length=max_length, padding=True)
    tokens_tensor = torch.tensor(tokens).to(device)
    with torch.no_grad():
        _, sentence_embeddings = model(tokens_tensor)
    embeddings = sentence_embeddings.cpu().numpy()
    return [_vector_to_list(embedding) for embedding in embeddings]
