import re
import nltk
from nltk.corpus import stopwords, wordnet
from nltk.tokenize import word_tokenize, sent_tokenize
from nltk.stem import WordNetLemmatizer
import pandas as pd
from dateutil import parser
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# Download required NLTK data
nltk.download('stopwords')
nltk.download('punkt')
nltk.download('punkt_tab')
nltk.download('wordnet')
nltk.download('omw-1.4')
nltk.download('averaged_perceptron_tagger_eng')


def get_wordnet_pos(word):
    """Map POS tag to WordNet format"""
    tag = nltk.pos_tag([word])[0][1][0].upper()
    tag_dict = {"J": wordnet.ADJ,
                "N": wordnet.NOUN,
                "V": wordnet.VERB,
                "R": wordnet.ADV}
    return tag_dict.get(tag, wordnet.NOUN)


def preprocess_and_lemmatize(text):
    """
    Process a single text string and return extracted/cleaned components.
    Returns: (lemmatized_document, links, hashtags, mentions, parsed_dates, clean_text)
    """
    if not isinstance(text, str):
        return [], [], [], [], [], ""

    links, hashtags, mentions, dates = [], [], [], []
    processed_text = text.lower()
    processed_text = processed_text.encode('ascii', 'ignore').decode('ascii')

    # Extract links, hashtags, mentions
    links     = re.findall(r'http\S+', processed_text)
    hashtags  = re.findall(r'#\S+', processed_text)
    mentions  = re.findall(r'@\S+', processed_text)

    # Remove links and hashtags from text
    processed_text = re.sub(r'http\S+', '', processed_text)
    processed_text = re.sub(r'#\S+',    '', processed_text)

    # Extract and remove dates
    date_patterns = [
        r'\d{1,2}[-/]\d{1,2}[-/]\d{2,4}',
        r'\d{4}[-/]\d{1,2}[-/]\d{1,2}',
        r'\b\d{1,2}\s+(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\s+\d{4}\b',
        r'\b(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*\s+\d{1,2},?\s+\d{4}\b',
    ]
    for pattern in date_patterns:
        matches = re.findall(pattern, text, re.IGNORECASE)
        dates.extend(matches)
        processed_text = re.sub(pattern, ' ', processed_text, flags=re.IGNORECASE)

    # Parse extracted date strings into datetime objects
    parsed_dates = []
    for date_str in dates:
        try:
            parsed_dates.append(parser.parse(date_str, fuzzy=True))
        except Exception:
            pass

    # Remove punctuation, digits, and extra whitespace
    processed_text = re.sub(r'[^a-zA-Z\s]', ' ', processed_text)
    processed_text = re.sub(r'\s+', ' ', processed_text).strip()

    # Remove stopwords
    stop_words = set(stopwords.words('english'))
    processed_text = ' '.join(
        word for word in processed_text.split() if word not in stop_words
    )

    # Tokenize sentences → lemmatize words
    lemmatizer = WordNetLemmatizer()
    lemmatized_document = []
    for sentence in sent_tokenize(processed_text):
        words = word_tokenize(sentence)
        lemmatized_sentence = [
            lemmatizer.lemmatize(word, get_wordnet_pos(word))
            for word in words if word.strip()
        ]
        if lemmatized_sentence:
            lemmatized_document.append(lemmatized_sentence)

    # Flatten lemmatized tokens into a single clean string
    clean_text = " ".join(token for sentence in lemmatized_document for token in sentence)

    return lemmatized_document, links, hashtags, mentions, parsed_dates, clean_text


def preprocess_dataframe(df: pd.DataFrame, text_column: str) -> pd.DataFrame:
    """
    Apply NLP preprocessing to a DataFrame text column and return an enriched DataFrame.

    Parameters
    ----------
    df          : Input DataFrame containing the text column.
    text_column : Name of the column holding raw text to process.

    Returns
    -------
    A copy of the original DataFrame with six new columns appended:
        - lemmatized_tokens : list[list[str]]  — sentences of lemmatized words
        - links             : list[str]        — extracted URLs
        - hashtags          : list[str]        — extracted hashtags
        - mentions          : list[str]        — extracted @mentions
        - parsed_dates      : list[datetime]   — parsed datetime objects
        - clean_text        : str              — flat, space-joined lemmatized string
    """
    result_df = df.copy()

    processed = result_df[text_column].apply(preprocess_and_lemmatize)

    result_df['lemmatized_tokens'] = processed.apply(lambda x: x[0])
    result_df['links']             = processed.apply(lambda x: x[1])
    result_df['hashtags']          = processed.apply(lambda x: x[2])
    result_df['mentions']          = processed.apply(lambda x: x[3])
    result_df['parsed_dates']      = processed.apply(lambda x: x[4])
    result_df['clean_text']        = processed.apply(lambda x: x[5])

    return result_df