from __future__ import annotations
import random
from typing import Callable
from retro.input import ProgrammaticInput
from retro.views.headless import HeadlessView
from retro_gamer.metadata import GameMetadata
from retro_gamer.observation import encode_observation


class GameEnvironment:
    """Gym-style wrapper around a retro game for RL training.

    The observation returned by reset()/step() comes from one of two mutually
    exclusive paths: metadata.observation_function, if set, fully replaces the
    built-in board/observe_state encoding — see GameMetadata for its contract.
    Otherwise the built-in encoder (encode_observation) is used, configured by
    observe_state/board/observe_state_sizes below.
    """

    def __init__(
        self,
        game_factory: Callable,
        metadata: GameMetadata,
        observe_state: list[str] | None = None,
        board: bool = True,
        observe_state_sizes: dict[str, int] | None = None,
    ):
        self.game_factory = game_factory
        self.metadata = metadata
        self.observe_state = observe_state or []
        self.board = board
        self.observe_state_sizes = observe_state_sizes or {}
        self._observation_fn = metadata.resolve_observation_function()
        if self._observation_fn is not None and self.observe_state:
            raise ValueError(
                "Both metadata.observation_function and [preprocessing].observe_state "
                "are set, but they're two conflicting ways of describing the\n"
                "observation. Use observation_function for full custom control, or\n"
                "observe_state (with the built-in board encoder) but not both."
            )
        self.game = None
        self.view: HeadlessView | None = None
        self.inp: ProgrammaticInput | None = None
        self._prev_reward: float = 0.0

    def reset(self):
        """Create a fresh game episode and return the initial observation.

        The observation's type depends on metadata.observation_function: a
        numpy array when using the built-in encoder (or a custom function
        built for DQN training), but it can be anything a custom function
        returns — e.g. a plain tuple for tabular use.
        """
        self.inp = ProgrammaticInput()
        self.view = HeadlessView()
        self.game = self.game_factory()
        self.game.input_source = self.inp
        self.game.view = self.view
        self.game.start()
        self._prev_reward = float(self.game.state.get(self.metadata.reward, 0))
        return self._observe()

    def step(self, action: str | None) -> tuple:
        """Advance one turn. Returns (observation, reward, done)."""
        self.inp.press(action)
        self.game.step()
        obs = self._observe()
        reward = self._delta_reward()
        done = not self.game.playing
        return obs, reward, done

    def _observe(self):
        if self._observation_fn is not None:
            return self._observation_fn(self.game)
        state = dict(self.game.state)
        return encode_observation(
            self.view.board_characters,
            state,
            self.metadata,
            self.observe_state,
            board=self.board,
        )

    def _delta_reward(self) -> float:
        current = float(self.game.state.get(self.metadata.reward, 0))
        delta = current - self._prev_reward
        self._prev_reward = current
        return delta

    def discover_character_set(self, exploration_turns: int) -> list[str]:
        """Run random turns to discover the characters that appear on the board."""
        self.reset()
        chars: set[str] = set()
        for _ in range(exploration_turns):
            for row in self.view.board_characters:
                chars.update(row)
            action = random.choice(self.metadata.actions + [None])
            _, _, done = self.step(action)
            if done:
                self.reset()
        chars.discard(' ')
        return sorted(chars)