"""NER engine implementation using spaCy.

Provides named entity extraction with lazy model loading and segment tracking.
"""

from __future__ import annotations

import asyncio
from functools import partial
from typing import TYPE_CHECKING, Final

from noteflow.config.constants import (
    SPACY_MODEL_LG,
    SPACY_MODEL_MD,
    SPACY_MODEL_SM,
    SPACY_MODEL_TRF,
)
from noteflow.domain.entities.named_entity import EntityCategory, NamedEntity
from noteflow.infrastructure.logging import get_logger, log_timing

if TYPE_CHECKING:
    from spacy.language import Language

logger = get_logger(__name__)

# Map spaCy entity types to our categories
_SPACY_CATEGORY_MAP: Final[dict[str, EntityCategory]] = {
    # People
    "PERSON": EntityCategory.PERSON,
    # Organizations
    "ORG": EntityCategory.COMPANY,
    # Products and creative works
    "PRODUCT": EntityCategory.PRODUCT,
    "WORK_OF_ART": EntityCategory.PRODUCT,
    # Locations
    "GPE": EntityCategory.LOCATION,  # Geo-political entity (countries, cities)
    "LOC": EntityCategory.LOCATION,  # Non-GPE locations (mountains, rivers)
    "FAC": EntityCategory.LOCATION,  # Facilities (buildings, airports)
    # Dates and times
    "DATE": EntityCategory.DATE,
    "TIME": EntityCategory.DATE,
    # Others (filtered out or mapped to OTHER)
    "MONEY": EntityCategory.OTHER,
    "PERCENT": EntityCategory.OTHER,
    "CARDINAL": EntityCategory.OTHER,
    "ORDINAL": EntityCategory.OTHER,
    "QUANTITY": EntityCategory.OTHER,
    "NORP": EntityCategory.OTHER,  # Nationalities, religions
    "EVENT": EntityCategory.OTHER,
    "LAW": EntityCategory.OTHER,
    "LANGUAGE": EntityCategory.OTHER,
}

# Entity types to skip (low value for meeting context)
_SKIP_ENTITY_TYPES: Final[frozenset[str]] = frozenset({
    "CARDINAL",
    "ORDINAL",
    "QUANTITY",
    "PERCENT",
    "MONEY",
})

# Valid model names
VALID_SPACY_MODELS: Final[tuple[str, ...]] = (
    SPACY_MODEL_SM,
    SPACY_MODEL_MD,
    SPACY_MODEL_LG,
    SPACY_MODEL_TRF,
)


class NerEngine:
    """Named entity recognition engine using spaCy.

    Lazy-loads the spaCy model on first use to avoid startup delay.
    Implements the NerPort protocol for hexagonal architecture.

    Uses chunking by segment (speaker turn/paragraph) to avoid OOM on
    long transcripts while maintaining segment tracking.
    """

    def __init__(self, model_name: str = SPACY_MODEL_TRF) -> None:
        """Initialize NER engine.

        Args:
            model_name: spaCy model to use. Defaults to transformer model
                       for higher accuracy.
        """
        if model_name not in VALID_SPACY_MODELS:
            raise ValueError(
                f"Invalid model name: {model_name}. "
                f"Valid models: {', '.join(VALID_SPACY_MODELS)}"
            )
        self._model_name = model_name
        self._nlp: Language | None = None

    def load_model(self) -> None:
        """Load the spaCy model.

        Raises:
            RuntimeError: If model loading fails.
        """
        import spacy

        with log_timing("ner_model_load", model_name=self._model_name):
            try:
                self._nlp = spacy.load(self._model_name)
            except OSError as e:
                msg = (
                    f"Failed to load spaCy model '{self._model_name}'. "
                    f"Run: python -m spacy download {self._model_name}"
                )
                raise RuntimeError(msg) from e

        logger.info("ner_model_loaded", model_name=self._model_name)

    def _ensure_loaded(self) -> Language:
        """Ensure model is loaded, loading if necessary.

        Returns:
            The loaded spaCy Language model.
        """
        if self._nlp is None:
            self.load_model()
        assert self._nlp is not None, "load_model() should set self._nlp"
        return self._nlp

    def is_ready(self) -> bool:
        """Check if model is loaded."""
        return self._nlp is not None

    def unload(self) -> None:
        """Unload the model to free memory."""
        self._nlp = None
        logger.info("spaCy model unloaded")

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

    def extract(self, text: str) -> list[NamedEntity]:
        """Extract named entities from text.

        Args:
            text: Input text to analyze.

        Returns:
            List of extracted entities (deduplicated by normalized text).
        """
        if not text or not text.strip():
            return []

        nlp = self._ensure_loaded()
        doc = nlp(text)

        entities: list[NamedEntity] = []
        seen: set[str] = set()

        for ent in doc.ents:
            # Normalize for deduplication
            normalized = ent.text.lower().strip()
            if not normalized or normalized in seen:
                continue

            # Skip low-value entity types
            if ent.label_ in _SKIP_ENTITY_TYPES:
                continue

            seen.add(normalized)
            category = _SPACY_CATEGORY_MAP.get(ent.label_, EntityCategory.OTHER)

            entities.append(
                NamedEntity.create(
                    text=ent.text,
                    category=category,
                    segment_ids=[],  # Filled by caller
                    confidence=0.8,  # spaCy doesn't provide per-entity confidence
                )
            )

        return entities

    def extract_from_segments(
        self,
        segments: list[tuple[int, str]],
    ) -> list[NamedEntity]:
        """Extract entities from multiple segments with segment tracking.

        Processes each segment individually (chunking by speaker turn/paragraph)
        to avoid OOM on long transcripts. Entities appearing in multiple segments
        are deduplicated with merged segment lists.

        Args:
            segments: List of (segment_id, text) tuples.

        Returns:
            Entities with segment_ids populated (deduplicated across segments).
        """
        if not segments:
            return []

        # Track entities and their segment occurrences
        # Key: normalized text, Value: NamedEntity
        all_entities: dict[str, NamedEntity] = {}

        for segment_id, text in segments:
            if not text or not text.strip():
                continue

            segment_entities = self.extract(text)

            for entity in segment_entities:
                key = entity.normalized_text

                if key in all_entities:
                    # Merge segment IDs
                    all_entities[key].merge_segments([segment_id])
                else:
                    # New entity - set its segment_ids
                    entity.segment_ids = [segment_id]
                    all_entities[key] = entity

        return list(all_entities.values())

    async def extract_async(self, text: str) -> list[NamedEntity]:
        """Extract entities asynchronously using executor.

        Offloads blocking extraction to a thread pool executor to avoid
        blocking the asyncio event loop.

        Args:
            text: Input text to analyze.

        Returns:
            List of extracted entities.
        """
        return await self._run_in_executor(partial(self.extract, text))

    async def extract_from_segments_async(
        self,
        segments: list[tuple[int, str]],
    ) -> list[NamedEntity]:
        """Extract entities from segments asynchronously.

        Args:
            segments: List of (segment_id, text) tuples.

        Returns:
            Entities with segment_ids populated.
        """
        return await self._run_in_executor(partial(self.extract_from_segments, segments))

    async def _run_in_executor(
        self,
        func: partial[list[NamedEntity]],
    ) -> list[NamedEntity]:
        """Run sync extraction in thread pool executor."""
        return await asyncio.get_running_loop().run_in_executor(None, func)
