from aiforge.core.registry import register
from aiforge.core.context import BaseOperator

@register("tokenize")
class Tokenize(BaseOperator):
    """Tokenize a string of text into words using NLTK."""
    def run(self, data):
        from nltk.tokenize import word_tokenize
        import nltk
        try:
            return word_tokenize(data)
        except LookupError:
            nltk.download('punkt', quiet=True)
            nltk.download('punkt_tab', quiet=True)
            return word_tokenize(data)

@register("lowercase")
class Lowercase(BaseOperator):
    """Convert strings or a list of strings to lowercase."""
    def run(self, data):
        if isinstance(data, list):
            return [str(item).lower() for item in data]
        return str(data).lower()

@register("join_strings")
class JoinStrings(BaseOperator):
    """Join a list of strings into a single text block (Bridge)."""
    def run(self, data, separator=" "):
        if isinstance(data, list):
            return separator.join([str(item) for item in data])
        return str(data)

@register("tfidf_vectorize")
class TfIdfVectorize(BaseOperator):
    """Convert text to numerical features (Bridge: NLP -> ML/DL)."""
    def run(self, data, max_features=100, fit=True, vectorizer=None):
        from sklearn.feature_extraction.text import TfidfVectorizer
        
        if isinstance(data, str):
            data = [data]
            
        if fit or vectorizer is None:
            vectorizer = TfidfVectorizer(max_features=max_features)
            features = vectorizer.fit_transform(data)
            return {"features": features.toarray(), "vectorizer": vectorizer}
        else:
            features = vectorizer.transform(data)
            return {"features": features.toarray(), "vectorizer": vectorizer}

@register("remove_stopwords")
class RemoveStopwords(BaseOperator):
    """Remove common English stopwords from text."""
    def run(self, data):
        from nltk.corpus import stopwords
        from nltk.tokenize import word_tokenize
        import nltk
        try:
            stop_words = set(stopwords.words('english'))
        except LookupError:
            nltk.download('stopwords', quiet=True)
            stop_words = set(stopwords.words('english'))
            
        def clean_text(text):
            try:
                tokens = word_tokenize(text)
            except LookupError:
                nltk.download('punkt', quiet=True)
                nltk.download('punkt_tab', quiet=True)
                tokens = word_tokenize(text)
            return " ".join([w for w in tokens if not w.lower() in stop_words])

        if isinstance(data, list):
            return [clean_text(str(item)) for item in data]
        return clean_text(str(data))

@register("summarize_text")
class SummarizeText(BaseOperator):
    """Summarize text using HuggingFace transformers pipeline."""
    def run(self, data, max_length=130, min_length=30):
        from transformers import pipeline
        # We initialize pipeline locally to avoid loading overhead if not used
        summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
        
        if isinstance(data, list):
            results = summarizer(data, max_length=max_length, min_length=min_length, do_sample=False)
            return [res['summary_text'] for res in results]
        
        result = summarizer(str(data), max_length=max_length, min_length=min_length, do_sample=False)
        return result[0]['summary_text']
