"""
GLiNER2 NER step — optional entity extraction for extractor v1.2.0.

Local inference only: fastino/gliner2-base-v1 (default) via [ner] extras.
Uses the gliner2 package (pip install gliner2), NOT the gliner package.

Graceful fallback:
  - If [ner] extras not installed: warn once, skip NER.
  - NER failure never fails extraction: try/except per batch → entities=[].

Usage:
    from extractor.ner import NERStep, preload_model
    step = NERStep(labels=["person", "organization"], threshold=0.50)
    elements = step.annotate(elements)
"""

from __future__ import annotations

import contextlib
import itertools
import logging
import re
import time
from collections.abc import Iterable, Iterator
from dataclasses import dataclass
from typing import Any

from .schema import BaseElement, EntityAnnotation

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

DEFAULT_LOCAL_MODEL_ID = "fastino/gliner2-base-v1"

DEFAULT_LABELS: list[str] = [
    "name of a specific person or researcher",
    "name of an organization, company, or institution",
    "name of a software framework, library, or tool",
    "name of a machine learning model or dataset",
    "name of a location, city, country, or region",
    "specific date, time period, or year",
    "numerical measurement with explicit units",
    "name of a conference, workshop, or event",
    "title of a research paper or publication",
]

# Per-label thresholds calibrated for fastino/gliner2-base-v1
# gliner2 scores cluster near 1.0; 0.70 cleanly separates signal from noise.
DEFAULT_LABEL_THRESHOLDS: dict[str, float] = {
    "name of a specific person or researcher": 0.70,
    "name of a software framework, library, or tool": 0.75,
    "name of a machine learning model or dataset": 0.70,
    "name of an organization, company, or institution": 0.70,
    "name of a location, city, country, or region": 0.70,
    "specific date, time period, or year": 0.70,
    "numerical measurement with explicit units": 0.75,
    "name of a conference, workshop, or event": 0.70,
    "title of a research paper or publication": 0.70,
}

# Element types that are never sent to NER
NER_SKIP_TYPES: frozenset[str] = frozenset(
    {
        "code:block",
        "code:inline",
        "code:cell",
        "code:diff",
        "scientific:equation_display",
        "scientific:equation_inline",
        "media:figure",
        "media:image",
        "media:audio",
        "media:video",
        "structural:page_header",
        "structural:page_footer",
        "structural:page_number",
        "structural:horizontal_rule",
        "structural:divider",
        "meta:citation",
        "meta:reference",
        "meta:footnote_ref",
        "meta:url",
        "meta:email",
        "form:field",
        "form:checkbox",
        "form:label",
    }
)

# Per-type confidence overrides (applied on top of global threshold)
NER_HIGH_THRESHOLD_TYPES: frozenset[str] = frozenset(
    {
        "structural:section_header",
        "structural:title",
        "structural:subtitle",
    }
)
NER_HIGH_THRESHOLD_DELTA = 0.15  # add to global threshold

# Max chars sent to NER (256 cl100k tokens ≈ 1024 chars)
_MAX_NER_CHARS = 1024

# Streaming batch size: number of elements collected per window before inference
_NER_STREAMING_BATCH = 32

# Module-level model singleton: model_id -> model instance
_MODEL_SINGLETON: dict[str, Any] = {}

# ---------------------------------------------------------------------------
# Text preprocessing
# ---------------------------------------------------------------------------

_RE_LATEX_DISPLAY = re.compile(r"\\\[.*?\\\]", re.DOTALL)
_RE_LATEX_INLINE_P = re.compile(r"\\\(.*?\\\)", re.DOTALL)
_RE_LATEX_DOLLAR = re.compile(r"\$[^$\n]+?\$")
_RE_LATEX_CMD = re.compile(r"\\[a-zA-Z]+\{[^}]*\}")
_RE_LATEX_BARE_CMD = re.compile(r"\\[a-zA-Z]+")
_RE_MD_BOLD_ITALIC = re.compile(r"\*{1,3}([^*\n]+?)\*{1,3}")
_RE_MD_CODE_INLINE = re.compile(r"`[^`\n]*`")
_RE_URL = re.compile(r"https?://\S+|www\.\S+", re.IGNORECASE)
_RE_EMAIL = re.compile(r"\S+@\S+\.\S+")
_RE_WHITESPACE = re.compile(r"[ \t]{2,}")

_RE_DIGITS_ONLY = re.compile(r"^\d[\d\s,._\-/]*$")
_RE_PUNCT_ONLY = re.compile(r"^[\W_]+$")
_ANAPHOR_PREFIX = re.compile(r"^(the|a|an)\s", re.IGNORECASE)
_RE_BARE_YEAR = re.compile(r"^\d{4}$")

# Common English words that are never meaningful named entities
_NOISE_WORDS: frozenset[str] = frozenset(
    {
        "model",
        "models",
        "network",
        "networks",
        "layer",
        "layers",
        "encoder",
        "decoder",
        "attention",
        "output",
        "input",
        "result",
        "results",
        "data",
        "dataset",
        "training",
        "test",
        "loss",
    }
)


def _preprocess_for_ner(text: str) -> str:
    """Strip LaTeX, markdown markers, URLs, and collapse whitespace before NER."""
    text = _RE_LATEX_DISPLAY.sub(" ", text)
    text = _RE_LATEX_INLINE_P.sub(" ", text)
    text = _RE_LATEX_DOLLAR.sub(" ", text)
    text = _RE_LATEX_CMD.sub(" ", text)
    text = _RE_LATEX_BARE_CMD.sub(" ", text)
    text = _RE_MD_CODE_INLINE.sub(" ", text)
    text = _RE_MD_BOLD_ITALIC.sub(r"\1", text)
    text = _RE_URL.sub(" ", text)
    text = _RE_EMAIL.sub(" ", text)
    text = _RE_WHITESPACE.sub(" ", text)
    return text.strip()


def _is_noise_annotation(text: str, entity_type: str) -> bool:
    """Return True if this annotation should be discarded post-inference."""
    stripped = text.strip()
    if not stripped:
        return True
    if _RE_PUNCT_ONLY.match(stripped):
        return True
    if _ANAPHOR_PREFIX.match(stripped):
        return True
    # Single lowercase generic words are never named entities
    if stripped.lower() in _NOISE_WORDS:
        return True
    # Bare 4-digit years are low-signal unless entity_type is explicitly date
    if _RE_BARE_YEAR.match(stripped) and "date" not in entity_type and "year" not in entity_type:
        return True
    if "measurement" in entity_type or entity_type == "quantity":
        if len(stripped) < 4:
            return True
        if _RE_DIGITS_ONLY.match(stripped):
            return True
    return False


# ---------------------------------------------------------------------------
# Stats
# ---------------------------------------------------------------------------


@dataclass
class NERStats:
    elements_processed: int = 0
    elements_skipped: int = 0
    elements_truncated: int = 0
    inference_errors: int = 0
    total_entities: int = 0
    duration_ms: float = 0.0
    model_id: str = ""
    tier: str = ""  # "api" or "local"


# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------


def _load_model(model_id: str = DEFAULT_LOCAL_MODEL_ID, use_fp16: bool = True) -> Any:
    """Load and cache the NER model singleton (local inference only)."""
    if model_id in _MODEL_SINGLETON:
        return _MODEL_SINGLETON[model_id]

    from gliner2 import GLiNER2  # type: ignore[import-untyped]

    model = GLiNER2.from_pretrained(model_id)
    if use_fp16:
        with contextlib.suppress(Exception):
            model = model.half()
    _MODEL_SINGLETON[model_id] = model
    logger.debug("NER: loaded local model %s", model_id)
    return model


def preload_model(model_id: str = DEFAULT_LOCAL_MODEL_ID, use_fp16: bool = True) -> None:
    """
    Eagerly load the NER model singleton before processing a batch.
    Call this before the document loop to avoid cold-start penalty on first document.

    Raises ImportError if [ner] extras are not installed.
    """
    _load_model(model_id=model_id, use_fp16=use_fp16)


# ---------------------------------------------------------------------------
# Text extraction helpers
# ---------------------------------------------------------------------------


def _get_ner_text(element: BaseElement) -> str:
    """
    Get the text to send to NER for a given element.
    For tables: use the structured cell text rather than Markdown.
    Text is preprocessed to strip LaTeX, markdown markers, and normalize whitespace.
    """
    if element.element_type.startswith("table:") and element.table is not None:
        rows = element.table.structured.rows
        headers = element.table.structured.headers
        parts: list[str] = []
        if headers:
            parts.append(" ".join(str(h) for h in headers))
        for row in rows:
            parts.append(" ".join(str(c) for c in row))
        return _preprocess_for_ner(" ".join(parts))
    return _preprocess_for_ner(element.text or "")


def _truncate_text(text: str) -> tuple[str, bool]:
    """
    Truncate text to _MAX_NER_CHARS characters.
    Returns (truncated_text, was_truncated).
    """
    if len(text) <= _MAX_NER_CHARS:
        return text, False
    return text[:_MAX_NER_CHARS], True


# ---------------------------------------------------------------------------
# Batch size heuristic
# ---------------------------------------------------------------------------


def _auto_batch_size(texts: list[str]) -> int:
    """
    Choose inference batch size based on average text length.
    Longer texts need smaller batches to avoid OOM.
    """
    if not texts:
        return 8
    avg_len = sum(len(t) for t in texts) / len(texts)
    if avg_len <= 256:
        return 16
    if avg_len <= 512:
        return 8
    if avg_len <= 1024:
        return 4
    return 2


# ---------------------------------------------------------------------------
# Core NER step
# ---------------------------------------------------------------------------


class NERStep:
    """
    Annotates a stream of BaseElement records with named entity spans.

    Usage:
        step = NERStep(labels=["person", "organization"], threshold=0.50)
        annotated_elements = list(step.annotate(elements))
        print(step.stats)
    """

    def __init__(
        self,
        labels: list[str] | None = None,
        threshold: float = 0.50,
        label_thresholds: dict[str, float] | None = None,
        model_id: str = DEFAULT_LOCAL_MODEL_ID,
        use_fp16: bool = True,
        streaming_batch: int = _NER_STREAMING_BATCH,
    ) -> None:
        self.labels = labels or DEFAULT_LABELS
        self.threshold = threshold
        if label_thresholds is not None:
            self.label_thresholds: dict[str, float] = label_thresholds
        elif labels is None:
            self.label_thresholds = DEFAULT_LABEL_THRESHOLDS
        else:
            self.label_thresholds = {}
        self.model_id = model_id
        self.use_fp16 = use_fp16
        self.streaming_batch = streaming_batch
        self.stats = NERStats(model_id=model_id)
        self._model: Any = None

    def _get_model(self) -> Any:
        if self._model is None:
            self._model = _load_model(self.model_id, self.use_fp16)
        return self._model

    def _effective_threshold(self, element: BaseElement, label: str | None = None) -> float:
        base = (
            self.label_thresholds.get(label, self.threshold)
            if label is not None
            else self.threshold
        )
        if element.element_type in NER_HIGH_THRESHOLD_TYPES:
            return min(1.0, base + NER_HIGH_THRESHOLD_DELTA)
        return base

    def annotate(self, elements: Iterable[BaseElement]) -> Iterator[BaseElement]:
        """
        Annotate elements with entity spans.

        Processes elements in fixed-size streaming batches to avoid materialising
        the entire document into memory. NER failure per inference batch is isolated:
        affected elements get entities=[] rather than propagating the exception.
        """
        t_start = time.monotonic()
        self.stats.tier = "local"

        # Process in fixed-size streaming windows
        element_iter = iter(elements)

        while True:
            # Collect next batch of raw elements
            batch = list(itertools.islice(element_iter, self.streaming_batch))
            if not batch:
                break

            # Separate eligible from skipped within this batch
            eligible_indices: list[int] = []
            eligible_texts: list[str] = []

            for i, el in enumerate(batch):
                if el.element_type in NER_SKIP_TYPES:
                    self.stats.elements_skipped += 1
                    continue

                text = _get_ner_text(el)
                if not text.strip():
                    batch[i] = el.model_copy(update={"entities": []})
                    self.stats.elements_processed += 1
                    continue

                text, was_truncated = _truncate_text(text)
                if was_truncated:
                    self.stats.elements_truncated += 1

                eligible_indices.append(i)
                eligible_texts.append(text)
                self.stats.elements_processed += 1

            # Run inference on this batch's eligible texts
            if eligible_texts:
                model = self._get_model()
                inf_batch_size = _auto_batch_size(eligible_texts)
                # all_predictions[i] = flat list of {"label", "text", "score", "start", "end"}
                all_predictions: list[list[dict]] = [[] for _ in eligible_indices]

                for b_start in range(0, len(eligible_texts), inf_batch_size):
                    b_end = b_start + inf_batch_size
                    b_texts = eligible_texts[b_start:b_end]
                    try:
                        # gliner2 API: batch_extract_entities returns list of
                        # {"entities": {"label": [{"text", "confidence", "start", "end"}]}}
                        raw_results = model.batch_extract_entities(
                            b_texts,
                            self.labels,
                            include_confidence=True,
                            include_spans=True,
                        )
                        for j, result in enumerate(raw_results):
                            flat: list[dict] = []
                            for label, spans in (result.get("entities") or {}).items():
                                for span in spans:
                                    flat.append(
                                        {
                                            "label": label,
                                            "text": span.get("text", ""),
                                            "score": span.get("confidence", 0.0),
                                            "start": span.get("start"),
                                            "end": span.get("end"),
                                        }
                                    )
                            all_predictions[b_start + j] = flat
                    except Exception as exc:
                        logger.warning(
                            "NER inference failed on batch %d-%d: %s", b_start, b_end, exc
                        )
                        self.stats.inference_errors += 1

                # Apply predictions back to batch elements
                for idx_in_eligible, elem_idx in enumerate(eligible_indices):
                    el = batch[elem_idx]
                    preds = all_predictions[idx_in_eligible]
                    annotations: list[EntityAnnotation] = []
                    for pred in preds:
                        score = float(pred.get("score", 0.0))
                        label = pred.get("label", "unknown")
                        entity_text = pred.get("text", "")
                        threshold = self._effective_threshold(el, label=label)
                        if score < threshold:
                            continue
                        if _is_noise_annotation(entity_text, label):
                            continue
                        try:
                            ann = EntityAnnotation(
                                entity_type=label,
                                text=entity_text,
                                start=pred.get("start"),
                                end=pred.get("end"),
                                score=round(score, 4),
                                extraction_method=self.model_id,
                            )
                            annotations.append(ann)
                        except Exception:
                            pass
                    self.stats.total_entities += len(annotations)
                    batch[elem_idx] = el.model_copy(update={"entities": annotations})

            # Yield the fully-annotated batch immediately
            yield from batch

        self.stats.duration_ms = (time.monotonic() - t_start) * 1000
