import random

import streamlit as st
from a3c.play import get_play_model_path, suggest
from wordle_env.wordle import get_env
from wordle_env.words import complete_vocabulary, target_vocabulary


def initial_config() -> None:
    """
    Function to configure the app
    """
    st.title("Wordle Game 🔠")
    if "secret_word" not in st.session_state:
        st.session_state["secret_word"] = random.choice(target_vocabulary)
    if "guesses" not in st.session_state:
        st.session_state["guesses"] = []
    if "won" not in st.session_state:
        st.session_state["won"] = False


def reset_state() -> None:
    """
    Deletes all the game state to start again
    """
    for key in st.session_state.keys():
        del st.session_state[key]


def game_not_finished() -> bool:
    """
    Condition that tells if the game is finished or not
    It could because there are not more attempts or because the user won
    """
    return not st.session_state.get("won") and len(st.session_state.guesses) < 6


def get_ai_guess() -> str:
    """
    Call the AI to make a guess based on the game state
    """
    env = get_env()
    model_path = get_play_model_path()
    words = list(map(lambda guess: guess[0], st.session_state.get("guesses")))
    states = list(map(lambda guess: guess[2], st.session_state.get("guesses")))
    ai_guess = suggest(env, words, states, model_path)
    return ai_guess


def render() -> None:
    """
    Function to render the app
    """
    st.markdown(
        "<style>.big-font {font-size:30px; display: inline-block; margin: 2px; padding: 10px; border-radius: 5px; text-align: center;}</style>",
        unsafe_allow_html=True,
    )

    if game_not_finished():
        guess_input = st.text_input("Enter your guess:", key="guess_input")
        col1, col2 = st.columns([0.2, 1])
        with col1:
            st.button("Submit", on_click=lambda: process_guess(guess_input))
        with col2:
            st.button("Ask AI ✨", on_click=lambda: process_guess(get_ai_guess()))
    elif st.session_state.get("won"):
        st.success("Congratulations! You guessed the word!")
    else:
        st.error("Game Over! The word was: " + st.session_state.secret_word)

    # Display guesses and feedback centered
    display_guesses()


def process_guess(guess: str) -> None:
    """
    Process the user guess and sets if the user won or not
    """
    guess = guess.upper()
    if not len(guess) == 5 or not guess.isalpha():
        st.warning("Please enter a valid 5-letter word.")
    elif guess not in complete_vocabulary:
        st.warning("That is not a valid word.")
    else:
        if guess == st.session_state.secret_word:
            st.session_state["won"] = True
        feedback, state = check_guess(guess)
        st.session_state.guesses.append((guess, feedback, state))


def display_guesses() -> None:
    """
    Display the user guesses with the feedback in a pretty format
    """
    for guess, feedback, _ in st.session_state.guesses:
        guess_display = '<div style="text-align: center;">'
        for i, letter in enumerate(guess):
            color = feedback[i]
            guess_display += f'<span class="big-font" style="background-color: {color};">{letter}</span>'
        guess_display += "</div>"
        st.markdown(guess_display, unsafe_allow_html=True)

    if game_not_finished():
        st.write("You have {} guesses left.".format(6 - len(st.session_state.guesses)))
    else:
        if st.button("Reset game"):
            reset_state()
            st.rerun()


def check_guess(guess: str) -> tuple[list[str], list[str]]:
    """
    Takes a guess and returns the guess feedback
    depending if the letters are in the right place or not
    """
    feedback = ["grey"] * 5
    word_list = list(st.session_state.secret_word)
    processed_letters = []
    state = ["0"] * 5

    # Check for correct position (green)
    for i, letter in enumerate(guess):
        if letter == word_list[i]:
            feedback[i] = "green"
            state[i] = "2"
            processed_letters.append(letter)

    # Check for correct letter but wrong position (yellow)
    for i, letter in enumerate(guess):
        if (
            feedback[i] == "grey"
            and letter in word_list
            and word_list.count(letter) > processed_letters.count(letter)
        ):
            feedback[i] = "yellow"
            state[i] = "1"
            processed_letters.append(letter)

    return feedback, state


# Main function
if __name__ == "__main__":
    initial_config()
    render()