from __future__ import annotations

import re
import unicodedata
from collections import Counter
from math import log

from .models import ParentDocument, RetrievalDocument, SearchHit

WORD_RE = re.compile(r"\w+", re.UNICODE)


def _normalized(text: str) -> str:
    return " ".join(WORD_RE.findall(unicodedata.normalize("NFKC", text).casefold()))


def _terms(text: str) -> set[str]:
    return set(_normalized(text).split())


def _char_ngrams(text: str, size: int = 3) -> set[str]:
    normalized = _normalized(text)
    if len(normalized) <= size:
        return {normalized} if normalized else set()
    return {
        normalized[index:index + size]
        for index in range(len(normalized) - size + 1)
    }


class InMemoryLexicalRetriever:
    """Deterministic test/fallback retriever, not the production neural backend."""

    def __init__(self, documents: list[RetrievalDocument]) -> None:
        self._documents = tuple(documents)
        self._term_counts = {
            document.doc_id: Counter(_normalized(document.text).split())
            for document in self._documents
        }
        self._average_length = (
            sum(sum(counts.values()) for counts in self._term_counts.values())
            / max(1, len(self._term_counts))
        )
        document_frequency: Counter[str] = Counter()
        for counts in self._term_counts.values():
            document_frequency.update(counts.keys())
        self._document_frequency = document_frequency

    def search(self, query: str, drug_id: str, limit: int) -> list[SearchHit]:
        query_terms = _terms(query)
        query_ngrams = _char_ngrams(query)
        candidates = []
        for document in self._documents:
            if document.drug_id != drug_id:
                continue
            counts = self._term_counts[document.doc_id]
            bm25 = self._bm25(query_terms, counts)
            ngrams = _char_ngrams(document.text)
            char_score = len(query_ngrams & ngrams) / max(1, len(query_ngrams))
            if bm25 > 0 or char_score > 0:
                candidates.append((document, bm25, char_score))
        max_bm25 = max((row[1] for row in candidates), default=0.0)
        hits = [
            SearchHit(
                document=document,
                score=0.8 * (bm25 / max_bm25 if max_bm25 else 0.0) + 0.2 * char_score,
            )
            for document, bm25, char_score in candidates
        ]
        return sorted(hits, key=lambda hit: (-hit.score, hit.document.doc_id))[:limit]

    def _bm25(self, query_terms: set[str], counts: Counter[str]) -> float:
        total_documents = len(self._documents)
        document_length = sum(counts.values())
        score = 0.0
        for term in query_terms:
            frequency = counts.get(term, 0)
            if not frequency:
                continue
            document_frequency = self._document_frequency[term]
            inverse_frequency = log(
                1 + (total_documents - document_frequency + 0.5)
                / (document_frequency + 0.5)
            )
            denominator = frequency + 1.5 * (
                1 - 0.75 + 0.75 * document_length / max(1.0, self._average_length)
            )
            score += inverse_frequency * frequency * 2.5 / denominator
        return score


class InMemoryParentStore:
    def __init__(self, parents: list[ParentDocument]) -> None:
        self._parents = {parent.parent_id: parent for parent in parents}

    def get(self, parent_id: str) -> ParentDocument | None:
        return self._parents.get(parent_id)