# ---------------------------------------------------------------------------
# Standard
# ---------------------------------------------------------------------------

from typing import TYPE_CHECKING, List
from .imports import (
    os,
    re,
    quote,
    Counter,
    get_torch,
    get_sentence_transformers,
    get_spacy,
    get_keybert

)

from .config import DEFAULT_PATHS

# ---------------------------------------------------------------------------
# Optional typing (no runtime import)
# ---------------------------------------------------------------------------

if TYPE_CHECKING:
    from sentence_transformers import SentenceTransformer
    import torch


# ---------------------------------------------------------------------------
# CONFIG
# ---------------------------------------------------------------------------

DEFAULT_KEYBERT_PATH = DEFAULT_PATHS["keybert"]




# ---------------------------------------------------------------------------
# SENTENCE-BERT LOADER
# ---------------------------------------------------------------------------

def load_sentence_bert_model(model_path: str = None):

    st = get_sentence_transformers()

    path = model_path or DEFAULT_KEYBERT_PATH

    word_embedding_model = get_sentence_transformers('models').Transformer(
        model_name_or_path=path,
        max_seq_length=256,
        do_lower_case=False
    )

    pooling_model = get_sentence_transformers('models').Pooling(
        word_embedding_model.get_word_embedding_dimension(),
        pooling_mode="mean"
    )

    normalize_model = get_sentence_transformers('models').Normalize()

    return st.SentenceTransformer(
        modules=[word_embedding_model, pooling_model, normalize_model]
    )


# ---------------------------------------------------------------------------
# NLP MANAGER (spaCy)
# ---------------------------------------------------------------------------

class nlpManager:

    _instance = None

    def __new__(cls):

        if cls._instance is None:

            spacy = get_spacy()

            obj = super().__new__(cls)

            obj.spacy = spacy
            obj.nlp = spacy.load("en_core_web_sm")

            cls._instance = obj

        return cls._instance


# ---------------------------------------------------------------------------
# KEYBERT MANAGER
# ---------------------------------------------------------------------------

class KeyBERTManager:

    _instance = None

    def __new__(cls, model_path=None):

        if cls._instance is None:

            obj = super().__new__(cls)

            sbert = load_sentence_bert_model(model_path)

            KeyBERT = get_keybert()

            obj._sbert = sbert
            obj._keybert = KeyBERT(sbert)

            cls._instance = obj

        return cls._instance

    @property
    def sbert(self):
        return self._sbert

    @property
    def keybert(self):
        return self._keybert


# ---------------------------------------------------------------------------
# MODEL ACCESSORS
# ---------------------------------------------------------------------------

def get_keybert_model(model_path: str = None):

    manager = KeyBERTManager(model_path=model_path)

    return manager.sbert


def get_keybert():

    return KeyBERTManager().keybert


# ---------------------------------------------------------------------------
# SENTENCE ENCODING
# ---------------------------------------------------------------------------

def encode_sentences(
    model=None,
    sentences: list[str] = None,
    model_path: str = None
):

    m = model or get_keybert_model(model_path=model_path)

    return m.encode(
        sentences,
        convert_to_tensor=True,
        show_progress_bar=True,
        normalize_embeddings=True
    )


# ---------------------------------------------------------------------------
# COSINE SIMILARITY
# ---------------------------------------------------------------------------

def compute_cosine_similarity(embeddings):

    return get_sentence_transformers('cos_sim')(embeddings, embeddings)


# ---------------------------------------------------------------------------
# KEYBERT EXTRACTION
# ---------------------------------------------------------------------------

def extract_keywords(
    text: str | list[str] = None,
    top_n: int = 5,
    diversity: float = 0.7,
    use_mmr: bool = True,
    stop_words: str = "english",
    keyphrase_ngram_range: tuple[int, int] = (1, 2),
    model: SentenceTransformer = None,
    model_path: str = None
) -> list[tuple[str, float]] | list[list[tuple[str, float]]]:

    if text is None or (isinstance(text, (str, list)) and not text):
        raise ValueError("No content provided for keyword extraction.")

    kw = get_keybert()

    docs = text if isinstance(text, list) else [text]

    return kw.extract_keywords(
        docs,
        keyphrase_ngram_range=keyphrase_ngram_range,
        stop_words=stop_words,
        top_n=top_n,
        use_mmr=use_mmr,
        diversity=diversity
    )


# ---------------------------------------------------------------------------
# RULE-BASED KEYWORDS (spaCy)
# ---------------------------------------------------------------------------

def extract_keywords_nlp(text: str, top_n: int = 5):

    if not isinstance(text, str):
        raise ValueError(f"extract_keywords_nlp expects a string, got {type(text)}")

    nlp_mgr = nlpManager()

    doc = nlp_mgr.nlp(text)

    word_counts = Counter(
        token.text.lower()
        for token in doc
        if token.pos_ in {"NOUN", "PROPN"} and not token.is_stop and len(token.text) > 3
    )

    entity_counts = Counter(
        ent.text.lower()
        for ent in doc.ents
        if len(ent.text.split()) >= 2 and ent.label_ in {"PERSON", "ORG", "GPE", "EVENT"}
    )

    combined = entity_counts + word_counts

    return [kw for kw, _ in combined.most_common(top_n)]


# ---------------------------------------------------------------------------
# KEYWORD DENSITY
# ---------------------------------------------------------------------------

def calculate_keyword_density(text: str, keywords: list[str]):

    if not text:
        return {kw: 0.0 for kw in keywords}

    words = [
        word.strip(".,!?;:()\"'").lower()
        for word in re.split(r"\s+", text)
        if word.strip()
    ]

    total_words = len(words)

    if total_words == 0:
        return {kw: 0.0 for kw in keywords}

    density = {}

    for kw in keywords:

        count = words.count(kw.lower())

        density[kw] = (count / total_words) * 100

    return density


# ---------------------------------------------------------------------------
# COMBINED KEYWORD PIPELINE
# ---------------------------------------------------------------------------

def refine_keywords(
    full_text,
    top_n=10,
    use_mmr=True,
    diversity=0.5,
    keyphrase_ngram_range=(1, 3),
    stop_words="english",
    model=None,
    model_path=None,
    info_data=None
):

    if info_data is None:
        info_data = {}

    nlp_kws = extract_keywords_nlp(full_text, top_n=top_n)

    info_data["keywords_nlp"] = nlp_kws

    keybert_kws = extract_keywords(
        text=full_text,
        top_n=top_n,
        diversity=diversity,
        use_mmr=use_mmr,
        stop_words=stop_words,
        keyphrase_ngram_range=keyphrase_ngram_range,
        model=model,
        model_path=model_path
    )

    keybert_phrases = [phrase for phrase, _ in keybert_kws]

    info_data["keywords_keybert"] = keybert_kws

    merged_set = set(
        [kw.lower() for kw in nlp_kws] +
        [kp.lower() for kp in keybert_phrases]
    )

    combined = list(merged_set)[:top_n]

    info_data["combined_keywords"] = combined

    densities = calculate_keyword_density(full_text, combined)

    info_data["keyword_density"] = densities

    return info_data
