from __future__ import annotations

import unittest

import numpy as np

from pong_state import PongState
from scripts.run_eval import rgb_to_frame


class VideoOverlayTests(unittest.TestCase):
    def test_rgb_overlay_marks_blank_frame(self) -> None:
        frame = np.zeros((1, 210, 160, 3), dtype=np.uint8)
        state = PongState(
            ball_x=80,
            ball_y=100,
            self_paddle_y=120,
            opponent_paddle_y=90,
            ball_active=True,
            ball_in_right_half=False,
        )

        image = rgb_to_frame(
            frame,
            scale=2,
            overlay={"step": 7, "action": 2, "reward": 0.0, "score": -1.0, "state": state},
        )

        self.assertEqual(image.size, (320, 420))
        self.assertGreater(np.asarray(image).sum(), 0)


if __name__ == "__main__":
    unittest.main()
