"""Entity linking from question text (or answer strings) to graph nodes.

Order of evidence: (1) exact match of a node name or alias as a whole word/phrase,
longest first, case-insensitive; (2) paper titles by fuzzy match (questions quote
titles); (3) rapidfuzz ratio of normalised strings for single answer strings.
"""

from __future__ import annotations

import re
from dataclasses import dataclass

from rapidfuzz import fuzz, process

from .canonical import normalise
from .graph import KG


@dataclass
class Link:
    node_id: str
    surface: str
    score: float
    method: str


class EntityLinker:
    def __init__(self, kg: KG, fuzzy_threshold: float = 90.0, min_alias_len: int = 2):
        self.kg = kg
        self.fuzzy_threshold = fuzzy_threshold
        self.alias_index: dict[str, set[str]] = {}
        for n in kg.nodes.values():
            if n.type == "Paper":
                continue
            # The canonical key itself (e.g. an abbreviation) may never occur as a surface form.
            for form in {n.name, *n.aliases, n.node_id.split(":", 1)[1]}:
                key = normalise(form, n.type)
                if len(key) >= min_alias_len:
                    self.alias_index.setdefault(key, set()).add(n.node_id)
                raw = normalise(form)
                if len(raw) >= min_alias_len:
                    self.alias_index.setdefault(raw, set()).add(n.node_id)
        self.titles = {normalise(n.name): n.node_id for n in kg.nodes.values() if n.type == "Paper"}
        alternation = "|".join(re.escape(k) for k in sorted(self.alias_index, key=len, reverse=True))
        self._pattern = re.compile(rf"(?<![a-z0-9])(?:{alternation})(?![a-z0-9])") if alternation else None

    def link_question(self, question: str) -> list[Link]:
        q = normalise(question)
        links: list[Link] = []
        taken: list[tuple[int, int]] = []
        # Paper titles first: a quoted title often contains method/dataset names.
        for m in re.finditer(r'"([^"]{8,})"', question):
            hit = self.match_title(m.group(1))
            if hit:
                links.append(hit)
                tq = normalise(m.group(1))
                p = q.find(tq)
                if p >= 0:
                    taken.append((p, p + len(tq)))
        if self._pattern is not None:
            for m in self._pattern.finditer(q):
                if any(a <= m.start() < b for a, b in taken):
                    continue
                for nid in sorted(self.alias_index[m.group(0)]):
                    links.append(Link(nid, m.group(0), 100.0, "exact"))
        seen: set[str] = set()
        out = []
        for lk in links:
            if lk.node_id not in seen:
                seen.add(lk.node_id)
                out.append(lk)
        return out

    def match_title(self, text: str) -> Link | None:
        if not self.titles:
            return None
        best = process.extractOne(normalise(text), list(self.titles), scorer=fuzz.ratio)
        if best and best[1] >= self.fuzzy_threshold:
            return Link(self.titles[best[0]], text, float(best[1]), "title")
        return None

    def link_name(self, name: str, etype: str | None = None) -> Link | None:
        """Link a single entity string (an answer item or a path element)."""
        key_typed = normalise(name, etype) if etype else None
        for key in filter(None, (key_typed, normalise(name), normalise(name, "Dataset"))):
            ids = sorted(self.alias_index.get(key, ()))
            if etype:
                ids = [i for i in ids if self.kg.nodes[i].type == etype] or ids
            if ids:
                return Link(ids[0], name, 100.0, "exact")
        title = self.match_title(name)
        if title:
            return title
        if not self.alias_index:
            return None
        best = process.extractOne(normalise(name), list(self.alias_index), scorer=fuzz.ratio)
        if best and best[1] >= self.fuzzy_threshold:
            return Link(min(self.alias_index[best[0]]), name, float(best[1]), "fuzzy")
        return None