"""
ClassifyStep — optional document classification for extractor v1.2.0.

Uses gliner2 batch_extract_entities with category label strings as entity type
queries. The category whose spans have the highest aggregate confidence score
wins. Runs on chunks (post-chunking). Graceful fallback if gliner2 not installed.

Usage:
    from extractor.classify import ClassifyStep
    step = ClassifyStep(labels=["news article", "scientific paper"])
    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, CategoryScore, ClassificationResult

logger = logging.getLogger(__name__)

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

DEFAULT_CLASSIFY_LABELS: list[str] = [
    "scientific paper or research article",
    "technical documentation or manual",
    "news article or press release",
    "legal document or contract",
    "financial report or earnings",
    "educational content or textbook",
    "product description or marketing",
    "code documentation or API reference",
    "meeting notes or transcript",
    "general prose or narrative",
]

# Element types that are never sent to classification
CLASSIFY_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_CLASSIFY_CHARS = 512
_CLASSIFY_STREAMING_BATCH = 32


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


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


# ---------------------------------------------------------------------------
# ClassifyStep
# ---------------------------------------------------------------------------


class ClassifyStep:
    """
    Annotates a stream of BaseElement/ChunkRecord records with a classification label.

    Uses gliner2 batch_extract_entities with category labels as entity type strings.
    The label with the highest max-span confidence across the text wins.

    Processes elements in fixed-size streaming batches.
    """

    def __init__(
        self,
        labels: list[str] | None = None,
        model_id: str = "fastino/gliner2-base-v1",
        use_fp16: bool = True,
        streaming_batch: int = _CLASSIFY_STREAMING_BATCH,
    ) -> None:
        self.labels = labels or DEFAULT_CLASSIFY_LABELS
        self.model_id = model_id
        self.use_fp16 = use_fp16
        self.streaming_batch = streaming_batch
        self.stats = ClassifyStats(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 a top classification label.

        Processes elements in fixed-size streaming batches.
        Classification failure per batch is isolated: affected elements are
        yielded without classification rather than propagating the exception.
        """
        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 CLASSIFY_SKIP_TYPES:
                    self.stats.elements_skipped += 1
                    continue
                text = (el.text or "").strip()
                if not text:
                    self.stats.elements_skipped += 1
                    continue
                text = text[:_MAX_CLASSIFY_CHARS]
                eligible_indices.append(i)
                eligible_texts.append(text)
                self.stats.elements_processed += 1

            if eligible_texts:
                try:
                    model = self._get_model()
                    # Use batch_extract_entities: labels treated as entity type queries.
                    # For each text, find the label with the highest span confidence.
                    raw_results = model.batch_extract_entities(
                        eligible_texts,
                        self.labels,
                        include_confidence=True,
                        include_spans=False,
                    )
                    for idx_in_eligible, elem_idx in enumerate(eligible_indices):
                        el = batch[elem_idx]
                        result = raw_results[idx_in_eligible]
                        entities_map: dict[str, Any] = result.get("entities") or {}

                        # Build per-label score: max confidence across all spans
                        label_scores: dict[str, float] = {}
                        for label in self.labels:
                            spans = entities_map.get(label) or []
                            if spans:
                                if isinstance(spans, list):
                                    confs = [
                                        s.get("confidence", 0.0) if isinstance(s, dict) else 0.0
                                        for s in spans
                                    ]
                                    label_scores[label] = max(confs) if confs else 0.0
                                else:
                                    label_scores[label] = 0.0
                            else:
                                label_scores[label] = 0.0

                        if not label_scores or max(label_scores.values()) == 0.0:
                            # No signal — leave classification as None
                            continue

                        top_label = max(label_scores, key=lambda k: label_scores[k])
                        top_score = label_scores[top_label]

                        scores = [
                            CategoryScore(label=lbl, score=round(sc, 4))
                            for lbl, sc in sorted(label_scores.items(), key=lambda x: -x[1])
                        ]

                        classification = ClassificationResult(
                            top_label=top_label,
                            top_score=round(top_score, 4),
                            scores=scores,
                            extraction_method=self.model_id,
                        )
                        batch[elem_idx] = el.model_copy(
                            update={
                                "classification": classification,
                                "doc_category": top_label,
                                "doc_category_score": round(top_score, 4),
                            }
                        )
                except Exception as exc:
                    logger.warning("ClassifyStep inference failed: %s", exc)
                    self.stats.inference_errors += 1

            yield from batch

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