"""
Text Cleaning Module
====================

This module provides text cleaning and preprocessing utilities
for the Sentiment140 dataset.
"""

import re
import html
import string
import logging
from typing import List
from unidecode import unidecode
import emoji
from sklearn.base import BaseEstimator, TransformerMixin

logger = logging.getLogger(__name__)


# Regular expressions for text cleaning
URL_RE = re.compile(r"http\S+|www\.\S+")
USER_RE = re.compile(r"@\w+")
HASHTAG_RE = re.compile(r"#(\w+)")
RT_RE = re.compile(r"\brt\b")
MULTISPACE_RE = re.compile(r"\s+")
HTML_TAG_RE = re.compile(r"<.*?>")


def split_camelcase(token: str) -> str:
    """Split camelCase tokens into separate words."""
    return re.sub(r"([a-z])([A-Z])", r"\1 \2", token)


def normalize_elongations(word: str) -> str:
    """Normalize repeated characters (e.g., 'cooool' -> 'cool')."""
    return re.sub(r"(.)\1{2,}", r"\1\1", word)


def demojize_text(text: str) -> str:
    """Convert emojis to text representation."""
    return emoji.demojize(text, language="es")


def clean_hashtags(text: str) -> str:
    """Clean hashtags by splitting camelCase and removing # symbol."""
    def repl(m):
        token = m.group(1)
        token = split_camelcase(token)
        return " " + token + " "
    return HASHTAG_RE.sub(repl, text)


def basic_clean(text: str) -> str:
    """
    Apply basic text cleaning operations.
    
    Args:
        text: Raw text to clean
        
    Returns:
        Cleaned text
    """
    if not isinstance(text, str):
        return ""
    
    # Decode HTML entities
    text = html.unescape(text)
    
    # Convert to lowercase
    text = text.lower()
    
    # Remove RT (retweet) markers
    text = RT_RE.sub(" ", text)
    
    # Remove URLs
    text = URL_RE.sub(" ", text)
    
    # Remove user mentions
    text = USER_RE.sub(" ", text)
    
    # Clean hashtags
    text = clean_hashtags(text)
    
    # Remove HTML tags
    text = HTML_TAG_RE.sub(" ", text)
    
    # Convert emojis to text
    text = demojize_text(text)
    
    # Remove accents
    text = unidecode(text)
    
    # Normalize elongations
    text = " ".join(normalize_elongations(w) for w in text.split())
    
    # Remove punctuation (keep apostrophes)
    punct_table = str.maketrans("", "", string.punctuation.replace("'", ""))
    text = text.translate(punct_table)
    
    # Normalize whitespace
    text = MULTISPACE_RE.sub(" ", text).strip()
    
    return text


class TextCleaner(BaseEstimator, TransformerMixin):
    """
    Scikit-learn compatible text cleaner transformer.
    
    This transformer applies comprehensive text cleaning operations
    and optional lemmatization using spaCy.
    """
    
    def __init__(self, use_spacy_lemma: bool = False):
        """
        Initialize the text cleaner.
        
        Args:
            use_spacy_lemma: Whether to use spaCy for lemmatization
        """
        self.use_spacy_lemma = use_spacy_lemma
        self.nlp = None
        
        if self.use_spacy_lemma:
            try:
                import spacy
                self.nlp = spacy.load(
                    "en_core_web_sm",
                    disable=["parser", "ner", "textcat"]
                )
                logger.info("spaCy loaded successfully for lemmatization")
            except Exception as e:
                logger.warning(f"spaCy not available: {e}. Continuing without lemmatization.")
                self.use_spacy_lemma = False
    
    def fit(self, X, y=None):
        """Fit method (does nothing, required by sklearn interface)."""
        return self
    
    def transform(self, X: List[str]) -> List[str]:
        """
        Transform a list of texts by cleaning them.
        
        Args:
            X: List of raw texts
            
        Returns:
            List of cleaned texts
        """
        logger.info(f"Cleaning {len(X)} texts...")
        
        # Apply basic cleaning
        cleaned = [basic_clean(t) for t in X]
        
        # Apply lemmatization if enabled
        if self.use_spacy_lemma and self.nlp is not None:
            logger.info("Applying lemmatization...")
            docs = self.nlp.pipe(cleaned, batch_size=512)
            cleaned = [
                " ".join(t.lemma_ for t in d if not t.is_space)
                for d in docs
            ]
        
        logger.info("Text cleaning completed")
        return cleaned


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # Test text cleaning
    test_texts = [
        "@user This is a #GreatProduct!!! http://example.com 😊",
        "RT @user: I looooove this!! <html>test</html>",
        "#CamelCaseHashtag is coooool"
    ]
    
    cleaner = TextCleaner(use_spacy_lemma=False)
    cleaned = cleaner.transform(test_texts)
    
    for original, cleaned_text in zip(test_texts, cleaned):
        print(f"Original: {original}")
        print(f"Cleaned:  {cleaned_text}")
        print()
