# Screen.py
# Andrew Davison, ad@coe.psu.ac.th, April 2026
'''
  The class manages a 2D grid of cells, where each cell stores a color. The grid is square (gridCells * gridCells), with each cell rendered as a square of size cellSize * cellSize pixels.
  - The display window is only made visible when first needed.
  - Coordinates automatically wrap modulo the grid size.
  
  API:
  - draw(x, y, colorIdx): set a cell and update the display
  - clear(): reset all cells to WHITE
  - close(): terminate the display
'''

import os
os.environ["PYGAME_HIDE_SUPPORT_PROMPT"] = "1"
import time
import pygame
import sys

class Screen:
  # Palette indices (must match self.palette below)
  WHITE = 0
  BLACK = 1
  RED   = 2
  GREEN = 3
  BLUE  = 4
  YELLOW = 5
  GRAY = 6


  def __init__(self, cellSize=10, gridCells=80):
    pygame.init()
    self.gridSize = gridCells
    self.cellSize = cellSize
    self.dim = self.gridSize * self.cellSize

    # Do not call set_mode yet; keep the screen as None
    self.screen = None

    # colors 
    self.palette = [
      (255, 255, 255),  # 0 WHITE
      (0, 0, 0),    # 1 BLACK
      (255, 0, 0),    # 2 RED
      (0, 180, 0),    # 3 GREEN
      (0, 0, 255),    # 4 BLUE
      (200, 200, 0),  # 5 YELLOW
      (120, 120, 120),  # 6 GRAY
    ]

    # Initialize grid state to WHITE
    self.cells = [
      [0 for _ in range(self.gridSize)]
      for _ in range(self.gridSize)
    ]

    self._isRunning = True
    self._isVisible = False # Track if the window has been shown yet



  # ----------- API --------

  def draw(self, x, y, colorIdx):
    # Draws a box and wraps coordinates modulo gridSize
    if not self._isRunning: 
      return
    self._showWindow() # Reveal window on first draw
    self._set(x, y, colorIdx)
    self.render()
    self.pump() # Check for close button


  def clear(self):
    # Clears the grid to WHITE
    if not self._isRunning: 
      return
    self._showWindow() # Reveal window on first clear
    for y in range(self.gridSize):
      for x in range(self.gridSize):
        self.cells[y][x] = 0
    self.render()
    self.pump() # Check for close button


  def close(self):
    self._isRunning = False
    pygame.quit()


  @property
  def isRunning(self):
    return self._isRunning


  def pump(self):
    # Checks for window events like the close button
    if not self._isVisible or not self._isRunning:
      return self._isRunning

    for event in pygame.event.get():
      if event.type == pygame.QUIT:
        self.close()
      elif event.type == pygame.VIDEOEXPOSE:
        self.render()
    return self._isRunning


  # -------------------------------------------

  def _showWindow(self):
    if not self._isVisible:
      # Create the window
      self.screen = pygame.display.set_mode((self.dim, self.dim))
      pygame.display.set_caption(
               f"Box Grid -- {self.gridSize} x {self.gridSize}")
      self._isVisible = True



  def render(self):
    # Render the white background, gray grid, and colored cells
    if not self._isRunning or not self._isVisible:
      return
    self.screen.fill(self.palette[0]) # White background
    for y in range(self.gridSize):
      for x in range(self.gridSize):
        self._renderCell(x, y)
    pygame.display.flip()


  def _renderCell(self, x, y):
    sz = self.cellSize
    x0, y0 = x * sz, y * sz

    # Light gray grid lines
    pygame.draw.rect(self.screen, (240,240,240), (x0,y0,sz,sz), 1)

    colorIdx = self.cells[y][x]
    if colorIdx != 0:
      color = self.palette[self._mapColor(colorIdx)]
      pygame.draw.rect(self.screen, color, 
                          (x0+1, y0+1, sz-1, sz-1))


  def _set(self, x, y, colorIdx):
    x, y = x % self.gridSize, y % self.gridSize
    self.cells[y][x] = self._mapColor(colorIdx)


  def _mapColor(self, idx):
    if not isinstance(idx, int): 
      return 0
    return idx % len(self.palette)



# ----------- test rig ----------------
if __name__ == "__main__":
  print("Initializing Screen...")
  grid = Screen(cellSize=10, gridCells=80)
  time.sleep(5) # Window remains hidden during this delay

  print("Step 1: Drawing a colorful border...")
  for i in range(80):
    grid.draw(i, 0, Screen.RED) # Top
    grid.draw(i, 79, Screen.BLUE) # Bottom
    grid.draw(0, i, Screen.GREEN) # Left
    grid.draw(79, i, Screen.YELLOW)# Right
    time.sleep(0.01) # Small delay to see it draw

  print("Step 2: Testing coordinate wrapping (drawing at 85, 40)...")
  grid.draw(85, 40, Screen.BLACK) # Should appear at (5, 40)
  time.sleep(1)

  print("Step 3: Clearing the screen...")
  grid.clear()
  time.sleep(1)

  print("Step 4: Drawing a center block...")
  for x in range(35, 45):
    for y in range(35, 45):
      grid.draw(x, y, Screen.GRAY)

  print("\nAll tests complete. Close the window to exit.")
  while grid.pump():   # keep things alive until the exit
    time.sleep(0.03)