"""NuExtract 1.5 Smol adapter for on-demand CPU extraction.

Provides an isolated interface for NuExtract inference:
- Production mode: loads numind/NuExtract-1.5-smol on CPU
- Test mode: deterministic schema-based mock extraction

This adapter uses the shared InferenceGateway with a configurable
target for CPU-only deployment. It is NOT always-resident — loaded
on demand for benchmark or promoted document classes only.

Requirement: 6.6
"""
from __future__ import annotations

import logging
import re
import time
from typing import Any

from services.intelligence_pipeline_v3.nuextract.models import (
    ExtractedField,
    NuExtractResult,
)

logger = logging.getLogger(__name__)

# Pinned model configuration
NUEXTRACT_MODEL_NAME = "numind/NuExtract-1.5-smol"
NUEXTRACT_MODEL_VERSION = "numind/NuExtract-1.5-smol@v1.5"


class NuExtractAdapter:
    """Adapter for NuExtract 1.5 Smol hierarchical extraction.

    Designed for CPU/on-demand use, not always-resident GPU deployment.

    Parameters
    ----------
    test_mode
        When True, uses deterministic schema-based extraction rather than
        loading the model. Useful for testing without model dependencies.
    max_length
        Maximum input text length in characters. Longer texts are processed
        in segments.
    """

    def __init__(
        self,
        test_mode: bool = True,
        max_length: int = 16_000,
    ) -> None:
        self._test_mode = test_mode
        self._max_length = max_length
        self._model = None
        self._tokenizer = None
        self._model_name = NUEXTRACT_MODEL_NAME
        self._model_version = NUEXTRACT_MODEL_VERSION
        self._loaded = False

        if not test_mode:
            self._load_model()

    @property
    def model_version(self) -> str:
        """Return the pinned model version string."""
        return self._model_version

    @property
    def model_name(self) -> str:
        """Return the model name."""
        return self._model_name

    @property
    def is_loaded(self) -> bool:
        """Return whether the model is currently loaded."""
        return self._loaded

    def _load_model(self) -> None:
        """Load the NuExtract model and tokenizer for CPU inference."""
        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer

            logger.info("Loading NuExtract model (CPU): %s", self._model_name)
            self._tokenizer = AutoTokenizer.from_pretrained(self._model_name)
            self._model = AutoModelForCausalLM.from_pretrained(
                self._model_name,
                device_map="cpu",
                torch_dtype="auto",
            )
            self._model.eval()
            self._loaded = True
            logger.info("NuExtract model loaded successfully on CPU")
        except ImportError:
            raise RuntimeError(
                "transformers and torch are required for NuExtract inference. "
                "Install with: pip install transformers torch"
            )
        except Exception as e:
            raise RuntimeError(f"Failed to load NuExtract model: {e}") from e

    def unload(self) -> None:
        """Unload model to free memory (on-demand lifecycle)."""
        if self._model is not None:
            del self._model
            del self._tokenizer
            self._model = None
            self._tokenizer = None
            self._loaded = False
            logger.info("NuExtract model unloaded")

    async def extract(
        self,
        text: str,
        schema: dict[str, Any],
        document_type: str = "",
    ) -> NuExtractResult:
        """Extract structured fields from text using the given schema.

        Parameters
        ----------
        text
            Source document text.
        schema
            JSON schema defining the fields to extract.
        document_type
            Document type identifier for reporting.

        Returns
        -------
        NuExtractResult
            Extracted fields with confidence, latency, and memory metrics.
        """
        start_time = time.perf_counter()

        try:
            if self._test_mode:
                fields = self._extract_test_mode(text, schema)
            else:
                fields = self._extract_production(text, schema)

            latency_ms = (time.perf_counter() - start_time) * 1000
            memory_mb = self._estimate_memory()

            # Compute overall confidence as mean of field confidences
            confidence = 0.0
            if fields:
                confidence = sum(f.confidence for f in fields) / len(fields)

            return NuExtractResult(
                fields=fields,
                spans=[
                    {"start": f.start_char, "end": f.end_char, "field": f.name}
                    for f in fields
                    if f.start_char is not None
                ],
                confidence=confidence,
                model_version=self._model_version,
                latency_ms=latency_ms,
                memory_mb=memory_mb,
                document_type=document_type,
                schema_used=schema,
            )
        except Exception as e:
            latency_ms = (time.perf_counter() - start_time) * 1000
            logger.error("NuExtract extraction failed: %s", e)
            return NuExtractResult(
                model_version=self._model_version,
                latency_ms=latency_ms,
                document_type=document_type,
                schema_used=schema,
                error=str(e),
            )

    def _extract_test_mode(
        self, text: str, schema: dict[str, Any]
    ) -> list[ExtractedField]:
        """Deterministic extraction for testing.

        Matches schema field names against text content using simple
        pattern matching to simulate extraction behavior.
        """
        fields: list[ExtractedField] = []
        properties = schema.get("properties", schema)

        for field_name, field_spec in properties.items():
            # Simple pattern: look for the field name or related keywords in text
            pattern = re.compile(
                rf"\b{re.escape(field_name.replace('_', ' '))}[:\s]+([^\n.;]+)",
                re.IGNORECASE,
            )
            match = pattern.search(text)

            if match:
                value = match.group(1).strip()
                fields.append(
                    ExtractedField(
                        name=field_name,
                        value=value,
                        start_char=match.start(1),
                        end_char=match.end(1),
                        confidence=0.85,
                    )
                )
            else:
                # Try extracting from nearby context for hierarchical schemas
                if isinstance(field_spec, dict) and "properties" in field_spec:
                    # Nested schema — attempt hierarchical extraction
                    nested = self._extract_test_mode(text, field_spec)
                    if nested:
                        fields.append(
                            ExtractedField(
                                name=field_name,
                                value={f.name: f.value for f in nested},
                                confidence=sum(f.confidence for f in nested) / len(nested),
                            )
                        )
                else:
                    # Field not found in text
                    fields.append(
                        ExtractedField(
                            name=field_name,
                            value=None,
                            confidence=0.0,
                        )
                    )

        return fields

    def _extract_production(
        self, text: str, schema: dict[str, Any]
    ) -> list[ExtractedField]:
        """Run NuExtract inference on text using the loaded model."""
        import json

        import torch

        if self._model is None or self._tokenizer is None:
            raise RuntimeError("Model not loaded. Initialize with test_mode=False.")

        # Format input in NuExtract's expected format
        schema_str = json.dumps(schema, indent=2)
        prompt = f"<|input|>\n### Template:\n{schema_str}\n### Text:\n{text[:self._max_length]}\n<|output|>\n"

        inputs = self._tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=4096,
        )

        with torch.no_grad():
            outputs = self._model.generate(
                **inputs,
                max_new_tokens=1024,
                temperature=0.0,
                do_sample=False,
            )

        # Decode and parse the output
        generated = self._tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True,
        )

        return self._parse_output(generated, text, schema)

    def _parse_output(
        self, output: str, source_text: str, schema: dict[str, Any]
    ) -> list[ExtractedField]:
        """Parse model output into structured fields with spans."""
        import json

        fields: list[ExtractedField] = []

        try:
            parsed = json.loads(output)
        except json.JSONDecodeError:
            logger.warning("Failed to parse NuExtract output as JSON")
            return fields

        properties = schema.get("properties", schema)
        for field_name in properties:
            if field_name in parsed:
                value = parsed[field_name]
                # Try to find the value in source text for span
                start_char = None
                end_char = None
                if isinstance(value, str) and value:
                    idx = source_text.find(value)
                    if idx >= 0:
                        start_char = idx
                        end_char = idx + len(value)

                fields.append(
                    ExtractedField(
                        name=field_name,
                        value=value,
                        start_char=start_char,
                        end_char=end_char,
                        confidence=0.8,
                    )
                )

        return fields

    def _estimate_memory(self) -> float:
        """Estimate current memory usage in MB."""
        if self._test_mode:
            return 0.0

        try:
            import torch

            if torch.cuda.is_available():
                return torch.cuda.memory_allocated() / (1024 * 1024)
            # For CPU, estimate from model parameters
            if self._model is not None:
                param_bytes = sum(
                    p.nelement() * p.element_size() for p in self._model.parameters()
                )
                return param_bytes / (1024 * 1024)
        except Exception:
            pass
        return 0.0
