"""
RelationStep — optional relation extraction for extractor v1.2.0.

Uses gliner2 batch_extract_entities with relation predicate strings as entity
type queries. Entity spans whose labels match known predicates are treated as
relation anchors; subject/object candidates are identified by proximity pairing
from surrounding entity spans.

Falls back gracefully if gliner2 not installed or inference fails.

Usage:
    from extractor.relations import RelationStep
    step = RelationStep(relation_types=["caused by", "part of"])
    chunks = list(step.annotate(chunks))
    print(step.stats)
"""

from __future__ import annotations

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

from .schema import BaseElement, RelationTriple

logger = logging.getLogger(__name__)

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

DEFAULT_RELATION_TYPES: list[str] = [
    "caused by",
    "part of",
    "references",
    "defined by",
    "approved by",
    "supersedes",
    "employed by",
    "located in",
    "cites",
    "regulates",
    "developed by",
    "published at",
    "uses",
    "outperforms",
    "founded by",
]

RELATION_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",
        "form:field",
        "form:checkbox",
        "form:label",
    }
)

_MAX_RELATION_CHARS = 1024
_RELATION_STREAMING_BATCH = 16
_MIN_RELATION_TEXT_LEN = 20
_MIN_SCORE = 0.50

# Entity labels used to discover subject/object candidates
_ENTITY_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",
    "name of a conference, workshop, or event",
    "title of a research paper or publication",
]


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


@dataclass
class RelationStats:
    elements_processed: int = 0
    elements_skipped: int = 0
    inference_errors: int = 0
    total_relations: int = 0
    duration_ms: float = 0.0
    model_id: str = ""


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _extract_spans(result: dict[str, Any], labels: list[str]) -> list[dict[str, Any]]:
    """Flatten entity spans from batch_extract_entities result into a sorted list."""
    spans: list[dict[str, Any]] = []
    entities_map = result.get("entities") or {}
    for label in labels:
        for span in entities_map.get(label) or []:
            if isinstance(span, dict):
                spans.append(
                    {
                        "label": label,
                        "text": span.get("text", ""),
                        "start": span.get("start", 0),
                        "end": span.get("end", 0),
                        "confidence": float(span.get("confidence", 0.0)),
                    }
                )
    spans.sort(key=lambda s: s["start"])
    return spans


def _build_triples(
    predicate_spans: list[dict[str, Any]],
    entity_spans: list[dict[str, Any]],
    source_text: str,
    min_score: float,
    model_id: str,
) -> list[RelationTriple]:
    """
    Pair predicate spans with nearest preceding and following entity spans
    to form (subject, predicate, object) triples.
    """
    triples: list[RelationTriple] = []
    for pred in predicate_spans:
        score = pred["confidence"]
        if score < min_score:
            continue
        pred_start = pred["start"]
        pred_end = pred["end"]

        # Subject: closest entity ending before the predicate
        subject_span = None
        for ent in reversed(entity_spans):
            if ent["end"] <= pred_start:
                subject_span = ent
                break

        # Object: closest entity starting after the predicate
        object_span = None
        for ent in entity_spans:
            if ent["start"] >= pred_end:
                object_span = ent
                break

        if subject_span is None or object_span is None:
            continue

        try:
            triple = RelationTriple(
                subject=subject_span["text"],
                predicate=pred["label"],
                object=object_span["text"],
                score=round(score, 4),
                subject_type=subject_span["label"],
                object_type=object_span["label"],
                source_text=source_text[:200],
                extraction_method=model_id,
            )
            triples.append(triple)
        except Exception:
            pass
    return triples


# ---------------------------------------------------------------------------
# RelationStep
# ---------------------------------------------------------------------------


class RelationStep:
    """
    Annotates a stream of BaseElement/ChunkRecord records with relation triples.

    Uses gliner2 batch_extract_entities:
    - First pass: relation predicate labels → locate predicate spans
    - Second pass: entity labels → locate subject/object candidates
    - Pairs are formed by proximity (nearest entity before/after predicate)
    """

    def __init__(
        self,
        relation_types: list[str] | None = None,
        model_id: str = "fastino/gliner2-base-v1",
        use_fp16: bool = True,
        streaming_batch: int = _RELATION_STREAMING_BATCH,
        min_score: float = _MIN_SCORE,
    ) -> None:
        self.relation_types = relation_types or DEFAULT_RELATION_TYPES
        self.model_id = model_id
        self.use_fp16 = use_fp16
        self.streaming_batch = streaming_batch
        self.min_score = min_score
        self.stats = RelationStats(model_id=model_id)
        self._model: Any = None

    def _get_model(self) -> Any:
        if self._model is None:
            from .gliner2_manager import get_model

            self._model = get_model(self.model_id, self.use_fp16)
        return self._model

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

        Processes elements in fixed-size streaming batches.
        Failure per batch is isolated; affected elements are yielded unmodified.
        """
        t_start = time.monotonic()
        element_iter = iter(elements)

        while True:
            batch = list(itertools.islice(element_iter, self.streaming_batch))
            if not batch:
                break

            eligible_indices: list[int] = []
            eligible_texts: list[str] = []

            for i, el in enumerate(batch):
                if el.element_type in RELATION_SKIP_TYPES:
                    self.stats.elements_skipped += 1
                    continue
                text = (el.text or "").strip()
                if len(text) < _MIN_RELATION_TEXT_LEN:
                    self.stats.elements_skipped += 1
                    continue
                text = text[:_MAX_RELATION_CHARS]
                eligible_indices.append(i)
                eligible_texts.append(text)
                self.stats.elements_processed += 1

            if eligible_texts:
                try:
                    model = self._get_model()
                    all_labels = self.relation_types + _ENTITY_LABELS

                    raw_results = model.batch_extract_entities(
                        eligible_texts,
                        all_labels,
                        include_confidence=True,
                        include_spans=True,
                    )

                    for idx_in_eligible, elem_idx in enumerate(eligible_indices):
                        el = batch[elem_idx]
                        result = raw_results[idx_in_eligible]
                        source_text = eligible_texts[idx_in_eligible]

                        predicate_spans = _extract_spans(result, self.relation_types)
                        entity_spans = _extract_spans(result, _ENTITY_LABELS)

                        triples = _build_triples(
                            predicate_spans,
                            entity_spans,
                            source_text,
                            self.min_score,
                            self.model_id,
                        )

                        if not triples:
                            batch[elem_idx] = el.model_copy(update={"relations": []})
                        else:
                            self.stats.total_relations += len(triples)
                            subjects = sorted({t.subject for t in triples})
                            objects = sorted({t.object for t in triples})
                            update: dict[str, Any] = {
                                "relations": triples,
                            }
                            # ChunkRecord has relation_subjects/relation_objects
                            if hasattr(el, "relation_subjects"):
                                update["relation_subjects"] = subjects
                                update["relation_objects"] = objects
                            batch[elem_idx] = el.model_copy(update=update)
                except Exception as exc:
                    logger.warning("RelationStep inference failed: %s", exc)
                    self.stats.inference_errors += 1

            yield from batch

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