#!/usr/bin/env python3
"""Play Pong through ALE's native Stella display window.

The Atari image is rendered by ALE/Stella when `display_screen=True`.
This script opens a tiny Tk controller window only to capture keyboard input.
Keep that controller window focused while playing.
"""

from __future__ import annotations

import argparse
import tkinter as tk
from dataclasses import dataclass

from ale_py import ALEInterface, roms


ACTION_NOOP = 0
ACTION_FIRE = 1
ACTION_UP = 2
ACTION_DOWN = 5


@dataclass
class KeyState:
    up: bool = False
    down: bool = False
    fire: bool = False

    def action(self) -> int:
        if self.fire:
            return ACTION_FIRE
        if self.up and not self.down:
            return ACTION_UP
        if self.down and not self.up:
            return ACTION_DOWN
        return ACTION_NOOP


class AlePongPlayer:
    def __init__(self, *, seed: int, fps: int) -> None:
        self.ale = ALEInterface()
        self.ale.setBool("display_screen", True)
        self.ale.setBool("sound", False)
        self.ale.setInt("random_seed", seed)
        self.ale.setInt("frame_skip", 1)
        self.ale.loadROM(str(roms.Pong))

        self.keys = KeyState()
        self.score = 0.0
        self.delay_ms = max(1, int(1000 / fps))

        self.root = tk.Tk()
        self.root.title("Pong Controls")
        self.root.resizable(False, False)
        self.label = tk.Label(
            self.root,
            padx=18,
            pady=14,
            justify="left",
            text=(
                "ALE native Pong is in the Atari window.\n\n"
                "Focus this controls window:\n"
                "  Up / Down: move paddle\n"
                "  Space: fire\n"
                "  R: reset\n"
                "  Esc: quit"
            ),
        )
        self.label.pack()

        self.root.bind("<KeyPress-Up>", lambda _: self._set_up(True))
        self.root.bind("<KeyRelease-Up>", lambda _: self._set_up(False))
        self.root.bind("<KeyPress-Down>", lambda _: self._set_down(True))
        self.root.bind("<KeyRelease-Down>", lambda _: self._set_down(False))
        self.root.bind("<KeyPress-space>", lambda _: self._set_fire(True))
        self.root.bind("<KeyRelease-space>", lambda _: self._set_fire(False))
        self.root.bind("<KeyPress-r>", lambda _: self.reset())
        self.root.bind("<Escape>", lambda _: self.close())
        self.root.protocol("WM_DELETE_WINDOW", self.close)

    def _set_up(self, value: bool) -> None:
        self.keys.up = value

    def _set_down(self, value: bool) -> None:
        self.keys.down = value

    def _set_fire(self, value: bool) -> None:
        self.keys.fire = value

    def reset(self) -> None:
        self.ale.reset_game()
        self.score = 0.0

    def close(self) -> None:
        self.root.destroy()

    def step(self) -> None:
        if self.ale.game_over():
            self.reset()
        reward = self.ale.act(self.keys.action())
        self.score += float(reward)
        self.label.configure(
            text=(
                "ALE native Pong is in the Atari window.\n\n"
                "Focus this controls window:\n"
                "  Up / Down: move paddle\n"
                "  Space: fire\n"
                "  R: reset\n"
                "  Esc: quit\n\n"
                f"score={self.score:.0f}"
            )
        )
        self.root.after(self.delay_ms, self.step)

    def run(self) -> None:
        self.step()
        self.root.mainloop()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--fps", type=int, default=60)
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    AlePongPlayer(seed=args.seed, fps=args.fps).run()
