"""
Named Entity Recognition (NER) Module
Extracts and categorizes entities: persons, organizations, dates,
monetary values, and locations.

Uses spaCy (en_core_web_trf or en_core_web_lg) for high accuracy NER.
Falls back to en_core_web_sm if transformer model is unavailable.
Combines regex patterns for money/date extraction to improve recall.
"""
import logging
import re
from collections import defaultdict
from typing import Dict, List

logger = logging.getLogger(__name__)


class EntityExtractor:
    """
    High-precision NER using spaCy with regex augmentation.
    Entity categories:
        - persons       → PERSON
        - organizations → ORG
        - dates         → DATE
        - amounts       → MONEY
        - locations     → GPE, LOC, FAC
    """

    # spaCy label → our category
    LABEL_MAP = {
        "PERSON": "persons",
        "ORG": "organizations",
        "DATE": "dates",
        "MONEY": "amounts",
        "GPE": "locations",
        "LOC": "locations",
        "FAC": "locations",
    }

    # Regex patterns for augmenting spaCy (higher recall for structured patterns)
    MONEY_PATTERNS = [
        r"\$\s*[\d,]+(?:\.\d{1,2})?(?:\s*(?:million|billion|trillion|thousand|M|B|K))?\b",
        r"(?:USD|EUR|GBP|INR|JPY)\s*[\d,]+(?:\.\d{1,2})?(?:\s*(?:million|billion|trillion))?\b",
        r"[\d,]+(?:\.\d{1,2})?\s*(?:dollars?|euros?|pounds?|rupees?|crore|lakh)\b",
    ]

    DATE_PATTERNS = [
        r"\b(?:January|February|March|April|May|June|July|August|September|October|November|December)"
        r"\s+\d{1,2},?\s+\d{4}\b",
        r"\b\d{1,2}[/-]\d{1,2}[/-]\d{2,4}\b",
        r"\b\d{4}[/-]\d{1,2}[/-]\d{1,2}\b",
        r"\b(?:Q[1-4]|FY|FQ)\s*\d{4}\b",
        r"\b\d{4}\b(?=\s*(?:to|through|until|-)\s*\d{4})",  # Year ranges
        r"\b(?:Jan|Feb|Mar|Apr|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\.?\s+\d{1,2},?\s+\d{4}\b",
    ]

    _nlp = None  # Class-level singleton

    def __init__(self):
        self._load_nlp()
        self._compile_patterns()

    @classmethod
    def _load_nlp(cls):
        """Load spaCy NLP model (tries best available model)."""
        if cls._nlp is not None:
            return

        model_preference = [
            "en_core_web_trf",   # Transformer-based (best accuracy)
            "en_core_web_lg",    # Large (good accuracy)
            "en_core_web_md",    # Medium
            "en_core_web_sm",    # Small fallback
        ]

        import spacy
        for model_name in model_preference:
            try:
                cls._nlp = spacy.load(model_name)
                logger.info(f"Loaded spaCy model: {model_name}")
                return
            except OSError:
                continue

        # Last resort: download small model
        logger.warning("No spaCy model found. Downloading en_core_web_sm...")
        try:
            import subprocess, sys
            subprocess.run(
                [sys.executable, "-m", "spacy", "download", "en_core_web_sm"],
                check=True, capture_output=True
            )
            cls._nlp = spacy.load("en_core_web_sm")
        except Exception as e:
            logger.error(f"Failed to load any spaCy model: {e}")
            cls._nlp = None

    def _compile_patterns(self):
        """Pre-compile regex patterns for performance."""
        self._money_re = [re.compile(p, re.IGNORECASE) for p in self.MONEY_PATTERNS]
        self._date_re = [re.compile(p, re.IGNORECASE) for p in self.DATE_PATTERNS]

    def extract(self, text: str) -> Dict[str, List[str]]:
        """
        Extract named entities from text.
        Returns a dict with keys: persons, organizations, dates, amounts, locations.
        """
        result = defaultdict(set)

        # ── spaCy NER ──────────────────────────────────────────────────────────
        if self._nlp and text.strip():
            spacy_entities = self._extract_spacy(text)
            for category, entities in spacy_entities.items():
                result[category].update(entities)

        # ── Regex augmentation ─────────────────────────────────────────────────
        regex_money = self._extract_regex(text, self._money_re)
        regex_dates = self._extract_regex(text, self._date_re)
        result["amounts"].update(regex_money)
        result["dates"].update(regex_dates)

        # ── Deduplicate and sort ───────────────────────────────────────────────
        return {
            "persons":       sorted(self._dedupe(result["persons"])),
            "organizations": sorted(self._dedupe(result["organizations"])),
            "dates":         sorted(self._dedupe(result["dates"])),
            "amounts":       sorted(self._dedupe(result["amounts"])),
            "locations":     sorted(self._dedupe(result["locations"])),
        }

    def _extract_spacy(self, text: str) -> Dict[str, List[str]]:
        """Run spaCy NER on text, handling long documents by chunking."""
        result = defaultdict(list)
        max_chars = 1_000_000  # spaCy default limit

        # Process in chunks for very long documents
        for chunk_start in range(0, len(text), max_chars - 100):
            chunk = text[chunk_start: chunk_start + max_chars]
            try:
                doc = self._nlp(chunk)
                for ent in doc.ents:
                    category = self.LABEL_MAP.get(ent.label_)
                    if category:
                        cleaned = self._clean_entity(ent.text)
                        if cleaned:
                            result[category].append(cleaned)
            except Exception as e:
                logger.warning(f"spaCy processing chunk failed: {e}")

        return result

    def _extract_regex(self, text: str, patterns: List[re.Pattern]) -> List[str]:
        """Extract entities using regex patterns."""
        found = []
        for pattern in patterns:
            matches = pattern.findall(text)
            found.extend(m.strip() for m in matches if m.strip())
        return found

    def _clean_entity(self, text: str) -> str:
        """
        Clean an entity string:
        - Strip whitespace and punctuation artifacts
        - Normalize internal whitespace
        - Skip single-character or pure-number entities for persons/orgs
        """
        text = text.strip()
        # Normalize internal whitespace
        text = re.sub(r"\s+", " ", text)
        # Remove leading/trailing punctuation
        text = text.strip(".,;:\"'()-")
        return text.strip()

    def _dedupe(self, items) -> List[str]:
        """
        Deduplicate entity list.
        Also removes sub-strings that are already covered by longer matches.
        """
        seen = list(dict.fromkeys(item for item in items if item))

        # Remove items that are substrings of other items
        result = []
        for item in seen:
            if not any(
                item.lower() != other.lower() and item.lower() in other.lower()
                for other in seen
            ):
                result.append(item)

        return result