"""
StructuredExtractStep — optional structured JSON field extraction for extractor v1.2.0.

Uses gliner2 batch_extract_json with a user-supplied JSON schema dict.
Runs on elements (pre-chunking), targeting specific element types.

Falls back gracefully if gliner2 not installed or inference fails.

Usage:
    from extractor.structured_extract import StructuredExtractStep
    schema = {"title": "string", "authors": "list[string]"}
    step = StructuredExtractStep(schema=schema)
    elements = list(step.annotate(elements))
    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

logger = logging.getLogger(__name__)

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

DEFAULT_ON_TYPES: frozenset[str] = frozenset(
    {
        "text:abstract",
        "text:narrative",
        "text:paragraph",
        "scientific:reference_entry",
    }
)

_MAX_STRUCT_CHARS = 1024
_STRUCT_STREAMING_BATCH = 8


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


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


# ---------------------------------------------------------------------------
# StructuredExtractStep
# ---------------------------------------------------------------------------


class StructuredExtractStep:
    """
    Annotates elements with structured field extraction using a JSON schema.

    Only processes element types in on_types. Writes results to
    element.structured_fields if the extracted result is a non-empty dict.
    """

    def __init__(
        self,
        schema: dict[str, Any],
        on_types: set[str] | None = None,
        model_id: str = "fastino/gliner2-base-v1",
        use_fp16: bool = True,
        streaming_batch: int = _STRUCT_STREAMING_BATCH,
    ) -> None:
        self.schema = schema
        self.on_types: frozenset[str] = (
            frozenset(on_types) if on_types is not None else DEFAULT_ON_TYPES
        )
        self.model_id = model_id
        self.use_fp16 = use_fp16
        self.streaming_batch = streaming_batch
        self.stats = StructuredExtractStats(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 structured fields.

        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 not in self.on_types:
                    self.stats.elements_skipped += 1
                    continue
                text = (el.text or "").strip()
                if not text:
                    self.stats.elements_skipped += 1
                    continue
                text = text[:_MAX_STRUCT_CHARS]
                eligible_indices.append(i)
                eligible_texts.append(text)
                self.stats.elements_processed += 1

            if eligible_texts:
                try:
                    model = self._get_model()
                    schemas = [self.schema] * len(eligible_texts)
                    raw_results = model.batch_extract_json(eligible_texts, schemas)

                    for idx_in_eligible, elem_idx in enumerate(eligible_indices):
                        el = batch[elem_idx]
                        result = raw_results[idx_in_eligible]
                        # result may be a dict or wrapped object
                        if isinstance(result, dict):
                            fields = result
                        elif hasattr(result, "get"):
                            fields = dict(result)
                        else:
                            fields = {}

                        if not fields:
                            continue

                        # Narrow values to allowed types for Pinecone compatibility
                        narrowed: dict[str, str | int | float | bool | list[str]] = {}
                        for k, v in fields.items():
                            if isinstance(v, (str, int, float, bool)):
                                narrowed[k] = v
                            elif isinstance(v, list):
                                narrowed[k] = [str(item) for item in v]
                            elif v is not None:
                                narrowed[k] = str(v)

                        if narrowed:
                            batch[elem_idx] = el.model_copy(update={"structured_fields": narrowed})
                except Exception as exc:
                    logger.warning("StructuredExtractStep inference failed: %s", exc)
                    self.stats.inference_errors += 1

            yield from batch

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