"""Hybrid NER: ML-powered with regex fallback."""

from __future__ import annotations

import logging
from typing import Any

from neurovault.cognitive.entity_extractor import EntityExtractor as RegexExtractor

logger = logging.getLogger(__name__)


class HybridNERExtractor:
    """ML-powered NER with regex fallback."""

    SPACY_TYPE_MAP = {
        "PERSON": "PERSON",
        "ORG": "ORGANIZATION",
        "GPE": "LOCATION",
        "LOC": "LOCATION",
        "DATE": "DATE",
        "MONEY": "METRIC",
        "PERCENT": "METRIC",
        "QUANTITY": "METRIC",
        "EVENT": "EVENT",
        "PRODUCT": "TECHNOLOGY",
        "WORK_OF_ART": "CONCEPT",
        "FAC": "LOCATION",
    }

    def __init__(self, spacy_model: str = "en_core_web_sm"):
        self._nlp = None
        self._regex = RegexExtractor()
        self._model_name = spacy_model

        try:
            import spacy

            self._nlp = spacy.load(spacy_model)
            logger.info("Loaded spaCy model: %s", spacy_model)
        except (ImportError, OSError) as exc:
            logger.info("spaCy not available (%s), using regex NER", exc)

    @property
    def has_ml(self) -> bool:
        return self._nlp is not None

    def extract(self, content: str) -> list[dict[str, Any]]:
        """Extract entities using ML + regex hybrid approach."""
        regex_entities = self._regex.extract(content)

        if self._nlp is None:
            return regex_entities

        doc = self._nlp(content[:5000])
        ml_entities: list[dict[str, Any]] = []

        for ent in doc.ents:
            mapped_type = self.SPACY_TYPE_MAP.get(ent.label_)
            if mapped_type and len(ent.text.strip()) > 1:
                ml_entities.append(
                    {
                        "name": ent.text.strip(),
                        "type": mapped_type,
                        "confidence": 0.85,
                        "start": ent.start_char,
                        "end": ent.end_char,
                        "source": "spacy",
                    }
                )

        return self._merge_entities(ml_entities, regex_entities)

    @staticmethod
    def _merge_entities(ml: list[dict[str, Any]], regex: list[dict[str, Any]]) -> list[dict[str, Any]]:
        seen = set()
        merged: list[dict[str, Any]] = []

        for e in ml:
            key = (str(e.get("name", "")).lower(), str(e.get("type", "")))
            if key not in seen:
                seen.add(key)
                merged.append(e)

        for e in regex:
            key = (str(e.get("name", "")).lower(), str(e.get("type", "")))
            if key not in seen:
                seen.add(key)
                merged.append(e)

        return merged
