import asyncio
import os
import re
from functools import lru_cache

import httpx
from fastembed import SparseTextEmbedding

from config import (
    EMBEDDINGS_DENSE_MODEL,
    EMBEDDINGS_DENSE_URL,
    SPARSE_MODEL_NAME,
    get_upstream_kwargs,
    logger,
)
from schemas import DenseEmbeddingResponse, Question, SparseVector


@lru_cache(maxsize=1)
def get_sparse_model() -> SparseTextEmbedding:
    logger.info("Loading local sparse model %s", SPARSE_MODEL_NAME)
    return SparseTextEmbedding(model_name=SPARSE_MODEL_NAME)


async def embed_dense(client: httpx.AsyncClient, text: str) -> list[float]:
    response = await client.post(
        str(EMBEDDINGS_DENSE_URL),
        **get_upstream_kwargs(),
        json={
            "model": os.getenv("EMBEDDINGS_DENSE_MODEL", EMBEDDINGS_DENSE_MODEL),
            "input": [text],
        },
    )
    response.raise_for_status()
    payload = DenseEmbeddingResponse.model_validate(response.json())
    if not payload.data:
        raise ValueError("Dense embedding response is empty")
    return payload.data[0].embedding


async def embed_dense_multi(client: httpx.AsyncClient, texts: list[str]) -> list[list[float]]:
    tasks = [embed_dense(client, t) for t in texts]
    return list(await asyncio.gather(*tasks))


def embed_sparse(text: str) -> SparseVector:
    vectors = list(get_sparse_model().embed([text]))
    if not vectors:
        raise ValueError("Sparse embedding response is empty")
    item = vectors[0]
    return SparseVector(
        indices=[int(i) for i in item.indices.tolist()],
        values=[float(v) for v in item.values.tolist()],
    )


def _normalize_query(text: str) -> str:
    return re.sub(r"\s+", " ", text).strip()


def build_primary_query(question: Question) -> str:
    q = question.search_text.strip() if question.search_text else ""
    if not q:
        q = question.text.strip()
    return _normalize_query(q)


def build_extra_dense_queries(question: Question) -> list[str]:
    extras: list[str] = []
    for v in question.variants or []:
        q = _normalize_query(v)
        if q:
            extras.append(q)
    for h in question.hyde or []:
        q = _normalize_query(h)
        if q:
            extras.append(q)
    return extras


def build_sparse_query(question: Question) -> str:
    kws = question.keywords or []
    if kws:
        return " ".join(kws)
    return build_primary_query(question)


def build_entity_tokens(question: Question) -> list[str]:
    tokens: list[str] = []
    if question.entities:
        for field in (
            question.entities.people,
            question.entities.emails,
            question.entities.documents,
            question.entities.names,
            question.entities.links,
        ):
            tokens.extend(field or [])
    return [t.strip() for t in tokens if t.strip()]