from pathlib import Path
from enum import Enum
from collections import defaultdict

SYSTEM_DICTIONARY = "/usr/share/dict/words"

class Hint(Enum):
    GRAY = 0
    YELLOW = 1
    GREEN = 2


def is_valid(word, i, letter, hint):
    if hint == Hint.GRAY:
        return letter not in word
    if hint == Hint.YELLOW:
        return letter in word and word[i] != letter
    if hint == Hint.GREEN:
        return word[i] == letter
    raise ValueError('wat?')


def filter_words(words, guess, hints):
    before_words = len(words)
    # Special case if we skip the guess
    if not hints:
        print(f"skipping: {guess}")
        words = [word for word in words if word != guess]
    else:
        for i, (letter, hint) in enumerate(zip(guess, hints)):
            words = [word for word in words if is_valid(word, i, letter, hint)]
    print(f"Pruned to {len(words)}/{before_words} words")
    return words


def best_guess(words):
    letter_freqs = defaultdict(int)
    for word in words:
        for letter in set(word):
            letter_freqs[letter] += 1
    word_scores = defaultdict(int)
    for word in words:
        for letter in set(word):
            word_scores[word] += letter_freqs[letter]
    return max(word_scores, key=word_scores.get)


def starting_words():
    words = Path(SYSTEM_DICTIONARY).read_text().splitlines()
    return [word for word in words if len(word) == 5]


def main():
    words = starting_words()
    print("Let's play wordl, use ctrl+c to exit")
    while True:
        guess = best_guess(words)
        print(f"Best guess: {guess}")
        hints = input(f"Enter hints (or blank to skip guess): ")
        hints = [Hint(int(hint)) for hint in hints]
        words = filter_words(words, guess, hints)


if __name__ == "__main__":
    main()