import math
import re

with open('wordle-words.txt', 'r') as file:
    WORD_LIST = [ word.strip() for word in file.readlines() ]

class Letter:
    char: str # character
    pos: int = -1 # -1 is unknown
    guessed: list[int] = [] # positions that were yellow
    found: bool = False

    def __init__(self, char: str, pos: int, found: bool = False):
        self.char = char
        self.guessed = []
        if found:
            self.pos = pos
        else:
            self.guessed.append(pos)

        self.found = found

    def __str__(self) -> str:
        return (f"Letter(char='{self.char}', pos={self.pos}, "
                f"guessed={self.guessed}, found={self.found})")

class WordleSolver:
    known_letters: list[str] = []
    word: str = '?????'
    contains: list[Letter] = []
    doesnt_contain: set[str] = set()
    possible_words: list[str] = WORD_LIST

    scores: list[tuple[str, int]] = []

    def __init__(self):
        self.known_letters = []
        self.word = '?????'
        self.contains = []
        self.doesnt_contain = set()
        self.possible_words = WORD_LIST
        self.scores = []

    # How many words does this guess rule out
    def score_word(self, guess: str) -> int:

        letters = set(guess)
        ruled_out = 0
        for word in self.possible_words:
            letters_in_word = False

            for letter in letters:
                if letter in word:
                    # print(f'{letter} in {word}')
                    letters_in_word = True
                    break

            if letters_in_word:
                ruled_out += 1

        return ruled_out

    def score_words(self) -> None:
        self.scores = []

        for word in self.possible_words:
            self.scores.append((word, self.score_word(word)))

        # Sort scores by the score value (descending)
        self.scores.sort(key=lambda x: x[1], reverse=True)

    def add_guess(self, guess: str, colors: str) -> None:
        # Pair each letter/color with its original index
        zipped = list(zip(guess, colors, range(len(guess))))
        order = {'g': 0, 'y': 1, 'b': 2}
        zipped_sorted = sorted(zipped, key=lambda x: order[x[1]])
        guess_sorted = ''.join([x[0] for x in zipped_sorted])
        colors_sorted = ''.join([x[1] for x in zipped_sorted])
        original_indices = [x[2] for x in zipped_sorted]

        # Now use guess_sorted, colors_sorted, and original_indices for processing
        for i, color in enumerate(colors_sorted):
            letter = guess_sorted[i]
            orig_idx = original_indices[i]
            match color:
                case 'y':
                    editted_existing = False

                    for containie in self.contains:
                        if containie.char == letter and not containie.found:
                            containie.guessed.append(orig_idx)
                            editted_existing = True

                    if not editted_existing:
                        self.contains.append(Letter(letter, orig_idx))
                case 'b':
                    self.doesnt_contain.add(letter)
                case 'g':
                    editted_existing = False

                    for containie in self.contains:
                        if containie.char == letter:
                            if containie.found and containie.pos == orig_idx:
                                editted_existing = True
                                break

                            containie.pos = orig_idx
                            containie.found = True
                            editted_existing = True

                    if not editted_existing:
                        self.contains.append(Letter(letter, orig_idx, found=True))

        self.prune_words()

    def prune_words(self) -> None:
        # Build regex pattern for green and yellow constraints
        pattern = ''
        for i in range(5):
            letter_at_pos = None
            for letter in self.contains:
                if letter.found and letter.pos == i:
                    letter_at_pos = letter.char
                    break
            if letter_at_pos:
                pattern += letter_at_pos
            else:
                # Exclude yellow letters at this position
                forbidden = [l.char for l in self.contains if i in l.guessed]
                if forbidden:
                    pattern += f"[^{''.join(forbidden)}]"
                else:
                    pattern += '.'

        regex = re.compile(f"^{pattern}$")

        filtered = []
        for word in self.possible_words:
            # Exclude words with black letters
            if any(b in word for b in self.doesnt_contain):
                continue
            # Must match green/yellow regex
            if not regex.match(word):
                continue
            # Each yellow letter must be present somewhere (not just at forbidden positions)
            valid = True
            for letter in self.contains:
                if letter.guessed:
                    # Count how many times this letter should appear (yellow + green)
                    required_count = len(letter.guessed)
                    if letter.found:
                        required_count += 1
                    if word.count(letter.char) < required_count:
                        valid = False
                        break
                    # Not at forbidden positions
                    for pos in letter.guessed:
                        if word[pos] == letter.char:
                            valid = False
                            break
            if valid:
                filtered.append(word)
        self.possible_words = filtered

ws = WordleSolver()

ws.add_guess('naieo', 'bbgyb')
ws.add_guess('beige', 'bygbg')
# # ws.add_guess('blowy', 'byybg')

for letter in ws.contains:
    print(letter)

ws.score_words()

with open('word_scores.txt', 'w') as file:
    for word, score in ws.scores:
        file.write(f'{word},{score}\n')