import re
import unicodedata
import numpy as np
from rank_bm25 import BM25Okapi
from .schema import Evidence


def normalized(text):
    return re.sub(r"\s+", " ", unicodedata.normalize("NFC", text)).strip()


def terms(text):
    text = "".join(str(unicodedata.decimal(c)) if c.isdecimal() else c for c in text)
    return re.findall(r"\w+(?:[./-]\w+)*", text.lower()) or ["empty"]


def evidence_for(doc):
    return [Evidence(id=f"{doc.id}:p{p.number}", document=doc.name, page=p.number,
                     text=p.english, original=p.original) for p in doc.pages
            if p.english.strip()]


class HybridIndex:
    """MiniLM cosine similarity + BM25, joined with reciprocal rank fusion."""
    def __init__(self, documents, embedder):
        self.embedder = embedder
        self.pages = {e.id: e for d in documents if d.ready for e in evidence_for(d)}
        self.rows = []
        for page in self.pages.values():
            offsets = embedder.tokenizer(page.text, add_special_tokens=False,
                                         return_offsets_mapping=True, truncation=False, verbose=False)["offset_mapping"]
            for start in range(0, len(offsets), 150):
                stop = min(start + 180, len(offsets))
                text = page.text[offsets[start][0]:offsets[stop - 1][1]]
                self.rows.append((page.id, text))
                if stop == len(offsets):
                    break
        texts = [row[1] for row in self.rows]
        self.vectors = embedder.encode(texts, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False)
        self.bm25 = BM25Okapi([terms(text) for text in texts])

    def search(self, question, document_ids, k=4):
        if not document_ids:
            return []
        positions = [i for i, row in enumerate(self.rows) if row[0].split(":")[0] in document_ids]
        if not positions:
            return []
        vector = self.embedder.encode([question], normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False)[0]
        semantic_scores = self.vectors @ vector
        keyword_scores = self.bm25.get_scores(terms(question))
        query_terms = set(terms(question))
        semantic = sorted(positions, key=lambda i: float(semantic_scores[i]), reverse=True)[:20]
        keyword = sorted(positions, key=lambda i: (float(keyword_scores[i]),
                         len(query_terms.intersection(terms(self.rows[i][1])))), reverse=True)[:20]
        scores = {}
        for ranking in [semantic, keyword]:
            for rank, index in enumerate(ranking, 1):
                scores[index] = scores.get(index, 0) + 1 / (60 + rank)
        selected = []
        for index in sorted(scores, key=scores.get, reverse=True):
            page_id = self.rows[index][0]
            if page_id not in selected:
                selected.append(page_id)
            if len(selected) == k:
                break
        # Full source pages preserve nearby conditions that a small chunk can omit.
        return [self.pages[key] for key in selected]