import os
import re
from typing import List, Set

class ArabicPreprocessor:
    """
    Arabic NLP Preprocessor with diacritics removal, letter normalization,
    kashida removal, light stemming, and stop-word filtering.
    """

    TASHKEEL_REGEX = re.compile(r'[\u0617-\u061A\u064B-\u0652]')
    TATWEEL_REGEX = re.compile(r'\u0640')
    PUNCTUATION_REGEX = re.compile(r'[،؛؟«»""\'\(\)\[\]\{\}\.,:!؟\-_/\\|*~`@#$%^&+=<>]')
    NON_ARABIC_REGEX = re.compile(r'[^\u0600-\u06FF\s]')

    # Affixes for Light Stemmer
    PREFIXES = ['كال', 'فال', 'بال', 'وال', 'ال', 'لل', 'فب', 'وب', 'وس', 'فس', 'ف', 'و', 'ب', 'ك', 'ل', 'س', 'ي', 'ت', 'ن', 'ا']
    SUFFIXES = ['كما', 'هما', 'هم', 'هن', 'كم', 'كن', 'نا', 'ها', 'ان', 'ين', 'ون', 'ات', 'تم', 'تن', 'وا', 'تا', 'تي', 'ني', 'ه', 'ي', 'ك', 'ت', 'ة']

    def __init__(self, stopwords_file: str = None):
        self.stopwords: Set[str] = set()
        if stopwords_file and os.path.exists(stopwords_file):
            self.load_stopwords(stopwords_file)
        else:
            default_path = os.path.join(os.path.dirname(__file__), '..', 'data', 'arabic_stopwords.txt')
            if os.path.exists(default_path):
                self.load_stopwords(default_path)

    def load_stopwords(self, filepath: str):
        with open(filepath, 'r', encoding='utf-8') as f:
            for line in f:
                word = line.strip()
                if word:
                    self.stopwords.add(self.normalize_letters(self.remove_tashkeel(word)))

    def remove_tashkeel(self, text: str) -> str:
        """Removes Arabic vowel marks (Harakat/Tashkeel)."""
        return self.TASHKEEL_REGEX.sub('', text)

    def remove_tatweel(self, text: str) -> str:
        """Removes elongation character (Kashida/Tatweel)."""
        return self.TATWEEL_REGEX.sub('', text)

    def normalize_letters(self, text: str) -> str:
        """Standardizes varied Arabic letter forms."""
        # Normalize Alefs
        text = re.sub(r'[إأآٱ]', 'ا', text)
        # Normalize Yaa / Alef Maksura
        text = re.sub(r'ى', 'ي', text)
        # Normalize Taa Marbuta
        text = re.sub(r'ة', 'ه', text)
        # Normalize Hamza variants
        text = re.sub(r'ؤ', 'ء', text)
        text = re.sub(r'ئ', 'ء', text)
        return text

    def clean_punctuation(self, text: str) -> str:
        """Removes punctuation symbols while keeping spaces."""
        return self.PUNCTUATION_REGEX.sub(' ', text)

    def light_stem(self, word: str) -> str:
        """
        Light stemming for Arabic words by systematically stripping
        common prefixes and suffixes while protecting root lengths.
        """
        if len(word) <= 3:
            return word

        # Strip prefixes
        for prefix in sorted(self.PREFIXES, key=len, reverse=True):
            if word.startswith(prefix) and len(word) - len(prefix) >= 3:
                word = word[len(prefix):]
                break

        # Strip suffixes
        for suffix in sorted(self.SUFFIXES, key=len, reverse=True):
            if word.endswith(suffix) and len(word) - len(suffix) >= 3:
                word = word[:-len(suffix)]
                break

        return word

    def preprocess_word(self, word: str, stem: bool = False) -> str:
        """Cleans and standardizes a single Arabic word."""
        word = self.remove_tashkeel(word)
        word = self.remove_tatweel(word)
        word = self.normalize_letters(word)
        word = self.clean_punctuation(word).strip()
        if stem:
            word = self.light_stem(word)
        return word

    def preprocess_text(self, text: str, remove_stopwords: bool = True, stem: bool = False) -> str:
        """Full pipeline for cleaning and standardizing raw Arabic text."""
        if not text:
            return ""
            
        text = self.remove_tashkeel(text)
        text = self.remove_tatweel(text)
        text = self.normalize_letters(text)
        text = self.clean_punctuation(text)
        
        words = text.split()
        processed_words = []
        for w in words:
            w_norm = self.normalize_letters(w)
            if remove_stopwords and (w in self.stopwords or w_norm in self.stopwords):
                continue
            if stem:
                w = self.light_stem(w)
            if len(w) > 0:
                processed_words.append(w)
                
        return ' '.join(processed_words)

    def tokenize_words(self, text: str, remove_stopwords: bool = False, stem: bool = False) -> List[str]:
        """Tokenizes text into a list of cleaned words."""
        cleaned = self.preprocess_text(text, remove_stopwords=remove_stopwords, stem=stem)
        return cleaned.split()