import random
from environments import register_env
import gradio as gr


class WordleEnv:
    """
    Demonstration env. Not a full game; 4-letter variant for brevity.

    Observations are emoji strings; actions are 4-letter lowercase words.
    Reward is 1.0 on success, else 0.0. Terminal on success or after 6 guesses.
    """

    dictionary = {"word", "wood", "ward", "sore", "bore", "bake", "bake", "bake", "earn"}

    def __init__(self, max_guesses: int = 6) -> None:
        self._max = max_guesses

    def reset(self) -> str:
        """Reset the environment and return the initial observation."""
        self._secret = random.choice(list(self.dictionary))
        self._n = 0
        self._obs = "⬜" * 4
        return self._obs

    def step(self, action: str) -> tuple[str, float, bool]:
        """
        Take an action (a 4-letter word) and return (observation, reward, done).
        If 
        When done is True, the episode has ended and reset() should be called to start a new episode.
        """
        if self._n >= self._max:
            return "The game is over. Please reset.", -1.0, True

        guess: str = str(action)
        guess = guess.strip().lower()

        if len(guess) != 4 or not guess.isalpha():
            return "Invalid guess. Must be a 4-letter word.", -1.0, False

        self._n += 1

        # Compute feedback
        feedback: list[str] = []
        for i, ch in enumerate(guess):
            if ch == self._secret[i]:
                feedback.append("🟩")
            elif ch in self._secret:
                feedback.append("🟨")
            else:
                feedback.append("⬜")
        self._obs = "".join(feedback)

        # Check for success or timeout and compute reward
        success = guess == self._secret
        timeout = self._n >= self._max
        done = success or timeout
        reward = 1.0 if success else 0.0
        if done and not success:
            self._obs += f" Game over. The word was '{self._secret}'."
        return self._obs, reward, done


with gr.Blocks() as demo:
    gr.Markdown(
        """# Wordle Environment

Usage:

```shell
pip install git+https://github.com/huggingface/environments.git
```

```python
>>> from environments import load
>>> env = load("qgallouedec/wordle")
Loaded as API: https://qgallouedec-counter.hf.space/ ✔
>>> env.reset()
('Counter reset to 0', 0.0, False)
>>> env.step()
('Counter is now 1', 1.0, False)
```
"""
    )
    register_env(WordleEnv)

demo.launch(mcp_server=True)