from pwn import *
import z3
import re
import os
from collections import defaultdict

# --- Configuration ---
# Update these if the challenge parameters change
HOST = "chall.polygl0ts.ch"
PORT = 6052
WORDLIST_PATH = "word_list.txt"

context.log_level = 'info'

def get_word_lengths(raw_bytes):
    """Parses the initial 'Structure: ...' line to get word lengths."""
    try:
        decoded = raw_bytes.decode('utf-8', errors='ignore')
    except:
        return []

    # Look for the visual structure line
    for line in decoded.splitlines():
        if "Structure:" in line:
            # Example: Structure: ■■■■■_■■■■
            pattern_str = line.split("Structure:")[1].strip()
            word_segments = pattern_str.split('_')
            # Length is number of block characters (or just length of segment if stripped)
            # UTF-8 block is 3 bytes, but decoded string is 1 char.
            return [len(segment) for segment in word_segments]
    return []

def load_wordlist(path):
    """Loads wordlist into a dictionary keyed by length."""
    if not os.path.exists(path):
        log.error(f"Wordlist {path} not found!")
        exit(1)
    
    words_by_len = defaultdict(list)
    with open(path, 'r', encoding='utf-8', errors='ignore') as f:
        for line in f:
            w = line.strip().lower()
            if w:
                words_by_len[len(w)].append(w)
    return words_by_len

def parse_ansi_feedback(raw_bytes):
    """Extracts (color, char) tuples from ANSI output."""
    try:
        decoded = raw_bytes.decode('utf-8', errors='ignore')
    except:
        return []

    # Heuristic: Find the line with ANSI codes and underscores
    lines = decoded.splitlines()
    target_line = ""
    for line in lines:
        if "\x1b[" in line and "_" in line:
            target_line = line.strip()
            break
    
    # If heuristic fails, scan whole block
    if not target_line:
        target_line = decoded

    raw_segments = target_line.split('_')
    parsed_segments = []
    # Matches ANSI code and the letter immediately after
    ansi_pattern = re.compile(r'\x1b\[(\d+)m([A-Z])')

    for raw_seg in raw_segments:
        matches = ansi_pattern.findall(raw_seg)
        if matches:
            parsed_segments.append(matches)
            
    return parsed_segments

def analyze_global_constraints(all_segments):
    """
    Returns constraints for *this specific guess result*.
    Format: {char: {'min': k, 'max': k or None}}
    """
    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
        
    constraints = {}
    for char, counts in char_data.items():
        greens = counts['92']
        yellows = counts['93']
        greys = counts['90']
        
        required = greens + yellows
        
        if greys > 0:
            # If we see a grey, the count is exact
            constraints[char] = {'min': required, 'max': required}
        else:
            # If no grey, we only know the minimum
            constraints[char] = {'min': required, 'max': None}
    return constraints

def filter_candidates(current_lists, segments, global_cons):
    """
    Filters the *existing* candidate lists based on positional feedback (Green/Yellow/Grey).
    This reduces the search space for Z3 significantly.
    """
    new_lists = []
    
    # Identify chars that are globally forbidden (max count is 0)
    forbidden_chars = {c for c, lim in global_cons.items() if lim['max'] == 0}

    for i, seg in enumerate(segments):
        # 1. Build Local Constraints for this slot
        local_fixed = {}        # Index -> Char (Green)
        local_not_at = []       # (Index, Char) (Yellow or Grey)
        
        for idx, (color, char) in enumerate(seg):
            c = char.lower()
            if color == '92': # Green
                local_fixed[idx] = c
            elif color == '93': # Yellow
                local_not_at.append((idx, c))
            elif color == '90': # Grey
                # Grey in Wordle means "not at this position" locally
                local_not_at.append((idx, c))

        # 2. Apply to current candidates
        filtered = []
        for w in current_lists[i]:
            # A. Global Forbidden Check
            if any(c in forbidden_chars for c in w): continue
            
            # B. Green Check (Must match exact position)
            if any(w[pos] != char for pos, char in local_fixed.items()): continue
            
            # C. Yellow/Grey Check (Must NOT be at this position)
            if any(w[pos] == char for pos, char in local_not_at): continue
            
            filtered.append(w)
        
        new_lists.append(filtered)
        log.info(f"Slot {i+1}: Candidates reduced from {len(current_lists[i])} to {len(filtered)}")
        
    return new_lists

def solve_with_z3(candidate_lists, global_constraints):
    """
    Uses Z3 to pick one word from each candidate list such that their combined 
    letter counts satisfy the global_constraints.
    """
    s = z3.Solver()
    selected_flags = []
    
    # 1. Create variables: One boolean per candidate word
    for slot_idx, words in enumerate(candidate_lists):
        slot_flags = []
        for word_idx, w in enumerate(words):
            v = z3.Bool(f"s{slot_idx}_{word_idx}")
            slot_flags.append(v)
        selected_flags.append(slot_flags)
        
        # Constraint: Select exactly one word per slot
        s.add(z3.PbEq([(f, 1) for f in slot_flags], 1))

    # 2. Add Global Count Constraints
    # We sum the character counts of selected words
    for char, limits in global_constraints.items():
        terms = []
        
        for slot_idx, words in enumerate(candidate_lists):
            # Pre-calculate counts for efficiency
            # If a word has 2 'e's, we add (2 * bool_var) to the sum
            for w_idx, w in enumerate(words):
                c_count = w.count(char)
                if c_count > 0:
                    terms.append(z3.If(selected_flags[slot_idx][w_idx], c_count, 0))
        
        if not terms:
            # Edge case: If min > 0 but no candidates have the char -> UNSAT
            if limits['min'] > 0:
                s.add(False)
            continue
            
        total = z3.Sum(terms)
        
        s.add(total >= limits['min'])
        if limits['max'] is not None:
            s.add(total <= limits['max'])

    log.info("Z3: Solving for valid combination...")
    if s.check() == z3.sat:
        m = s.model()
        result_words = []
        for slot_idx, flags in enumerate(selected_flags):
            for w_idx, f in enumerate(flags):
                if z3.is_true(m[f]):
                    result_words.append(candidate_lists[slot_idx][w_idx])
                    break
        return "_".join(result_words)
    else:
        return None

def main():
    # 1. Connect
    r = remote(HOST, PORT)
    
    # 2. Read Initial State
    log.info("Reading initial structure...")
    try:
        initial_output = r.recvuntil(b"Your guess: ", timeout=5)
    except:
        initial_output = r.recv()
        
    log.info(f"Initial Output: {initial_output.decode(errors='ignore')}")
    
    # 3. Parse Lengths & Setup
    lengths = get_word_lengths(initial_output)
    if not lengths:
        log.error("Could not parse word lengths! Check connection or format.")
        return
    log.success(f"Target Lengths: {lengths}")

    word_dict = load_wordlist(WORDLIST_PATH)
    
    # Initialize candidates for each slot based on length
    candidates = []
    for l in lengths:
        if l in word_dict:
            candidates.append(word_dict[l])
        else:
            log.error(f"No words of length {l} in wordlist!")
            return

    # Cumulative Global Constraints (start empty)
    # format: {char: {'min': 0, 'max': None}}
    cum_global_cons = defaultdict(lambda: {'min': 0, 'max': None})

    # 4. Game Loop
    # We loop until solved.
    # First guess: Just pick the first word from each list (random start)
    current_guess = "_".join([c[0] for c in candidates])
    
    step = 0
    while True:
        step += 1
        log.info(f"--- Round {step} ---")
        log.info(f"Sending guess: {current_guess}")
        
        r.sendline(current_guess.encode())
        
        # Read response
        try:
            response = r.recvuntil(b"Your guess: ", timeout=2)
        except:
            # Connection might close on win, or lag
            response = r.recvall(timeout=1)
            if b"Correct" in response or b"Flag" in response or b"flag" in response:
                log.success(f"WINNER! \n{response.decode(errors='ignore')}")
                break
            else:
                log.info("Stream ended or timed out.")
                log.info(response.decode(errors='ignore'))
                break

        # Check for Win text
        if b"Correct!" in response:
            log.success("Correct solution found!")
            log.info(response.decode(errors='ignore'))
            break

        # 5. Process Feedback
        segments = parse_ansi_feedback(response)
        if not segments:
            log.warning("Could not parse ANSI feedback. Retrying read...")
            continue
            
        # 6. Update Constraints & Candidates
        
        # Update Globals (Accumulate)
        new_globals = analyze_global_constraints(segments)
        for char, limits in new_globals.items():
            # Min can only go up (if we need at least 1, then at least 2...)
            cum_global_cons[char]['min'] = max(cum_global_cons[char]['min'], limits['min'])
            
            # Max can only go down (if max 3, then max 2...)
            if limits['max'] is not None:
                current_max = cum_global_cons[char]['max']
                if current_max is None:
                    cum_global_cons[char]['max'] = limits['max']
                else:
                    cum_global_cons[char]['max'] = min(current_max, limits['max'])

        # Filter Candidates (Locals)
        candidates = filter_candidates(candidates, segments, cum_global_cons)
        
        if any(len(c) == 0 for c in candidates):
            log.error("Empty candidate list! The correct word is missing from word_list.txt.")
            break

        # 7. Generate Next Guess using Z3
        guess_str = solve_with_z3(candidates, cum_global_cons)
        
        if not guess_str:
            log.error("Z3 says UNSAT. Contradictory constraints or missing words.")
            break
            
        current_guess = guess_str

if __name__ == "__main__":
    main()