import spacy

from analogai.foundation.settings import config


_nlp = None

_PLACEHOLDER_PREFIXES = ("respondent:", "agent:", "company:")
_RESTRICTED_WORDS = {
    "birth",
    "lain",
    "laylie",
    "lie",
}


def _is_placeholder_identifier(text: str) -> bool:
    if not text or not text.strip():
        return False
    lowered = text.strip().lower()
    return any(lowered.startswith(prefix) for prefix in _PLACEHOLDER_PREFIXES)


def _get_nlp():
    global _nlp
    if _nlp is None:
        _nlp = spacy.load("en_core_web_sm")
    return _nlp


def _is_restricted_word(text: str) -> bool:
    if not text or not text.strip():
        return False
    return text.strip().lower() in _RESTRICTED_WORDS


def _is_named_entity_token(token) -> bool:
    return token.ent_iob_ != "O" or bool(token.ent_type_)


def _contains_named_entity(doc) -> bool:
    if not doc:
        return False
    return any(_is_named_entity_token(token) for token in doc)


def _contains_restricted_word(doc) -> bool:
    if not doc:
        return False
    return any(_is_restricted_word(token.text) for token in doc)


def is_likely_proper_noun(token) -> bool:
    if token.pos_ == "PROPN":
        return True
    
    text = token.text
    if text and text[0].isupper() and len(text) > 1:
        return True
    
    return False


def lemmatize_text(text: str, preserve_proper: bool = True) -> str:
    if not text or not text.strip():
        return ""
    if _is_placeholder_identifier(text):
        return text.strip()
    if not config.get("lemmatize", False):
        return text.strip()
    
    nlp = _get_nlp()
    doc = nlp(text.strip())
    result_tokens = []
    
    for token in doc:
        if token.is_punct or token.is_space:
            continue

        if token.text.lower() == "born":
            lemma = "birth"
        elif _is_named_entity_token(token) or _is_restricted_word(token.text):
            lemma = token.text
        elif preserve_proper and is_likely_proper_noun(token):
            lemma = token.text
        else:
            lemma = token.lemma_.lower()
        
        result_tokens.append(lemma)
    
    return " ".join(result_tokens)


if __name__ == "__main__":
    test_cases = [
        "Socrates",
        "iPhone",
        "John",
        "New York",
        "apples",
        "humans",
        "philosophers",
        "born humans",
        "George bought an iPhone",
        "Giorgi was born at 1991",
        "born",
    ]

    for text in test_cases:
        lemmatized = lemmatize_text(text)
        print(f"Original: '{text}' -> Lemmatized: '{lemmatized}'")