import itertools
import math
import re
from typing import Optional, List, Tuple

from iso639.exceptions import InvalidLanguageValue
from pydantic import BaseModel, ConfigDict, Field

from presidio_analyzer import AnalyzerEngine, RecognizerRegistry, AnalysisExplanation, \
    EntityRecognizer, RecognizerResult
from presidio_analyzer.nlp_engine import NlpArtifacts, NlpEngineProvider
from presidio_analyzer.predefined_recognizers import SpacyRecognizer, PhoneRecognizer
import streamlit as st
import nltk
from flair.data import Sentence
from flair.models import SequenceTagger

from sumy.parsers.plaintext import PlaintextParser
from sumy.nlp.tokenizers import Tokenizer
from sumy.nlp.stemmers import Stemmer
from sumy.summarizers.lsa import LsaSummarizer
from sumy.summarizers.luhn import LuhnSummarizer
from sumy.utils import get_stop_words

from .language import to_name, SPACY_MODELS

BASE_ALLOWED_LABELS = ["PERSON", "ORG", "LOC", "DATE", "PHONE", "IP_ADDRESS", "EMAIL", "URL",
                       "CRYPTO", "IBAN", "CREDIT_CARD", "US_SSN", "US_DRIVER_LICENSE", "US_PASSPORT", "MEDICAL_LICENSE"]

PRESIDIO_TO_BASE_ALIASES = {
    "PHONE_NUMBER": "PHONE",
    "EMAIL_ADDRESS": "EMAIL",
    "IBAN_CODE": "IBAN",
    "DRIVER_LICENSE": "US_DRIVER_LICENSE",
    "US_DRIVER_LICENSE": "US_DRIVER_LICENSE",
    "US_DRIVERS_LICENSE": "US_DRIVER_LICENSE",
    "PASSPORT": "US_PASSPORT",
    "CREDIT_CARD": "CREDIT_CARD",
    "URL": "URL",
    "IP_ADDRESS": "IP_ADDRESS",
    "CRYPTO": "CRYPTO",
    "CRYPTO_WALLET": "CRYPTO",
    "CRYPTO_WALLET_ADDRESS": "CRYPTO",
    "DATE_TIME": "DATE",
    "LOCATION": "LOC",
    "ORGANIZATION": "ORG",
    "GPE": "LOC",
    "NORP": "ORG",
}

BASE_TO_RECOGNIZER_EXPANSIONS = {
    "ORG": ["ORG", "ORGANIZATION", "NORP"],
    "LOC": ["LOC", "LOCATION", "GPE"],
    "PHONE": ["PHONE", "PHONE_NUMBER"],
    "EMAIL": ["EMAIL", "EMAIL_ADDRESS"],
    "IBAN": ["IBAN", "IBAN_CODE"],
    "US_DRIVER_LICENSE": ["US_DRIVER_LICENSE", "US_DRIVERS_LICENSE", "DRIVER_LICENSE"],
    "US_PASSPORT": ["US_PASSPORT", "PASSPORT"],
    "DATE": ["DATE", "DATE_TIME"],
    "PERSON": ["PERSON"],
    "URL": ["URL"],
    "IP_ADDRESS": ["IP_ADDRESS"],
    "CRYPTO": ["CRYPTO", "CRYPTO_WALLET", "CRYPTO_WALLET_ADDRESS"],
    "CREDIT_CARD": ["CREDIT_CARD"],
    "US_SSN": ["US_SSN"],
    "MEDICAL_LICENSE": ["MEDICAL_LICENSE"],
}


class FlairRecognizer(EntityRecognizer):
    ENTITIES = [
        "LOC",
        "PERSON",
        "ORG",
    ]

    DEFAULT_EXPLANATION = "Identified as {} by Flair's Named Entity Recognition"

    CHECK_LABEL_GROUPS = [
        ({"LOC"}, {"LOC", "LOCATION"}),
        ({"PERSON"}, {"PER", "PERSON"}),
        ({"ORG"}, {"ORG", "ORGANIZATION"}),
    ]

    PRESIDIO_EQUIVALENCES = {
        "PER": "PERSON",
        "LOC": "LOC",
        "ORG": "ORG"
    }

    def __init__(
            self,
            model: SequenceTagger = None,
            supported_language: str = "en",
            supported_entities: Optional[List[str]] = None,
            check_label_groups: Optional[Tuple[set, set]] = None,
    ):
        self.check_label_groups = (
            check_label_groups if check_label_groups else self.CHECK_LABEL_GROUPS
        )

        supported_entities = supported_entities if supported_entities else self.ENTITIES
        self.model = model

        super().__init__(
            supported_entities=supported_entities,
            supported_language=supported_language,
            name="Flair Analytics",
        )

    def load(self) -> None:
        pass

    def get_supported_entities(self) -> List[str]:
        return self.supported_entities

    def analyze(self, text: str, entities: List[str], nlp_artifacts: NlpArtifacts = None) -> List[RecognizerResult]:
        results = []

        sentences = Sentence(text)
        self.model.predict(sentences)

        if not entities:
            entities = self.supported_entities

        for entity in entities:
            if entity not in self.supported_entities:
                continue

            for ent in sentences.get_spans("ner"):
                if not self.__check_label(
                        entity, ent.labels[0].value, self.check_label_groups
                ):
                    continue
                textual_explanation = self.DEFAULT_EXPLANATION.format(
                    ent.labels[0].value
                )
                explanation = self.build_flair_explanation(
                    round(ent.score, 2), textual_explanation
                )
                flair_result = self._convert_to_recognizer_result(ent, explanation)

                results.append(flair_result)

        return results

    def build_flair_explanation(self, original_score: float, explanation: str) -> AnalysisExplanation:
        explanation = AnalysisExplanation(
            recognizer=self.__class__.__name__,
            original_score=original_score,
            textual_explanation=explanation,
        )
        return explanation

    def _convert_to_recognizer_result(self, entity, explanation) -> RecognizerResult:
        entity_type = self.PRESIDIO_EQUIVALENCES.get(entity.tag, entity.tag)
        flair_score = round(entity.score, 2)

        flair_results = RecognizerResult(
            entity_type=entity_type,
            start=entity.start_position,
            end=entity.end_position,
            score=flair_score,
            analysis_explanation=explanation,
        )

        return flair_results

    @staticmethod
    def __check_label(
            entity: str, label: str, check_label_groups: Tuple[set, set]
    ) -> bool:
        return any(
            [entity in egrp and label in lgrp for egrp, lgrp in check_label_groups]
        )


def normalize_label(label: str) -> str:
    return PRESIDIO_TO_BASE_ALIASES.get(label, label)


def expand_entities_for_analyzer(entities_list):
    expanded = set()
    for e in entities_list:
        vals = BASE_TO_RECOGNIZER_EXPANSIONS.get(e, [e])
        for v in vals:
            expanded.add(v)
    return list(expanded)


def _sumy__get_best_sentences(sentences, rating, *args, **kwargs):
    from operator import attrgetter
    from sumy.summarizers._summarizer import SentenceInfo

    rate = rating
    if isinstance(rating, dict):
        assert not args and not kwargs
        rate = lambda s: rating[s]
    infos = (SentenceInfo(s, o, rate(s, *args, **kwargs)) for o, s in enumerate(sentences))
    infos = sorted(infos, key=attrgetter("rating"), reverse=True)
    return tuple((i.sentence, i.rating, i.order) for i in infos)


def _sumy__lsa_call(summarizer, document):
    summarizer._ensure_dependecies_installed()
    dictionary = summarizer._create_dictionary(document)
    if not dictionary:
        return ()
    matrix = summarizer._create_matrix(document, dictionary)
    matrix = summarizer._compute_term_frequency(matrix)
    from numpy.linalg import svd as singular_value_decomposition

    u, sigma, v = singular_value_decomposition(matrix, full_matrices=False)
    ranks = iter(summarizer._compute_ranks(sigma, v))
    return _sumy__get_best_sentences(document.sentences, lambda s: next(ranks))


def _sumy__luhn_call(summarizer, document):
    words = summarizer._get_significant_words(document.words)
    return _sumy__get_best_sentences(document.sentences, summarizer.rate_sentence, words)


def get_nltk_tokenizer(language: str) -> Tokenizer:
    try:
        nltk_lang = to_name(language).lower()

    except InvalidLanguageValue:
        nltk_lang = language

    try:
        nltk.data.find("tokenizers/punkt")
    except LookupError:
        nltk.download("punkt")

    return Tokenizer(nltk_lang)


class NERObject(BaseModel):
    name: str
    label: str
    score: float = 0.0
    start: int
    count: int
    context: str | None = None
    comentions: list[str] = Field(default_factory=list)
    model_config = ConfigDict(extra="allow")

    def __repr__(self):
        return f"NERObject(label={self.label},name={self.name})"


def postprocess_ner(entities: list[NERObject], whitelisted_labels=None, max_entities=None):
    if whitelisted_labels is not None:
        entities = [e for e in entities if e.label in whitelisted_labels]
    entities = sorted(entities, key=lambda x: x.name)
    final_entities = []
    for _, group in itertools.groupby(entities, key=lambda x: x.name):
        group = list(group)
        best_entity = max(group, key=lambda x: x.score * x.count)
        merged_data = {
            "name": best_entity.name,
            "label": best_entity.label,
            "score": best_entity.score,
            "context": best_entity.context,
            "count": sum(e.count for e in group),
            "start": best_entity.start,
        }
        all_fields = type(best_entity).model_fields.keys()
        for field in all_fields:
            if field in merged_data:
                continue
            values = [getattr(e, field, None) for e in group if getattr(e, field, None) is not None]
            if not values:
                continue
            if isinstance(values[0], list):
                merged_data[field] = list(set(itertools.chain.from_iterable(values or [])))
            else:
                merged_data[field] = getattr(best_entity, field, None)
        final_entities.append(NERObject(**merged_data))
    final_entities = sorted(final_entities, key=lambda x: x.score * x.count, reverse=True)
    if max_entities and len(final_entities) > max_entities:
        final_entities = final_entities[:max_entities]
    return final_entities


# Pattern-based recognizers that work in any language (regex-based)
REGEX_RECOGNIZER_NAMES = [
    "EmailRecognizer",
    "IbanRecognizer", 
    "CreditCardRecognizer",
    "PhoneRecognizer",
    "IpRecognizer",
    "UrlRecognizer",
    "CryptoRecognizer",
    "UsSsnRecognizer",
    "UsLicenseRecognizer",
    "UsPassportRecognizer",
    "MedicalLicenseRecognizer",
    "DateRecognizer",
    "UsItinRecognizer",
    "SgFinRecognizer",
    "AuAbnRecognizer",
    "AuAcnRecognizer",
    "AuTfnRecognizer",
    "AuMedicareRecognizer",
    "InPanRecognizer",
]

# Entities that regex recognizers handle
REGEX_ENTITIES = [
    "PHONE", "EMAIL", "IBAN", "CREDIT_CARD", "URL", "IP_ADDRESS",
    "CRYPTO", "US_SSN", "US_DRIVER_LICENSE", "US_PASSPORT", 
    "MEDICAL_LICENSE", "DATE",
]

# Entities that NER models handle
NER_ENTITIES = ["PERSON", "ORG", "LOC"]


def build_presidio_analyzer(language: str, engine_type: str = "spacy", model=None) -> AnalyzerEngine:
    """
    Build a Presidio analyzer with EXCLUSIVE engine types.
    
    Args:
        language: Language code (e.g., "en", "de", "fr", "it")
        engine_type: One of:
            - "regex": Only pattern-based recognizers (phones, emails, IBANs, etc.)
            - "flair": Only Flair NER (PERSON, ORG, LOC)
            - "spacy": Only spaCy NER (PERSON, ORG, LOC, etc.)
        model: For "flair" - Flair SequenceTagger or model name
               For "spacy" - spaCy model name
               For "regex" - ignored (can be None)
    
    Returns:
        AnalyzerEngine configured for the specified engine type only
    """
    registry = RecognizerRegistry()
    
    if engine_type == "regex":
        default_registry = RecognizerRegistry()
        default_registry.load_predefined_recognizers(languages=["en"])
        for recognizer in default_registry.recognizers:
            if recognizer.name in REGEX_RECOGNIZER_NAMES:
                # Special handling for PhoneRecognizer - use leniency=0 for more permissive matching
                if recognizer.name == "PhoneRecognizer":
                    phone_recognizer = PhoneRecognizer(
                        supported_language=language,
                        leniency=0
                    )
                    registry.add_recognizer(phone_recognizer)
                else:
                    recognizer.supported_language = language
                    if hasattr(recognizer, 'supported_languages'):
                        recognizer.supported_languages = [language]
                    registry.add_recognizer(recognizer)
        registry.supported_languages = [language]
        spacy_model = SPACY_MODELS.get(language)
        if spacy_model is None:
            spacy_model = SPACY_MODELS.get("en", "en_core_web_sm")
            nlp_lang = "en"
        else:
            nlp_lang = language
        configuration = {
            "nlp_engine_name": "spacy",
            "models": [{"lang_code": nlp_lang, "model_name": spacy_model}],
        }
        provider = NlpEngineProvider(nlp_configuration=configuration)
        nlp_engine = provider.create_engine()
        return AnalyzerEngine(nlp_engine=nlp_engine, registry=registry, supported_languages=[language])
    
    elif engine_type == "flair":
        if model is None:
            raise ValueError("Flair model must be provided")
        if isinstance(model, str):
            flair_model = SequenceTagger.load(model)
        else:
            flair_model = model
        
        flair_recognizer = FlairRecognizer(
            model=flair_model,
            supported_language=language
        )
        registry.add_recognizer(flair_recognizer)
        registry.supported_languages = [language]
        spacy_model = SPACY_MODELS.get(language, SPACY_MODELS.get("en", "en_core_web_sm"))
        configuration = {
            "nlp_engine_name": "spacy",
            "models": [{"lang_code": language, "model_name": spacy_model}],
        }
        provider = NlpEngineProvider(nlp_configuration=configuration)
        nlp_engine = provider.create_engine()
        return AnalyzerEngine(nlp_engine=nlp_engine, registry=registry, supported_languages=[language])
    
    elif engine_type == "spacy":
        if model is None:
            raise ValueError("SpaCy model name must be provided")
        spacy_recognizer = SpacyRecognizer(supported_language=language)
        registry.add_recognizer(spacy_recognizer)
        registry.supported_languages = [language]
        configuration = {
            "nlp_engine_name": "spacy",
            "models": [{"lang_code": language, "model_name": model}],
        }
        provider = NlpEngineProvider(nlp_configuration=configuration)
        nlp_engine = provider.create_engine()
        return AnalyzerEngine(nlp_engine=nlp_engine, registry=registry, supported_languages=[language])
    else:
        raise ValueError(f"Unknown engine_type: {engine_type}. Must be 'regex', 'flair', or 'spacy'")


def compute_ner_presidio(
        text,
        language,
        analyzer,
        engine_type="spacy",
        entities=None,
        score_threshold=0.5,
        context_width=150,
        with_comentions=True,
        with_context=True,
        batch_size=32,
        n_process=4
):
    # Prevent CUDA fork issues
    if engine_type == "flair" and n_process > 1:
        raise ValueError("n_process must be 1 for Flair")
    
    # Single text string processing
    results = analyzer.analyze(
        text=text,
        language=language,
        entities=(expand_entities_for_analyzer(entities) if entities else None)
    )
    ner_objects = []
    for result in results:
        if result.score >= score_threshold:
            context_start = max(0, result.start - math.floor(context_width / 2))
            context_end = min(len(text), result.end + math.ceil(context_width / 2))
            context = text[context_start:context_end] if with_context else None
            ner_objects.append(NERObject(
                name=text[result.start:result.end],
                label=normalize_label(result.entity_type),
                score=float(result.score),
                start=int(result.start),
                count=1,
                context=context
            ))
    if with_comentions:
        for i in range(len(ner_objects)):
            entity = ner_objects[i]
            comentions = [
                ner_objects[j].name
                for j in range(len(ner_objects))
                if j != i and abs(ner_objects[j].start - entity.start) < math.ceil(context_width / 2)
            ]
            ner_objects[i].comentions = comentions
    return ner_objects


def get_extractive_summary(text, language, max_chars, fast=False, with_scores=False):
    tokenizer = get_nltk_tokenizer(language)
    stemmer = Stemmer(language)
    parser = PlaintextParser.from_string(text, tokenizer)
    if fast:
        summarizer = LuhnSummarizer(stemmer)
        summarizer.stop_words = get_stop_words(language)
        scored_sentences = iter(_sumy__luhn_call(summarizer, parser.document))
    else:
        summarizer = LsaSummarizer(stemmer)
        summarizer.stop_words = get_stop_words(language)
        scored_sentences = iter(_sumy__lsa_call(summarizer, parser.document))
    summary = []
    summary_chars = 0
    summary_chars_penultimate = 0
    while summary_chars < max_chars:
        try:
            next_sentence = next(scored_sentences)
            summary.append(next_sentence)
            summary_chars_penultimate = summary_chars
            summary_chars += len(" " + next_sentence[0]._text)
        except StopIteration:
            break
    summary = sorted(summary, key=lambda x: x[2])
    summary = [(sentence[0]._text, sentence[1]) for sentence in summary]
    if summary_chars > max_chars:
        summary[-1] = (
            summary[-1][0][: max_chars - summary_chars_penultimate],
            summary[-1][1],
        )
    if not with_scores:
        summary = " ".join([s[0] for s in summary])
    else:
        min_score = min([s[1] for s in summary]) if summary else 0
        max_score = max([min_score] + [s[1] for s in summary])
        score_range = 1 if min_score == max_score else (max_score - min_score)
        summary = [(s[0], (s[1] - min_score) / score_range) for s in summary]
    return summary


def _preprocess_newlines_for_ner(text: str) -> str:
    """
    Replace newline-containing whitespace sequences with ' — ' (space + em dash + space).
    
    This helps NER models treat text on separate lines as distinct entities,
    preventing merging of e.g. "Jennifer Williams\\nMegaCorp" into a single entity.
    
    Handles:
    - Single newlines: \\n
    - Carriage returns: \\r\\n, \\r
    - Mixed whitespace with newlines: '  \\n  ', '\\n\\n', etc.
    
    All are collapsed into a single ' — ' separator.
    """
    pattern = r'[ \t]*[\r\n]+[ \t\r\n]*'
    return re.sub(pattern, ' — ', text)


def _strip_honorifics_for_ner(text: str) -> str:
    """
    Remove abbreviated honorifics/titles from text to prevent them from interfering with NER.
    
    Returns:
        Text with honorifics removed, preserving the rest of the content
    """
    honorifics = [
        # English
        r'Dr', r'Mr', r'Mrs', r'Ms', r'Miss', r'Prof', r'Rev', r'Revd',
        r'Hon', r'Sr', r'Jr', r'Esq',
        
        # Italian
        r'Dott', r'Dott\.ssa', r'Dott\.re',
        r'Prof', r'Prof\.ssa',
        r'Sig', r'Sig\.ra', r'Sig\.na',
        r'Avv', r'Ing', r'Arch', r'Geom',
        r'Rag',

        # German
        r'Herr', r'Frau', r'Frl', r'Fräulein',
        r'Dipl\.-Ing', r'Dipl\.-Kfm', r'Dipl',
        
        # French
        r'M', r'Mme', r'Mlle', r'Mgr', r'Me',
        
        # Spanish
        r'D', r'Dª', r'Doña', r'Don',
        r'Dra', r'Sra', r'Srta',
        
        # General academic/professional
        r'Ph\.D', r'M\.D', r'B\.Sc', r'M\.Sc', r'B\.A', r'M\.A',
    ]
    pattern = r'\b(' + '|'.join(honorifics) + r')\.?\s+'
    result = re.sub(pattern, '', text, flags=re.IGNORECASE)
    result = re.sub(r'\s{2,}', ' ', result)
    return result


def ner_pipe(
        text,
        language,
        model,
        engine_type="spacy",
        fast=False,
        compression_ratio="auto",
        with_scores=False,
        with_comentions=True,
        with_context=True,
        entities=None,
        score_threshold=0.5,
        batch_size=32,
        n_process=4,
        preprocess_newlines=True
):
    """
    Run NER pipeline on text.
    
    Args:
        text: Input text (str). For multiple texts, iterate and call this function for each.
        language: Language code (e.g., 'en', 'de', 'fr')
        model: Model name for spacy/flair engine
        engine_type: 'regex', 'flair', or 'spacy'
        fast: Use fast summarization for long texts
        compression_ratio: Compression ratio for long texts ('auto' or float)
        with_scores: Include confidence scores (not implemented)
        with_comentions: Include co-mentioned entities
        with_context: Include surrounding context
        entities: List of entity types to detect (None = all)
        score_threshold: Minimum confidence score
        batch_size: Batch size for processing
        n_process: Number of parallel processes
        preprocess_newlines: Replace newlines with ' — ' to prevent entity merging
    """
    if with_scores:
        raise NotImplementedError("with_scores functionality is not implemented yet")
    
    if not isinstance(text, str):
        raise TypeError(f"text must be str, not {type(text).__name__}")

    analyzer = build_presidio_analyzer(
        language=language,
        engine_type=engine_type,
        model=model,
    )

    if preprocess_newlines:
        text = _preprocess_newlines_for_ner(text)
    text = _strip_honorifics_for_ner(text)

    if compression_ratio == "auto":
        compression_ratio = max(1.0, len(text) / 15000) if fast else 1.0

    if compression_ratio > 1.0:
        sentences = get_extractive_summary(text, language, int(len(text) / compression_ratio), fast=fast,
                                           with_scores=True)
        text = " ".join([s[0] for s in sentences])

    ner = compute_ner_presidio(
        text,
        language,
        analyzer,
        engine_type,
        entities,
        score_threshold,
        with_comentions=with_comentions,
        with_context=with_context,
        batch_size=batch_size,
        n_process=n_process
    )

    return ner


def get_ner_handler(
        language,
        model,
        engine_type="spacy",
        fast=False,
        entities=None,
        score_threshold=0.5,
        batch_size=32,
        n_process=4,
        preprocess_newlines=True
):
    try:
        get_nltk_tokenizer(language)
    except LookupError:
        language = "en"

    return lambda text, compression_ratio="auto", with_scores=False, with_comentions=True, with_context=True: ner_pipe(
        text,
        language,
        model,
        engine_type,
        fast,
        compression_ratio,
        with_scores,
        with_comentions,
        with_context,
        entities,
        score_threshold,
        batch_size,
        n_process,
        preprocess_newlines
    )


@st.cache_resource
def get_cached_ner_handler(language, model):
    return get_ner_handler(language, model)
