"""Word expansion logic for morphological analysis."""

from typing import Dict, List, Optional, Set, Tuple
import logging
from dataclasses import dataclass

from .affix import AffixFile, AffixRule, AffixType
from .dictionary import Dictionary
from .exceptions import ExpansionError


@dataclass
class ExpansionOptions:
    """Options for controlling word expansion behavior."""
    max_recursion_depth: int = 2
    enable_deep_recursion: bool = False
    process_intermediate_flags: bool = False
    extensive_cross_products: bool = False
    max_forms_per_word: int = 10000  # Safety limit
    
    @classmethod
    def conservative(cls) -> 'ExpansionOptions':
        """Conservative expansion options (default behavior)."""
        return cls()
    
    @classmethod
    def deep(cls) -> 'ExpansionOptions':
        """Deep expansion options for maximum morphological coverage."""
        return cls(
            max_recursion_depth=4,
            enable_deep_recursion=True,
            process_intermediate_flags=True,
            extensive_cross_products=True,
            max_forms_per_word=50000
        )
    
    @classmethod
    def balanced(cls) -> 'ExpansionOptions':
        """Balanced expansion options."""
        return cls(
            max_recursion_depth=3,
            enable_deep_recursion=True,
            process_intermediate_flags=True,
            extensive_cross_products=False,
            max_forms_per_word=20000
        )


class WordExpander:
    """Handles word expansion using affix rules."""
    
    def __init__(self, affix_file: Optional[AffixFile] = None, dictionary: Optional[Dictionary] = None, 
                 expansion_options: Optional[ExpansionOptions] = None):
        self.affix_file = affix_file
        self.dictionary = dictionary
        self.options = expansion_options or ExpansionOptions.conservative()
        # Legacy compatibility
        self.max_recursion_depth = self.options.max_recursion_depth
        
    def set_affix_file(self, affix_file: AffixFile) -> None:
        """Set the affix file to use for expansion."""
        self.affix_file = affix_file
        
    def set_dictionary(self, dictionary: Dictionary) -> None:
        """Set the dictionary to use for validation."""
        self.dictionary = dictionary
    
    def expand_word(self, word: str, flags: Optional[List[str]] = None) -> Set[str]:
        """
        Expand a word using affix rules.
        
        Args:
            word: The root word to expand
            flags: List of flags for the word (if not provided, uses dictionary)
            
        Returns:
            Set of all expanded word forms
        """
        if not self.affix_file:
            raise ExpansionError("No affix file loaded")
        
        if flags is None and self.dictionary:
            flags = self.dictionary.get_flags(word)
        elif flags is None:
            flags = []
        
        # Resolve flag aliases
        resolved_flags = self.affix_file.resolve_flag_aliases(flags)
        
        expanded = set()
        expanded.add(word)  # Always include the original word
        
        # Apply suffix rules first
        suffixed_forms = self._apply_suffix_rules(word, resolved_flags)
        expanded.update(suffixed_forms)
        
        # Apply prefix rules to original word
        prefixed_forms = self._apply_prefix_rules(word, resolved_flags)
        expanded.update(prefixed_forms)
        
        # Apply cross-product rules (prefix + suffix combinations)
        cross_product_forms = self._apply_cross_product_rules(word, resolved_flags)
        expanded.update(cross_product_forms)
        
        return expanded
    
    def _apply_suffix_rules(self, word: str, flags: List[str]) -> Set[str]:
        """Apply suffix rules to a word."""
        results = set()
        
        for flag in flags:
            suffix_rules = self.affix_file.get_suffix_rules(flag)
            for rule in suffix_rules:
                expanded_forms = self._apply_affix_rule(word, rule, AffixType.SUFFIX)
                results.update(expanded_forms)
        
        return results
    
    def _apply_prefix_rules(self, word: str, flags: List[str]) -> Set[str]:
        """Apply prefix rules to a word."""
        results = set()
        
        for flag in flags:
            prefix_rules = self.affix_file.get_prefix_rules(flag)
            for rule in prefix_rules:
                expanded_forms = self._apply_affix_rule(word, rule, AffixType.PREFIX)
                results.update(expanded_forms)
        
        return results
    
    def _apply_cross_product_rules(self, word: str, flags: List[str]) -> Set[str]:
        """Apply cross-product rules (combinations of prefixes and suffixes)."""
        results = set()
        
        # First, get all suffixed forms that allow cross-products
        suffixed_with_crossproduct = set()
        for flag in flags:
            suffix_rules = self.affix_file.get_suffix_rules(flag)
            for rule in suffix_rules:
                if rule.cross_product:
                    expanded_forms = self._apply_affix_rule(word, rule, AffixType.SUFFIX)
                    suffixed_with_crossproduct.update(expanded_forms)
        
        # Then apply prefix rules to these suffixed forms
        for suffixed_word in suffixed_with_crossproduct:
            for flag in flags:
                prefix_rules = self.affix_file.get_prefix_rules(flag)
                for rule in prefix_rules:
                    if rule.cross_product:
                        expanded_forms = self._apply_affix_rule(suffixed_word, rule, AffixType.PREFIX)
                        results.update(expanded_forms)
        
        return results
    
    def _apply_affix_rule(self, word: str, rule: AffixRule, affix_type: AffixType) -> Set[str]:
        """Apply a single affix rule to a word."""
        results = set()
        
        # Check if the rule's condition matches
        if not rule.matches_condition(word, affix_type):
            return results
        
        if affix_type == AffixType.PREFIX:
            # Check if we can strip the required characters from the beginning
            if len(word) >= len(rule.strip):
                if not rule.strip or word.startswith(rule.strip):
                    # Remove the strip part and add the prefix
                    remaining = word[len(rule.strip):]
                    new_word = rule.add + remaining
                    if new_word and new_word != word:  # Avoid infinite loops
                        results.add(new_word)
        else:  # SUFFIX
            # Check if we can strip the required characters from the end
            if len(word) >= len(rule.strip):
                if not rule.strip or word.endswith(rule.strip):
                    # Remove the strip part and add the suffix
                    remaining = word[:-len(rule.strip)] if rule.strip else word
                    new_word = remaining + rule.add
                    if new_word and new_word != word:  # Avoid infinite loops
                        results.add(new_word)
        
        return results
    
    def expand_dictionary_word(self, word: str) -> Set[str]:
        """
        Expand a word that exists in the dictionary.
        
        This method is specifically for dictionary words and will use
        their associated flags for expansion.
        """
        if not self.dictionary:
            raise ExpansionError("No dictionary loaded")
        
        if not self.dictionary.has_word(word):
            return {word}  # Return just the word if not in dictionary
        
        flags = self.dictionary.get_flags(word)
        return self.expand_word(word, flags)
    
    def expand_all_dictionary(self) -> Dict[str, Set[str]]:
        """
        Expand all words in the dictionary.
        
        Returns a dictionary mapping each root word to its expanded forms.
        """
        if not self.dictionary:
            raise ExpansionError("No dictionary loaded")
        
        results = {}
        
        for word, flags in self.dictionary.iter_entries():
            expanded = self.expand_word(word, flags)
            results[word] = expanded
        
        return results
    
    def expand_word_simple(self, word: str) -> Set[str]:
        """
        Expand a word using only affix rules, without dictionary validation.
        
        This method tries to generate possible word forms by applying
        available affix rules, useful for the --expand functionality.
        """
        if not self.affix_file:
            raise ExpansionError("No affix file loaded")
        
        if self.options.enable_deep_recursion:
            return self._expand_word_deep(word, set())
        else:
            return self._expand_word_simple_legacy(word)
    
    def _expand_word_simple_legacy(self, word: str) -> Set[str]:
        """Legacy simple expansion logic (original behavior)."""
        expanded = set()
        expanded.add(word)  # Always include the original word
        
        # Try all possible flags with this word
        all_flags = list(self.affix_file.get_all_flags())
        
        # Apply suffix rules
        for flag in all_flags:
            suffix_rules = self.affix_file.get_suffix_rules(flag)
            for rule in suffix_rules:
                new_forms = self._apply_affix_rule(word, rule, AffixType.SUFFIX)
                expanded.update(new_forms)
        
        # Apply prefix rules  
        for flag in all_flags:
            prefix_rules = self.affix_file.get_prefix_rules(flag)
            for rule in prefix_rules:
                new_forms = self._apply_affix_rule(word, rule, AffixType.PREFIX)
                expanded.update(new_forms)
        
        # Apply some cross-product combinations (limited to avoid explosion)
        word_variants = list(expanded)[:10]  # Limit to first 10 variants to avoid explosion
        
        for variant in word_variants:
            if variant == word:
                continue
                
            # Try adding prefixes to suffixed forms
            for flag in all_flags[:5]:  # Limit flags to avoid explosion
                prefix_rules = self.affix_file.get_prefix_rules(flag)
                for rule in prefix_rules[:3]:  # Limit rules per flag
                    if rule.cross_product:
                        new_forms = self._apply_affix_rule(variant, rule, AffixType.PREFIX)
                        expanded.update(list(new_forms)[:5])  # Limit results
        
        return expanded
    
    def _expand_word_deep(self, word: str, applied_flags: Set[str], depth: int = 0) -> Set[str]:
        """
        Deep recursive expansion with intermediate flag processing.
        
        This method implements the advanced features:
        1. Deep recursive expansion with multiple rule applications
        2. Processes intermediate flags from rules like word0/yoc  
        3. Extensively combines prefix+suffix+intermediate transformations
        """
        if depth > self.options.max_recursion_depth:
            return {word}
        
        expanded = set()
        expanded.add(word)
        
        if len(expanded) > self.options.max_forms_per_word:
            return expanded
        
        # Try all available flags
        all_flags = list(self.affix_file.get_all_flags())
        
        for flag in all_flags:
            if flag in applied_flags and depth > 0:
                continue  # Avoid immediate re-application of same flag
            
            # Apply suffix rules
            suffix_rules = self.affix_file.get_suffix_rules(flag)
            for rule in suffix_rules:
                new_forms_with_flags = self._apply_affix_rule_with_intermediate_flags(word, rule, AffixType.SUFFIX)
                
                for new_form, intermediate_flags in new_forms_with_flags:
                    if new_form and new_form != word:
                        expanded.add(new_form)
                        
                        # Process intermediate flags if enabled
                        if self.options.process_intermediate_flags and intermediate_flags:
                            new_applied = applied_flags | {flag}
                            recursive_forms = self._expand_word_deep(new_form, new_applied, depth + 1)
                            expanded.update(recursive_forms)
                            
                            # Apply intermediate flags recursively
                            for inter_flag in intermediate_flags:
                                if inter_flag not in new_applied:
                                    inter_forms = self._expand_word_deep(new_form, new_applied | {inter_flag}, depth + 1)
                                    expanded.update(inter_forms)
            
            # Apply prefix rules
            prefix_rules = self.affix_file.get_prefix_rules(flag)
            for rule in prefix_rules:
                new_forms_with_flags = self._apply_affix_rule_with_intermediate_flags(word, rule, AffixType.PREFIX)
                
                for new_form, intermediate_flags in new_forms_with_flags:
                    if new_form and new_form != word:
                        expanded.add(new_form)
                        
                        # Process intermediate flags if enabled
                        if self.options.process_intermediate_flags and intermediate_flags:
                            new_applied = applied_flags | {flag}
                            recursive_forms = self._expand_word_deep(new_form, new_applied, depth + 1)
                            expanded.update(recursive_forms)
                            
                            # Apply intermediate flags recursively
                            for inter_flag in intermediate_flags:
                                if inter_flag not in new_applied:
                                    inter_forms = self._expand_word_deep(new_form, new_applied | {inter_flag}, depth + 1)
                                    expanded.update(inter_forms)
        
        # Extensive cross-product combinations
        if self.options.extensive_cross_products and depth < self.options.max_recursion_depth - 1:
            expanded.update(self._apply_extensive_cross_products(word, applied_flags, depth))
        
        return expanded
    
    def _apply_affix_rule_with_intermediate_flags(self, word: str, rule: AffixRule, affix_type: AffixType) -> List[Tuple[str, List[str]]]:
        """
        Apply affix rule and extract intermediate flags.
        
        Returns list of (new_word, intermediate_flags) tuples.
        """
        results = []
        
        # Check if the rule's condition matches
        if not rule.matches_condition(word, affix_type):
            return results
        
        if affix_type == AffixType.PREFIX:
            # Check if we can strip the required characters from the beginning
            if len(word) >= len(rule.strip):
                if not rule.strip or word.startswith(rule.strip):
                    # Remove the strip part and add the prefix
                    remaining = word[len(rule.strip):]
                    
                    # Handle intermediate flags in the add part (e.g., "0/yoc")
                    add_part, intermediate_flags = self._parse_add_with_flags(rule.add)
                    new_word = add_part + remaining
                    
                    if new_word and new_word != word:  # Avoid infinite loops
                        results.append((new_word, intermediate_flags))
        else:  # SUFFIX
            # Check if we can strip the required characters from the end
            if len(word) >= len(rule.strip):
                if not rule.strip or word.endswith(rule.strip):
                    # Remove the strip part and add the suffix
                    remaining = word[:-len(rule.strip)] if rule.strip else word
                    
                    # Handle intermediate flags in the add part (e.g., "0/yoc")
                    add_part, intermediate_flags = self._parse_add_with_flags(rule.add)
                    new_word = remaining + add_part
                    
                    if new_word and new_word != word:  # Avoid infinite loops
                        results.append((new_word, intermediate_flags))
        
        return results
    
    def _parse_add_with_flags(self, add_string: str) -> Tuple[str, List[str]]:
        """
        Parse add string that might contain intermediate flags like '0/yoc'.
        
        Returns (add_part, intermediate_flags).
        """
        if '/' in add_string:
            add_part, flags_part = add_string.split('/', 1)
            intermediate_flags = self.affix_file._parse_flags(flags_part) if self.affix_file else []
            return add_part, intermediate_flags
        else:
            return add_string, []
    
    def _apply_extensive_cross_products(self, word: str, applied_flags: Set[str], depth: int) -> Set[str]:
        """Apply extensive cross-product combinations of prefixes and suffixes."""
        results = set()
        
        if not self.affix_file:
            return results
        
        # Get all cross-product enabled rules
        cross_product_prefixes = []
        cross_product_suffixes = []
        
        for flag in self.affix_file.get_all_flags():
            if flag in applied_flags:
                continue
                
            prefix_rules = self.affix_file.get_prefix_rules(flag)
            suffix_rules = self.affix_file.get_suffix_rules(flag)
            
            for rule in prefix_rules:
                if rule.cross_product:
                    cross_product_prefixes.append((flag, rule))
            
            for rule in suffix_rules:
                if rule.cross_product:
                    cross_product_suffixes.append((flag, rule))
        
        # Apply suffix first, then prefix (standard cross-product)
        for suf_flag, suf_rule in cross_product_suffixes[:10]:  # Limit to avoid explosion
            suf_forms_with_flags = self._apply_affix_rule_with_intermediate_flags(word, suf_rule, AffixType.SUFFIX)
            
            for suf_form, suf_intermediate in suf_forms_with_flags:
                if suf_form and suf_form != word:
                    results.add(suf_form)
                    
                    # Apply prefixes to the suffixed form
                    for pre_flag, pre_rule in cross_product_prefixes[:10]:
                        if pre_flag != suf_flag:  # Don't apply same flag
                            pre_forms_with_flags = self._apply_affix_rule_with_intermediate_flags(suf_form, pre_rule, AffixType.PREFIX)
                            
                            for pre_form, pre_intermediate in pre_forms_with_flags:
                                if pre_form and pre_form != suf_form:
                                    results.add(pre_form)
                                    
                                    # Recursively apply intermediate flags if any
                                    if self.options.process_intermediate_flags:
                                        all_intermediate = suf_intermediate + pre_intermediate
                                        for inter_flag in all_intermediate:
                                            if inter_flag not in applied_flags | {suf_flag, pre_flag}:
                                                new_applied = applied_flags | {suf_flag, pre_flag, inter_flag}
                                                inter_forms = self._expand_word_deep(pre_form, new_applied, depth + 1)
                                                results.update(inter_forms)

        return results

    def find_base_words(self, word: str) -> List[str]:
        """
        Find possible base words that could generate the given inflected word.

        This reverses the expansion process by:
        1. Trying to remove possible suffixes and prefixes
        2. Checking if resulting candidates exist in dictionary
        3. Verifying the original word can be generated from those candidates

        Args:
            word: The inflected word to analyze

        Returns:
            List of possible base words from the dictionary
        """
        if not self.affix_file or not self.dictionary:
            return []

        candidates = set()

        # Try removing suffixes
        for flag in self.affix_file.get_all_flags():
            suffix_rules = self.affix_file.get_suffix_rules(flag)
            for rule in suffix_rules:
                base_candidates = self._reverse_suffix_rule(word, rule)
                candidates.update(base_candidates)

        # Try removing prefixes
        for flag in self.affix_file.get_all_flags():
            prefix_rules = self.affix_file.get_prefix_rules(flag)
            for rule in prefix_rules:
                base_candidates = self._reverse_prefix_rule(word, rule)
                candidates.update(base_candidates)

        # Try removing combinations (prefix + suffix)
        candidates.update(self._reverse_cross_product_rules(word))

        # Filter candidates that exist in dictionary and can generate the original word
        valid_bases = []
        for candidate in candidates:
            if self.dictionary.has_word(candidate):
                # Verify this base can actually generate the original word
                expanded = self.expand_dictionary_word(candidate)
                if word in expanded:
                    valid_bases.append(candidate)

        return sorted(valid_bases)

    def _reverse_suffix_rule(self, word: str, rule: AffixRule) -> Set[str]:
        """Try to reverse a suffix rule to find possible base words."""
        candidates = set()

        # If word ends with the rule's add part, try removing it
        if rule.add and word.endswith(rule.add):
            # Remove the added suffix
            remaining = word[:-len(rule.add)]
            # Add back the stripped part
            candidate = remaining + rule.strip

            # Check if this candidate would match the rule's condition
            if candidate and rule.matches_condition(candidate, AffixType.SUFFIX):
                candidates.add(candidate)

        # Also try the case where nothing was added (rule.add is empty)
        elif not rule.add:
            # Add back the stripped part
            candidate = word + rule.strip
            if candidate and rule.matches_condition(candidate, AffixType.SUFFIX):
                candidates.add(candidate)

        return candidates

    def _reverse_prefix_rule(self, word: str, rule: AffixRule) -> Set[str]:
        """Try to reverse a prefix rule to find possible base words."""
        candidates = set()

        # If word starts with the rule's add part, try removing it
        if rule.add and word.startswith(rule.add):
            # Remove the added prefix
            remaining = word[len(rule.add):]
            # Add back the stripped part
            candidate = rule.strip + remaining

            # Check if this candidate would match the rule's condition
            if candidate and rule.matches_condition(candidate, AffixType.PREFIX):
                candidates.add(candidate)

        # Also try the case where nothing was added (rule.add is empty)
        elif not rule.add:
            # Add back the stripped part
            candidate = rule.strip + word
            if candidate and rule.matches_condition(candidate, AffixType.PREFIX):
                candidates.add(candidate)

        return candidates

    def _reverse_cross_product_rules(self, word: str) -> Set[str]:
        """Try to reverse cross-product rules (prefix + suffix combinations)."""
        candidates = set()

        # Try all combinations of prefix and suffix rules that allow cross-products
        for prefix_flag in self.affix_file.get_all_flags():
            prefix_rules = self.affix_file.get_prefix_rules(prefix_flag)
            for prefix_rule in prefix_rules:
                if not prefix_rule.cross_product:
                    continue

                # Try removing this prefix
                if prefix_rule.add and word.startswith(prefix_rule.add):
                    after_prefix_removal = word[len(prefix_rule.add):]

                    # Now try removing suffixes from what remains
                    for suffix_flag in self.affix_file.get_all_flags():
                        suffix_rules = self.affix_file.get_suffix_rules(suffix_flag)
                        for suffix_rule in suffix_rules:
                            if not suffix_rule.cross_product:
                                continue

                            if suffix_rule.add and after_prefix_removal.endswith(suffix_rule.add):
                                # Remove suffix, add back stripped parts
                                after_suffix_removal = after_prefix_removal[:-len(suffix_rule.add)]
                                candidate = prefix_rule.strip + after_suffix_removal + suffix_rule.strip

                                if (candidate and
                                    prefix_rule.matches_condition(candidate, AffixType.PREFIX) and
                                    suffix_rule.matches_condition(candidate, AffixType.SUFFIX)):
                                    candidates.add(candidate)

        return candidates

    def expand_from_base(self, word: str) -> Set[str]:
        """
        Find base words for the input and expand them using dictionary flags.

        This combines base word detection with proper expansion:
        1. Find possible base words that could generate the input
        2. Expand those base words using their dictionary flags
        3. Return all possible forms

        Args:
            word: The word to analyze and expand from its base forms

        Returns:
            Set of all expanded forms from valid base words
        """
        if not self.dictionary:
            # If no dictionary, fall back to simple expansion
            return self.expand_word_simple(word)

        base_words = self.find_base_words(word)

        if not base_words:
            # If no base words found, the input might already be a base word
            if self.dictionary.has_word(word):
                return self.expand_dictionary_word(word)
            else:
                # Fall back to simple expansion
                return self.expand_word_simple(word)

        # Expand all found base words
        all_forms = set()
        for base_word in base_words:
            expanded = self.expand_dictionary_word(base_word)
            all_forms.update(expanded)

        return all_forms