import asyncio
import logging

import numpy as np

logger = logging.getLogger(__name__)

_model = None
_predict_lock = asyncio.Lock()


def load_reranker() -> None:
    global _model
    from axon.config import settings
    from sentence_transformers import CrossEncoder

    logger.info("Loading reranker model: %s", settings.reranker_model)
    _model = CrossEncoder(settings.reranker_model)
    logger.info("Reranker model loaded")


def get_reranker() -> "CrossEncoder | None":
    return _model


def chunk_text(text: str, max_tokens: int = 512, max_chunks: int = 4) -> list[str]:
    if not text:
        return [""]
    words = text.split()
    words_per_chunk = max(1, int(max_tokens * 0.75))
    chunks = [
        " ".join(words[i : i + words_per_chunk])
        for i in range(0, len(words), words_per_chunk)
    ]
    return (chunks or [""])[:max_chunks]


def _sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-x))


async def rerank_results(
    query: str,
    contents: list[tuple[str, str]],
) -> list[tuple[str, str, float]]:
    if not contents:
        return []

    # Build (url, chunks) pairs and flat passage list
    url_chunks: list[tuple[str, list[str]]] = []
    flat_passages: list[str] = []
    for url, text in contents:
        chunks = chunk_text(text)
        url_chunks.append((url, chunks))
        flat_passages.extend(chunks)

    pairs = [[query, p] for p in flat_passages]

    def _predict():
        return _model.predict(pairs)

    async with _predict_lock:
        raw_scores = await asyncio.to_thread(_predict)
    scores: list[float] = _sigmoid(np.array(raw_scores, dtype=float)).tolist()

    # Pick best chunk per URL
    results: list[tuple[str, str, float]] = []
    idx = 0
    for url, chunks in url_chunks:
        chunk_scores = scores[idx : idx + len(chunks)]
        best_i = int(np.argmax(chunk_scores))
        results.append((url, chunks[best_i], round(chunk_scores[best_i], 6)))
        idx += len(chunks)

    return sorted(results, key=lambda t: t[2], reverse=True)