processors/language_detector.py

# DEPENDENCIES
import re
import torch
import string
from typing import Dict
from typing import List
from typing import Tuple
from loguru import logger
from typing import Optional
from config.enums import Script
from dataclasses import dataclass
from config.enums import Language
from config.schemas import LanguageDetectionResult
from config.constants import language_detection_params


# Try to import optional libraries
try:
    import langdetect
    from langdetect import detect, detect_langs, DetectorFactory
    # Seed for reproducibility
    DetectorFactory.seed = 0
    LANGDETECT_AVAILABLE = True

except ImportError:
    logger.warning("langdetect not available. Install: pip install langdetect")
    LANGDETECT_AVAILABLE = False

try:
    from models.model_manager import get_model_manager
    MODEL_MANAGER_AVAILABLE = True
    
except ImportError:
    logger.warning("model_manager not available, using fallback methods")
    MODEL_MANAGER_AVAILABLE = False


class LanguageDetector:
    """
    Detects the language of input text using multiple strategies with fallbacks.

    Features:
    - Learned language representations (when available)
    - Statistical language probability estimation
    - Script and character distribution analysis
    - Multi-signal aggregation with fallbacks

    Supported Languages:
    - Broad multilingual coverage via learned language representations
    - Deterministic support via script and statistical analysis
    """
    # Use constants from config
    MIN_TEXT_LENGTH = language_detection_params.MINIMUM_TEXT_LENGTH
    LANGUAGE_NAMES  = language_detection_params.LANGUAGE_NAMES
    
    # Map Script enum to string keys for SCRIPT_RANGES
    SCRIPT_RANGES   = {Script.LATIN      : language_detection_params.SCRIPT_RANGES["latin"],
                       Script.CYRILLIC   : language_detection_params.SCRIPT_RANGES["cyrillic"],
                       Script.ARABIC     : language_detection_params.SCRIPT_RANGES["arabic"],
                       Script.CHINESE    : language_detection_params.SCRIPT_RANGES["chinese"],
                       Script.JAPANESE   : language_detection_params.SCRIPT_RANGES["japanese"],
                       Script.KOREAN     : language_detection_params.SCRIPT_RANGES["korean"],
                       Script.DEVANAGARI : language_detection_params.SCRIPT_RANGES["devanagari"],
                       Script.GREEK      : language_detection_params.SCRIPT_RANGES["greek"],
                       Script.HEBREW     : language_detection_params.SCRIPT_RANGES["hebrew"],
                       Script.THAI       : language_detection_params.SCRIPT_RANGES["thai"],
                      }
    

    def __init__(self, use_model: bool = True):
        """
        Initialize language detector
        
        Arguments:
        ----------
            use_model : Use ML model for detection (more accurate)
        """
        self.use_model             = use_model and MODEL_MANAGER_AVAILABLE
        self.model_manager         = None
        self.classifier            = None
        self.is_initialized        = False
        
        logger.info(f"LanguageDetector initialized (use_model={self.use_model})")
    

    def initialize(self) -> bool:
        """
        Initialize the ML model (if using)
        
        Returns:
        --------
            { bool } : True if successful, False otherwise
        """
        if not self.use_model:
            self.is_initialized = True
            return True
        
        try:
            logger.info("Initializing language detection model...")
            
            self.model_manager  = get_model_manager()
            self.classifier     = self.model_manager.load_pipeline(model_name = "language_detector",
                                                                   task       = "text-classification",
                                                                  )
            
            self.is_initialized = True
            logger.success("Language detector initialized successfully")
            return True
            
        except Exception as e:
            logger.error(f"Failed to initialize language detector: {repr(e)}")
            logger.warning("Falling back to langdetect library")
            self.use_model      = False
            self.is_initialized = True
            return False
    

    def detect(self, text: str, **kwargs) -> LanguageDetectionResult:
        """
        Detect language of input text
        
        Arguments:
        ----------
            text     { str } : Input text to analyze

            **kwargs         : Additional options
            
        Returns:
        --------
            LanguageDetectionResult object
        """
        warnings = list()
        
        # Validate input
        if not text or not isinstance(text, str):
            return self._create_unknown_result(text     = "",
                                               warnings = ["Empty or invalid text"],
                                              )
        
        # Clean text for analysis
        cleaned_text = self._clean_text(text)
        char_count   = len(cleaned_text)
        word_count   = len(cleaned_text.split())
        
        # Check minimum length
        if (char_count < self.MIN_TEXT_LENGTH):
            warnings.append(f"Text too short ({char_count} chars, minimum {self.MIN_TEXT_LENGTH}). Detection may be unreliable.")
        
        # Detect script first
        script = self._detect_script(cleaned_text)
        
        # Try detection methods in order
        result = None
        
        # Method 1: ML Model
        if self.use_model and self.is_initialized:
            try:
                result                  = self._detect_with_model(text = cleaned_text)
                result.detection_method = "learned-language-representation"
            
            except Exception as e:
                logger.warning(f"Model detection failed: {repr(e)}, trying fallback")
                warnings.append("Model detection failed, using fallback")
        
        # Method 2: langdetect library
        if result is None and LANGDETECT_AVAILABLE:
            try:
                result                  = self._detect_with_langdetect(text = cleaned_text)
                result.detection_method = "statistical-language-estimation"
            
            except Exception as e:
                logger.warning(f"langdetect failed: {repr(e)}, trying heuristics")
                warnings.append("langdetect failed, using heuristics")
        
        # Method 3: Character-based heuristics
        if result is None:
            result                  = self._detect_with_heuristics(cleaned_text, script)
            result.detection_method = "character-distribution-analysis"
        
        # Add metadata
        result.script     = script
        result.char_count = char_count
        result.word_count = word_count
        result.warnings.extend(warnings)
        
        # Check for multilingual content using constant
        threshold         = language_detection_params.MULTILINGUAL_THRESHOLD

        if len([v for v in result.all_languages.values() if v > threshold]) > 1:
            result.is_multilingual = True
            warnings.append("Text appears to contain multiple languages")
        
        logger.info(f"Detected language: {result.primary_language.value} (evidence_strength: {result.evidence_strength:.2f}, method: {result.detection_method})")
        
        return result
    

    def _detect_with_model(self, text: str) -> LanguageDetectionResult:
        """
        Detect language using XLM-RoBERTa model with sentence-based chunking for more accurate detection on long texts
        """
        if not self.is_initialized:
            if not self.initialize():
                raise RuntimeError("Model not initialized")
        
        try:
            # Strategy: Use multiple text chunks for better accuracy
            chunks            = self._split_text_into_chunks(text = text)
            logger.info(f"Split text into {len(chunks)} chunks for language detection")
            
            all_chunk_results = list()
            
            for i, chunk in enumerate(chunks):
                try:
                    chunk_result = self._process_single_chunk(chunk = chunk)
                    all_chunk_results.append(chunk_result)
                    
                except Exception as e:
                    logger.warning(f"Chunk {i+1} processing failed: {repr(e)}")
                    continue
            
            if not all_chunk_results:
                raise RuntimeError("All chunks failed processing")
            
            # Aggregate results from all chunks
            return self._aggregate_chunk_results(chunk_results = all_chunk_results)
            
        except Exception as e:
            logger.error(f"Chunk-based model detection failed: {repr(e)}")
            raise


    def _split_text_into_chunks(self, text: str, max_chunk_length: int = language_detection_params.MAX_CHUNK_LENGTH, min_chunk_length: int = language_detection_params.MIN_CHUNK_LENGTH) -> List[str]:
        """
        Split text into meaningful chunks for language detection
        
        Arguments:
        ----------
            text             { str } : Input text

            max_chunk_length { int } : Maximum characters per chunk
            
            min_chunk_length { int } : Minimum characters per chunk
            
        Returns:
        --------
            List of text chunks
        """
        if (len(text) <= max_chunk_length):
            return [text]
        
        # Strategy 1: Split by sentences first
        sentences     = re.split(r'[.!?]+', text)
        sentences     = [s.strip() for s in sentences if s.strip()]
        
        chunks        = list()
        current_chunk = ""
        
        for sentence in sentences:
            # If adding this sentence doesn't exceed max length
            if ((len(current_chunk) + len(sentence) + 1) <= max_chunk_length):
                if current_chunk:
                    current_chunk += " " + sentence
                
                else:
                    current_chunk = sentence
            
            else:
                # Current chunk is full, save it
                if (current_chunk and (len(current_chunk) >= min_chunk_length)):
                    chunks.append(current_chunk)
                
                # Start new chunk with current sentence
                current_chunk = sentence
        
        # Add the last chunk if it meets minimum length
        if (current_chunk and (len(current_chunk) >= min_chunk_length)):
            chunks.append(current_chunk)
        
        # Strategy 2: If sentence splitting didn't work well, use fixed-length chunks
        if ((len(chunks) == 0) or ((len(chunks) == 1) and (len(chunks[0]) > max_chunk_length))):
            chunks = self._split_fixed_length(text, max_chunk_length)
        
        logger.debug(f"Split {len(text)} chars into {len(chunks)} chunks: {[len(c) for c in chunks]}")
        return chunks


    def _split_fixed_length(self, text: str, chunk_size: int = language_detection_params.FIXED_CHUNK_SIZE) -> List[str]:
        """
        Fallback: Split text into fixed-length chunks
        """
        chunks              = list()
        word_boundary_ratio = language_detection_params.WORD_BOUNDARY_RATIO

        for i in range(0, len(text), chunk_size):
            chunk = text[i:i + chunk_size]
            
            # Try to break at word boundaries
            if ((i + chunk_size) < len(text)):
                last_space = chunk.rfind(' ')
                # If we found a space in the last 30% (using word_boundary_ratio)
                if (last_space > chunk_size * word_boundary_ratio):
                    chunk = chunk[:last_space].strip()

            chunks.append(chunk)
        
        return chunks


    def _process_single_chunk(self, chunk: str) -> Dict:
        """
        Process a single chunk through the language detection model
        """
        # Get the tokenizer from the pipeline
        tokenizer  = self.classifier.tokenizer
        
        # Tokenize with explicit length limits using constant
        max_length = language_detection_params.MODEL_MAX_LENGTH
        inputs     = tokenizer(chunk,
                               return_tensors     = "pt",
                               truncation         = True,
                               max_length         = max_length,
                               padding            = True,
                               add_special_tokens = True,
                              )
        
        # Get model from pipeline
        model      = self.classifier.model
        device     = next(model.parameters()).device
        
        # Move inputs to correct device
        inputs     = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs     = model(**inputs)
            predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        
        # Get top predictions for this chunk using constant
        top_k           = language_detection_params.TOP_K_PREDICTIONS
        top_predictions = torch.topk(predictions[0], k=top_k)
        
        chunk_results = dict()

        for i in range(len(top_predictions.indices)):
            lang_idx  = top_predictions.indices[i].item()
            score     = top_predictions.values[i].item()
            
            # Get language label from model config
            lang_code = model.config.id2label[lang_idx]
            
            # Handle model output format
            if ('_' in lang_code):
                lang_code = lang_code.split('_')[0]
            
            chunk_results[lang_code] = score
        
        return chunk_results


    def _map_language_code(self, code: str) -> Language:
        """
        Map language code string to Language enum
        """
        code = code.lower()

        for lang in Language:
            if (lang.value == code):
                return lang

        return Language.UNKNOWN


    def _aggregate_chunk_results(self, chunk_results: List[Dict]) -> LanguageDetectionResult:
        """
        Aggregate results from multiple chunks using weighted averaging
        """
        # Combine scores from all chunks
        all_scores    = dict()
        chunk_weights = list()
        
        for chunk_result in chunk_results:
            # Calculate chunk weight based on evidence_strength and diversity
            top_score    = max(chunk_result.values()) if chunk_result else 0
            # Weight by evidence_strength
            chunk_weight = top_score  
            
            chunk_weights.append(chunk_weight)
            
            for lang_code, score in chunk_result.items():
                if lang_code not in all_scores:
                    all_scores[lang_code] = list()

                all_scores[lang_code].append(score)
        
        # Calculate weighted average for each language
        weighted_scores = dict()

        for lang_code, scores in all_scores.items():
            if (len(scores) != len(chunk_weights)):
                # Use simple average if weight mismatch
                weighted_scores[lang_code] = sum(scores) / len(scores)
            
            else:
                # Weighted average
                weighted_sum               = sum(score * weight for score, weight in zip(scores, chunk_weights))
                total_weight               = sum(chunk_weights)
                weighted_scores[lang_code] = (weighted_sum / total_weight if total_weight > 0 else sum(scores) / len(scores))
        
        # Find primary language
        primary_lang = None
        primary_conf = 0.0
        
        for lang_code, score in weighted_scores.items():
            if (score > primary_conf):
                primary_evidence_strength = score
                primary_lang              = lang_code
        
        # Convert to Language enum
        try:
            primary_language = self._map_language_code(code = primary_lang)
        
        except ValueError:
            primary_language = Language.UNKNOWN
        
        # Calculate detection quality metrics
        detection_quality = self._assess_detection_quality(chunk_results, weighted_scores)
        
        warnings          = list()

        if detection_quality.get('evidence_strength', False):
            warnings.append("Low evidence_strength across multiple chunks")

        if detection_quality.get('inconsistent', False):
            warnings.append("Inconsistent language detection across chunks")
        
        return LanguageDetectionResult(primary_language  = primary_language,
                                       evidence_strength = primary_evidence_strength,
                                       all_languages     = weighted_scores,
                                       script            = Script.UNKNOWN,
                                       is_multilingual   = detection_quality.get('multilingual', False),
                                       detection_method  = "model-chunked",
                                       char_count        = 0,
                                       word_count        = 0,
                                       warnings          = warnings,
                                      )


    def _assess_detection_quality(self, chunk_results: List[Dict], final_scores: Dict[str, float]) -> Dict[str, bool]:
        """
        Assess the quality and consistency of language detection across chunks
        """
        quality_metrics = {'low_evidence_strength' : False,
                           'inconsistent'          : False,
                           'multilingual'          : False,
                          }
        
        if not chunk_results:
            return quality_metrics
        
        # Check for low evidence_strength using constant
        avg_top_evidence_strength       = sum(max(chunk.values()) for chunk in chunk_results) / len(chunk_results)
        
        low_evidence_strength_threshold = language_detection_params.LOW_CONFIDENCE_THRESHOLD
        
        if (avg_top_evidence_strength < low_evidence_strength_threshold):
            quality_metrics['low_evidence_strength'] = True
        
        # Check for inconsistency (different primary languages across chunks)
        chunk_primaries = list()

        for chunk in chunk_results:
            if chunk:
                primary = max(chunk.items(), key=lambda x: x[1])[0]
                chunk_primaries.append(primary)
        
        if (len(set(chunk_primaries)) > 1):
            quality_metrics['inconsistent'] = True
        
        # Check for multilingual content using constant
        multilingual_threshold = language_detection_params.MULTILINGUAL_THRESHOLD
        strong_languages       = [lang for lang, score in final_scores.items() if score > multilingual_threshold]

        if (len(strong_languages) > 1):
            quality_metrics['multilingual'] = True
        
        return quality_metrics

    
    def _detect_with_langdetect(self, text: str) -> LanguageDetectionResult:
        """
        Detect language using langdetect library
        """
        # Get all language probabilities
        lang_probs    = detect_langs(text)
        
        all_languages = dict()

        for prob in lang_probs:
            all_languages[prob.lang] = prob.prob
        
        # Primary language
        primary = lang_probs[0]
        
        try:
            primary_language = self._map_language_code(code = primary.lang)
        
        except ValueError:
            primary_language = Language.UNKNOWN
        
        return LanguageDetectionResult(primary_language  = primary_language,
                                       evidence_strength = primary.prob,
                                       all_languages     = all_languages,
                                       script            = Script.UNKNOWN,
                                       is_multilingual   = False,
                                       detection_method  = "langdetect",
                                       char_count        = 0,
                                       word_count        = 0,
                                       warnings          = [],
                                      )
                                

    def _detect_with_heuristics(self, text: str, script: Script) -> LanguageDetectionResult:
        """
        Detect language using character-based heuristics
        """
        # Script-based language mapping
        script_to_language = {Script.CHINESE:    Language.CHINESE,
                              Script.JAPANESE:   Language.JAPANESE,
                              Script.KOREAN:     Language.KOREAN,
                              Script.ARABIC:     Language.ARABIC,
                              Script.CYRILLIC:   Language.RUSSIAN,
                              Script.DEVANAGARI: Language.HINDI,
                              Script.GREEK:      Language.GREEK,
                              Script.HEBREW:     Language.HEBREW,
                              Script.THAI:       Language.THAI,
                             }
        
        # If script clearly indicates language
        if script in script_to_language:
            primary_language  = script_to_language[script]
            # Moderate evidence_strength for heuristics
            evidence_strength = 0.7

        else:
            # For Latin script, check common words
            primary_language  = self._detect_latin_language(text)
            # Lower evidence_strength
            evidence_strength = 0.5
        
        return LanguageDetectionResult(primary_language  = primary_language,
                                       evidence_strength = evidence_strength,
                                       all_languages     = {primary_language.value: evidence_strength},
                                       script            = script,
                                       is_multilingual   = False,
                                       detection_method  = "heuristics",
                                       char_count        = 0,
                                       word_count        = 0,
                                       warnings          = ["Detection using heuristics, accuracy may be limited"],
                                      )
    

    def _detect_latin_language(self, text: str) -> Language:
        """
        Detect Latin-script language using common word patterns
        """
        text_lower = text.lower()
        
        # Common word patterns for major Latin-script languages
        patterns   = {Language.ENGLISH    : ['the', 'and', 'is', 'in', 'to', 'of', 'a', 'that', 'it', 'with', 'for', 'on', 'this', 'are', 'was', 'be', 'have', 'from', 'or', 'by'],
                      Language.SPANISH    : ['el', 'la', 'de', 'que', 'y', 'en', 'un', 'por', 'con', 'no', 'una', 'para', 'es', 'al', 'como', 'del', 'los', 'se', 'las', 'su'],
                      Language.FRENCH     : ['le', 'de', 'un', 'être', 'et', 'à', 'il', 'avoir', 'ne', 'je', 'son', 'que', 'ce', 'du', 'quel', 'elle', 'dans', 'pour', 'au', 'avec'],
                      Language.GERMAN     : ['der', 'die', 'und', 'in', 'den', 'von', 'zu', 'das', 'mit', 'sich', 'des', 'auf', 'für', 'ist', 'im', 'dem', 'nicht', 'ein', 'eine', 'als'],
                      Language.ITALIAN    : ['di', 'e', 'il', 'la', 'che', 'per', 'un', 'in', 'è', 'a', 'non', 'una', 'da', 'sono', 'come', 'del', 'ma', 'si', 'nel', 'anche'],
                      Language.PORTUGUESE : ['de', 'a', 'o', 'que', 'e', 'do', 'da', 'em', 'um', 'para', 'é', 'com', 'não', 'uma', 'os', 'no', 'se', 'na', 'por', 'mais'],
                     }
        
        # Count matches for each language
        scores     = dict()
        words      = set(text_lower.split())
        
        for lang, common_words in patterns.items():
            score        = sum(1 for word in common_words if word in words)
            scores[lang] = score
        
        # Return language with highest score
        if scores:
            best_lang = max(scores.items(), key=lambda x: x[1])
            # At least 3 matches
            if (best_lang[1] > 2):
                return best_lang[0]
        
        # Default to English for Latin script
        return Language.ENGLISH
    

    def _detect_script(self, text: str) -> Script:
        """
        Detect the writing script used in text
        """
        # Count characters in each script
        script_counts = {script: 0 for script in Script if script not in [Script.MIXED, Script.UNKNOWN]}
        
        for char in text:
            if char in string.whitespace or char in string.punctuation:
                continue
            
            code_point = ord(char)
            
            for script, ranges in self.SCRIPT_RANGES.items():
                for start, end in ranges:
                    if (start <= code_point <= end):
                        script_counts[script] += 1
                        break
        
        # Find dominant script
        total_chars = sum(script_counts.values())
        
        if (total_chars == 0):
            return Script.UNKNOWN
        
        # Calculate percentages
        script_percentages = {script: count / total_chars for script, count in script_counts.items() if count > 0}
        
        # Check if mixed using constant
        dominance_threshold = language_detection_params.SCRIPT_DOMINANCE_THRESHOLD
        if (len(script_percentages) > 1):
            max_percentage = max(script_percentages.values())
            if (max_percentage < dominance_threshold):
                return Script.MIXED
        
        # Return dominant script
        if script_percentages:
            return max(script_percentages.items(), key=lambda x: x[1])[0]
        
        return Script.UNKNOWN
    

    def _clean_text(self, text: str) -> str:
        """
        Clean text for language detection
        """
        # Remove URLs
        text = re.sub(r'https?://\S+', '', text)
        text = re.sub(r'www\.\S+', '', text)
        
        # Remove emails
        text = re.sub(r'\S+@\S+', '', text)
        
        # Remove excessive whitespace
        text = re.sub(r'\s+', ' ', text)
        
        return text.strip()
    

    def _create_unknown_result(self, text: str, warnings: List[str]) -> LanguageDetectionResult:
        """
        Create result for unknown language
        """
        return LanguageDetectionResult(primary_language  = Language.UNKNOWN,
                                       evidence_strength = 0.0,
                                       all_languages     = {},
                                       script            = Script.UNKNOWN,
                                       is_multilingual   = False,
                                       detection_method  = "none",
                                       char_count        = len(text),
                                       word_count        = len(text.split()),
                                       warnings          = warnings,
                                      )

    
    def is_language(self, text: str, target_language: Language, threshold: float = language_detection_params.LANGUAGE_MATCH_THRESHOLD
    ) -> bool:
        """
        Check if text is in a specific language
        
        Arguments:
        ----------
            text            : Input text

            target_language : Language to check for
            
            threshold       : Minimum evidence_strength threshold
            
        Returns:
        --------
            { bool }        : True if text is in target language with sufficient evidence_strength
        """
        result = self.detect(text)
        return ((result.primary_language == target_language) and (result.evidence_strength >= threshold))

    
    def get_supported_languages(self) -> List[str]:
        """
        Get list of supported language codes
        """
        return [lang.value for lang in Language if (lang != Language.UNKNOWN)]
    

    def cleanup(self):
        """
        Clean up resources
        """
        self.classifier     = None
        self.is_initialized = False



def quick_detect(text: str, **kwargs) -> LanguageDetectionResult:
    """
    Quick language detection with default settings
    
    Arguments:
    ----------
        text     : Input text
        **kwargs : Override settings
        
    Returns:
    --------
        LanguageDetectionResult object
    """
    detector = LanguageDetector(**kwargs)
    
    if detector.use_model:
        detector.initialize()

    return detector.detect(text)


def is_english(text: str, threshold: float = language_detection_params.LANGUAGE_MATCH_THRESHOLD) -> bool:
    """
    Quick check if text is English
    """
    detector   = LanguageDetector(use_model=True)
    is_english = detector.is_language(text, Language.ENGLISH, threshold)
    
    return is_english


# Export
__all__ = ['Script',
           'Language',
           'is_english',
           'quick_detect',
           'LanguageDetector',
           'LanguageDetectionResult',
          ]