#!/usr/bin/env python3
"""Play EnvPool Pong-v5 in a small Tk window."""

from __future__ import annotations

import argparse
import tkinter as tk
from dataclasses import dataclass

import envpool
import numpy as np
from PIL import Image, ImageTk


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


def latest_rgb_frame(obs: np.ndarray) -> np.ndarray:
    arr = np.asarray(obs)
    if arr.ndim == 4:
        arr = arr[0]
    if arr.ndim == 3 and arr.shape[0] in (1, 3, 4):
        if arr.shape[0] == 1:
            arr = np.repeat(arr[0][..., None], 3, axis=-1)
        else:
            arr = np.moveaxis(arr[:3], 0, -1)
    elif arr.ndim == 3 and arr.shape[-1] >= 3:
        arr = arr[..., :3]
    elif arr.ndim == 2:
        arr = np.repeat(arr[..., None], 3, axis=-1)
    else:
        raise ValueError(f"Unsupported observation shape: {arr.shape}")
    return arr.astype(np.uint8, copy=False)


@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 PongWindow:
    def __init__(self, *, scale: int, fps: int, seed: int) -> None:
        self.env = envpool.make_gym(
            "Pong-v5",
            num_envs=1,
            batch_size=1,
            seed=seed,
            img_height=210,
            img_width=160,
            stack_num=1,
            gray_scale=False,
            frame_skip=1,
            noop_max=1,
            use_fire_reset=True,
            episodic_life=False,
            reward_clip=False,
            repeat_action_probability=0.0,
            full_action_space=False,
        )
        self.obs, self.info = self.env.reset()
        self.score = 0.0
        self.keys = KeyState()
        self.delay_ms = max(1, int(1000 / fps))
        self.scale = scale

        self.root = tk.Tk()
        self.root.title("EnvPool Pong-v5")
        self.root.resizable(False, False)
        self.image_label = tk.Label(self.root)
        self.image_label.pack()
        self.status = tk.Label(
            self.root,
            text="Up/Down arrows move paddle | Space fires | R resets | Esc quits",
            anchor="w",
        )
        self.status.pack(fill="x")

        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.obs, self.info = self.env.reset()
        self.score = 0.0

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

    def step(self) -> None:
        frame = latest_rgb_frame(self.obs)
        image = Image.fromarray(frame).resize(
            (frame.shape[1] * self.scale, frame.shape[0] * self.scale),
            Image.Resampling.NEAREST,
        )
        self.photo = ImageTk.PhotoImage(image)
        self.image_label.configure(image=self.photo)

        result = self.env.step(np.asarray([self.keys.action()], dtype=np.int32))
        self.obs, reward, terminated, truncated, self.info = result
        self.score += float(np.asarray(reward)[0])
        done = bool(np.logical_or(terminated, truncated)[0])
        self.status.configure(
            text=(
                "Up/Down move | Space fire | R reset | Esc quit"
                f"    score={self.score:.0f}"
            )
        )
        if done:
            self.reset()
        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("--scale", type=int, default=3)
    parser.add_argument("--fps", type=int, default=60)
    parser.add_argument("--seed", type=int, default=0)
    return parser.parse_args()


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