import re
import os
from collections import Counter, defaultdict
import itertools

def parse_ansi_feedback(raw_bytes):
    """
    Parses the raw bytes containing ANSI codes to extract letters and their colors.
    Returns a list of word_segments. 
    Each segment is a list of tuples: (color_code, char)
    """
    try:
        decoded = raw_bytes.decode('utf-8')
    except UnicodeDecodeError:
        decoded = raw_bytes.decode('utf-8', errors='ignore')

    print("--- Raw Decoded String ---")
    print(decoded)
    print("--------------------------")

    # Split into word chunks based on the underscore
    raw_segments = decoded.strip().split('_')
    
    parsed_segments = []

    # Regex to find ANSI code and the character immediately following
    # Matches: \x1b[93mU (captures '93' and 'U')
    ansi_pattern = re.compile(r'\x1b\[(\d+)m([A-Z])')

    for raw_seg in raw_segments:
        matches = ansi_pattern.findall(raw_seg)
        if matches:
            # matches is a list of ('93', 'U'), ('90', 'V'), etc.
            parsed_segments.append(matches)
            
    return parsed_segments

def analyze_global_constraints(all_segments):
    """
    Derives global letter counts from all segments combined based on the rule:
    - Green/Yellow count contributes to the total.
    - Presence of Grey implies the count is EXACTLY (Green + Yellow).
    - Absence of Grey implies the count is AT LEAST (Green + Yellow).
    """
    # Flatten all entries to count global occurrences in the guess
    all_entries = [item for seg in all_segments for item in seg]
    
    char_data = defaultdict(lambda: {'92': 0, '93': 0, '90': 0})
    
    for color, char in all_entries:
        char_data[char.lower()][color] += 1
        
    global_constraints = {}
    
    for char, counts in char_data.items():
        greens = counts['92']
        yellows = counts['93']
        greys = counts['90']
        
        required = greens + yellows
        
        if greys > 0:
            # We hit the limit (Grey exists), so the total count in solution is exactly 'required'
            global_constraints[char] = {'min': required, 'max': required}
        else:
            # No Greys seen, so there could be more. Count is at least 'required'
            global_constraints[char] = {'min': required, 'max': float('inf')}
            
    return global_constraints

def generate_local_constraints(segment_data):
    """
    Generates positional constraints for a single word slot.
    """
    constraints = {
        'length': len(segment_data),
        'fixed': {},        # index: char (Green)
        'not_at_pos': [],   # list of (index, char) (Yellow or Grey)
    }
    
    for i, (color, char) in enumerate(segment_data):
        char = char.lower()
        if color == '92': # Green
            constraints['fixed'][i] = char
        elif color == '93': # Yellow
            constraints['not_at_pos'].append((i, char))
        elif color == '90': # Grey
            # Grey means "not a match at this position".
            constraints['not_at_pos'].append((i, char))
            
    return constraints

def get_candidates_for_slots(wordlist_path, segments, global_constraints):
    """
    1. Filters wordlist for each slot based on local constraints.
    2. Prunes words containing letters that are globally forbidden (max=0).
    """
    if not os.path.exists(wordlist_path):
        print(f"Error: {wordlist_path} not found.")
        return []

    try:
        with open(wordlist_path, 'r', encoding='utf-8') as f:
            all_words = [w.strip().lower() for w in f if w.strip()]
    except Exception as e:
        print(f"Error reading wordlist: {e}")
        return []

    # Identify globally forbidden characters (max count is 0)
    forbidden_chars = {c for c, limits in global_constraints.items() if limits['max'] == 0}

    all_slots_candidates = []

    for i, seg in enumerate(segments):
        local = generate_local_constraints(seg)
        candidates = []
        
        print(f"\nScanning candidates for Slot {i+1} (Length {local['length']})...")

        for w in all_words:
            # 1. Length
            if len(w) != local['length']: continue
            
            # 2. Global "Zero Tolerance" check
            if any(c in forbidden_chars for c in w): continue

            # 3. Green (Fixed Position)
            if any(w[idx] != char for idx, char in local['fixed'].items()): continue
            
            # 4. Not At Position (Yellow/Grey)
            if any(w[idx] == char for idx, char in local['not_at_pos']): continue
            
            candidates.append(w)
            
        print(f"-> {len(candidates)} candidates found.")
        all_slots_candidates.append(candidates)
        
    return all_slots_candidates

def solve_combination(candidate_lists, global_constraints):
    """
    Finds a combination of one word from each list such that the total letter counts 
    satisfy the global constraints.
    """
    print("\n=== Solving for Valid Combination ===")
    
    # We use a recursive backtracking approach or itertools.product.
    # Given the likely search space, product might be heavy if candidates are many.
    # However, Python generators handle it decently if solutions appear early.
    
    # Pre-check: if any list is empty, fail early
    if any(not lst for lst in candidate_lists):
        print("Error: One or more slots have no candidates.")
        return None

    # Limit total combinations to avoid hanging if constraints are loose
    max_checks = 2000000 
    checks = 0

    import time
    start_time = time.time()

    for combination in itertools.product(*candidate_lists):
        checks += 1
        if checks % 100000 == 0:
            print(f"Checked {checks} combinations...")
        if checks > max_checks:
            print("Time limit exceeded searching for combinations.")
            break

        # Combine all letters
        combined_text = "".join(combination)
        counts = Counter(combined_text)
        
        valid = True
        # Check against global constraints
        for char, limits in global_constraints.items():
            c_count = counts[char]
            if c_count < limits['min']:
                valid = False; break
            if limits['max'] is not None and c_count != limits['max']:
                valid = False; break
        
        if valid:
            print(f"Solution Found in {time.time() - start_time:.2f}s!")
            return "_".join(combination)

    print("No valid combination found matching all constraints.")
    return None

# --- Main Execution ---

# The specific ANSI byte string provided
raw_ansi_output = b'\x1b[93mU\x1b[0m\x1b[93mN\x1b[0m\x1b[93mI\x1b[0m\x1b[90mV\x1b[0m\x1b[92mE\x1b[0m\x1b[93mR\x1b[0m\x1b[93mS\x1b[0m\x1b[90mI\x1b[0m\x1b[93mT\x1b[0m\x1b[90mY\x1b[0m_\x1b[90mC\x1b[0m\x1b[93mO\x1b[0m\x1b[93mN\x1b[0m\x1b[93mT\x1b[0m\x1b[90mA\x1b[0m\x1b[90mC\x1b[0m\x1b[90mT\x1b[0m_\x1b[90mI\x1b[0m\x1b[93mN\x1b[0m\x1b[90mT\x1b[0m\x1b[92mE\x1b[0m\x1b[90mR\x1b[0m\x1b[93mN\x1b[0m\x1b[90mA\x1b[0m\x1b[90mT\x1b[0m\x1b[90mI\x1b[0m\x1b[93mO\x1b[0m\x1b[90mN\x1b[0m\x1b[90mA\x1b[0m\x1b[93mL\x1b[0m\n'

# 1. Parse
segments = parse_ansi_feedback(raw_ansi_output)

# 2. Analyze Global Constraints
global_cons = analyze_global_constraints(segments)
print("\nGlobal Constraints (Total Counts in Solution):")
for c, lim in global_cons.items():
    mx = lim['max'] if lim['max'] is not None else "Inf"
    print(f"  {c.upper()}: {lim['min']} - {mx}")

# 3. Get Candidates per Slot
candidate_lists = get_candidates_for_slots('wordlist.txt', segments, global_cons)

# 4. Find valid combination
final_guess = solve_combination(candidate_lists, global_cons)

if final_guess:
    print("\n--------------------------------")
    print(f"FORMATTED GUESS: {final_guess}")
    print("--------------------------------")