"""Specialist extraction engine — wraps GLiNER2 or a mock for testing."""

from __future__ import annotations

import logging
import os
from typing import Any

from services.specialist.batching import DynamicBatcher
from services.specialist.models import (
    MODEL_VERSION,
    SCHEMA_VERSION,
    ClassificationResult,
    EntityResult,
    RelationResult,
    StructuredResult,
)

logger = logging.getLogger(__name__)


class SpecialistEngine:
    """CPU-first extraction engine backed by GLiNER2 Large or a mock.

    In test mode (SPECIALIST_TEST_MODE=1), uses a mock engine that produces
    deterministic results without downloading model weights.

    Integrates a DynamicBatcher for each extraction type to collect concurrent
    requests and process them as batches.
    """

    def __init__(
        self,
        model_name: str = "urchade/gliner_large-v2.1",
        max_batch_size: int = 32,
        max_wait_ms: float = 50.0,
        max_queue_size: int = 256,
    ) -> None:
        self.model_name = model_name
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
        self.max_queue_size = max_queue_size
        self._model: Any = None
        self._test_mode = os.environ.get("SPECIALIST_TEST_MODE", "1") == "1"
        self._ready = False

        # Dynamic batchers (created but not started until load_model)
        self._entity_batcher: DynamicBatcher | None = None
        self._classification_batcher: DynamicBatcher | None = None
        self._relation_batcher: DynamicBatcher | None = None
        self._structured_batcher: DynamicBatcher | None = None

    @property
    def is_ready(self) -> bool:
        return self._ready

    @property
    def batcher_metrics(self) -> dict[str, Any]:
        """Return aggregated metrics from all batchers."""
        metrics: dict[str, Any] = {
            "total_batches": 0,
            "total_items": 0,
            "total_rejections": 0,
            "queue_depth": 0,
        }
        for batcher in [
            self._entity_batcher,
            self._classification_batcher,
            self._relation_batcher,
            self._structured_batcher,
        ]:
            if batcher:
                metrics["total_batches"] += batcher.total_batches_processed
                metrics["total_items"] += batcher.total_items_processed
                metrics["total_rejections"] += batcher.total_rejections
                metrics["queue_depth"] += batcher.queue_size
        return metrics

    def load_model(self, model_name: str | None = None) -> None:
        """Load the specialist model or mock engine."""
        if model_name:
            self.model_name = model_name

        if self._test_mode:
            logger.info("Specialist engine starting in TEST mode (mock)")
            self._model = _MockGLiNER()
            self._ready = True
            return

        try:
            from gliner import GLiNER  # type: ignore[import-untyped]

            logger.info("Loading GLiNER model: %s", self.model_name)
            self._model = GLiNER.from_pretrained(self.model_name)
            self._ready = True
            logger.info("GLiNER model loaded successfully")
        except ImportError:
            logger.warning(
                "gliner package not available — falling back to mock engine"
            )
            self._model = _MockGLiNER()
            self._ready = True
        except Exception:
            logger.exception("Failed to load GLiNER model")
            self._model = _MockGLiNER()
            self._ready = True

    def warm_up(self) -> None:
        """Run a dummy inference to warm up model weights and caches."""
        if not self._ready:
            self.load_model()
        dummy_text = "Apple Inc reported Q3 revenue of $81.4 billion."
        dummy_labels = ["company", "financial_metric", "date"]
        self.extract_entities([dummy_text], dummy_labels)
        logger.info("Specialist engine warm-up complete")

    def extract_entities(
        self, texts: list[str], labels: list[str]
    ) -> list[list[EntityResult]]:
        """Extract entities from a batch of texts."""
        if not self._ready:
            raise RuntimeError("Engine not initialized — call load_model() first")

        results: list[list[EntityResult]] = []
        for text in texts:
            entities = self._predict_entities(text, labels)
            results.append(entities)
        return results

    def classify_texts(
        self, texts: list[str], labels: list[str]
    ) -> list[list[ClassificationResult]]:
        """Classify texts against the given labels."""
        if not self._ready:
            raise RuntimeError("Engine not initialized — call load_model() first")

        results: list[list[ClassificationResult]] = []
        for text in texts:
            classifications = self._predict_classification(text, labels)
            results.append(classifications)
        return results

    def extract_relations(
        self, texts: list[str], labels: list[str]
    ) -> list[list[RelationResult]]:
        """Extract relations from a batch of texts."""
        if not self._ready:
            raise RuntimeError("Engine not initialized — call load_model() first")

        results: list[list[RelationResult]] = []
        for text in texts:
            relations = self._predict_relations(text, labels)
            results.append(relations)
        return results

    def extract_structured(
        self, texts: list[str], labels: list[str]
    ) -> list[list[StructuredResult]]:
        """Extract structured key-value facts from texts."""
        if not self._ready:
            raise RuntimeError("Engine not initialized — call load_model() first")

        results: list[list[StructuredResult]] = []
        for text in texts:
            structured = self._predict_structured(text, labels)
            results.append(structured)
        return results

    # ------------------------------------------------------------------
    # Internal prediction methods
    # ------------------------------------------------------------------

    def _predict_entities(self, text: str, labels: list[str]) -> list[EntityResult]:
        """Run entity prediction for a single text."""
        if self._test_mode or isinstance(self._model, _MockGLiNER):
            return self._model.predict_entities(text, labels)

        # Real GLiNER inference
        raw_entities = self._model.predict_entities(text, labels)
        return [
            EntityResult(
                text=ent["text"],
                entity_type=ent["label"],
                start_char=ent["start"],
                end_char=ent["end"],
                score=round(float(ent["score"]), 4),
                model_version=MODEL_VERSION,
                schema_version=SCHEMA_VERSION,
            )
            for ent in raw_entities
        ]

    def _predict_classification(
        self, text: str, labels: list[str]
    ) -> list[ClassificationResult]:
        """Run classification for a single text using zero-shot NER heuristic."""
        if self._test_mode or isinstance(self._model, _MockGLiNER):
            return self._model.predict_classification(text, labels)

        # Use entity extraction as a proxy for classification
        raw_entities = self._model.predict_entities(text, labels)
        # Group by label and take the highest score per label
        label_scores: dict[str, float] = {}
        for ent in raw_entities:
            lbl = ent["label"]
            score = float(ent["score"])
            if lbl not in label_scores or score > label_scores[lbl]:
                label_scores[lbl] = score

        return [
            ClassificationResult(
                text=text[:200],
                label=lbl,
                score=round(score, 4),
                model_version=MODEL_VERSION,
                schema_version=SCHEMA_VERSION,
            )
            for lbl, score in label_scores.items()
        ]

    def _predict_relations(
        self, text: str, labels: list[str]
    ) -> list[RelationResult]:
        """Run relation extraction for a single text."""
        if self._test_mode or isinstance(self._model, _MockGLiNER):
            return self._model.predict_relations(text, labels)

        # Real GLiNER does not natively support relations — use mock pattern
        return self._model.predict_relations(text, labels)

    def _predict_structured(
        self, text: str, labels: list[str]
    ) -> list[StructuredResult]:
        """Run structured fact extraction for a single text."""
        if self._test_mode or isinstance(self._model, _MockGLiNER):
            return self._model.predict_structured(text, labels)

        # Real GLiNER structured extraction uses entity spans as key-value pairs
        raw_entities = self._model.predict_entities(text, labels)
        return [
            StructuredResult(
                text=ent["text"],
                field=ent["label"],
                value=ent["text"],
                start_char=ent["start"],
                end_char=ent["end"],
                score=round(float(ent["score"]), 4),
                model_version=MODEL_VERSION,
                schema_version=SCHEMA_VERSION,
            )
            for ent in raw_entities
        ]


class _MockGLiNER:
    """Mock GLiNER model for testing without downloading weights."""

    def predict_entities(self, text: str, labels: list[str]) -> list[EntityResult]:
        """Produce deterministic mock entities based on simple heuristics."""
        results: list[EntityResult] = []

        # Simple keyword-based mock extraction
        _MOCK_ENTITIES = {
            "company": ["Apple", "Google", "Microsoft", "Tesla", "Amazon"],
            "person": ["Elon Musk", "Tim Cook", "Satya Nadella"],
            "financial_metric": ["revenue", "earnings", "EPS", "profit"],
            "date": ["Q1", "Q2", "Q3", "Q4", "2024", "2025"],
            "currency": ["$", "€", "£"],
            "percentage": ["%"],
        }

        for label in labels:
            keywords = _MOCK_ENTITIES.get(label, [])
            for keyword in keywords:
                start = text.find(keyword)
                if start >= 0:
                    results.append(
                        EntityResult(
                            text=keyword,
                            entity_type=label,
                            start_char=start,
                            end_char=start + len(keyword),
                            score=0.85,
                            model_version=MODEL_VERSION,
                            schema_version=SCHEMA_VERSION,
                        )
                    )
        return results

    def predict_classification(
        self, text: str, labels: list[str]
    ) -> list[ClassificationResult]:
        """Produce deterministic mock classification results."""
        results: list[ClassificationResult] = []
        # Assign first label with high score, rest with decreasing
        for i, label in enumerate(labels):
            score = max(0.3, 0.9 - i * 0.2)
            results.append(
                ClassificationResult(
                    text=text[:200],
                    label=label,
                    score=round(score, 4),
                    model_version=MODEL_VERSION,
                    schema_version=SCHEMA_VERSION,
                )
            )
        return results

    def predict_relations(
        self, text: str, labels: list[str]
    ) -> list[RelationResult]:
        """Produce deterministic mock relation results."""
        results: list[RelationResult] = []

        # Simple mock: if text mentions two companies, create a relation
        companies = ["Apple", "Google", "Microsoft", "Tesla", "Amazon"]
        found: list[tuple[str, int]] = []
        for company in companies:
            idx = text.find(company)
            if idx >= 0:
                found.append((company, idx))

        if len(found) >= 2 and labels:
            subj_name, subj_start = found[0]
            obj_name, obj_start = found[1]
            results.append(
                RelationResult(
                    subject=subj_name,
                    subject_type="company",
                    subject_start=subj_start,
                    subject_end=subj_start + len(subj_name),
                    relation=labels[0],
                    object=obj_name,
                    object_type="company",
                    object_start=obj_start,
                    object_end=obj_start + len(obj_name),
                    score=0.78,
                    model_version=MODEL_VERSION,
                    schema_version=SCHEMA_VERSION,
                )
            )
        return results

    def predict_structured(
        self, text: str, labels: list[str]
    ) -> list[StructuredResult]:
        """Produce deterministic mock structured results."""
        results: list[StructuredResult] = []

        # Extract number-like patterns as structured facts
        import re

        for label in labels:
            # Find dollar amounts
            pattern = r"\$[\d,]+\.?\d*\s*(?:billion|million|thousand)?"
            matches = list(re.finditer(pattern, text))
            for match in matches:
                results.append(
                    StructuredResult(
                        text=match.group(),
                        field=label,
                        value=match.group(),
                        start_char=match.start(),
                        end_char=match.end(),
                        score=0.82,
                        model_version=MODEL_VERSION,
                        schema_version=SCHEMA_VERSION,
                    )
                )
                break  # one per label for mock
        return results