import kandinsky as kd
import ion
import random
import time

# Screen dimensions for the NumWorks calculator
SCREEN_WIDTH = 320
SCREEN_HEIGHT = 240
CELL_SIZE = 10
GRID_WIDTH = SCREEN_WIDTH // CELL_SIZE
GRID_HEIGHT = SCREEN_HEIGHT // CELL_SIZE

# Tetromino shapes
SHAPES = [
    [[1, 1, 1, 1]],               # I
    [[1, 1], [1, 1]],             # O
    [[0, 1, 0], [1, 1, 1]],       # T
    [[1, 1, 0], [0, 1, 1]],       # Z
    [[0, 1, 1], [1, 1, 0]],       # S
    [[1, 1, 1], [1, 0, 0]],       # L
    [[1, 1, 1], [0, 0, 1]]        # J
]

# Colors for each tetromino
COLORS = [
    (0, 255, 255),  # Cyan for I
    (255, 255, 0),  # Yellow for O
    (128, 0, 128),  # Purple for T
    (255, 0, 0),    # Red for Z
    (0, 255, 0),    # Green for S
    (255, 165, 0),  # Orange for L
    (0, 0, 255)     # Blue for J
]

# Draw a single cell at a grid position
def draw_cell(x, y, color):
    kd.fill_rect(x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE, CELL_SIZE, color)

# Clear the screen
def clear_screen():
    kd.fill_rect(0, 0, SCREEN_WIDTH, SCREEN_HEIGHT, kd.white)

# Draw the entire grid
def draw_grid(grid):
    for y in range(GRID_HEIGHT):
        for x in range(GRID_WIDTH):
            if grid[y][x] != 0:
                draw_cell(x, y, COLORS[grid[y][x] - 1])

# Initialize a blank grid
def create_grid():
    return [[0] * GRID_WIDTH for _ in range(GRID_HEIGHT)]

# Check for collision
def check_collision(grid, shape, offset):
    off_x, off_y = offset
    for y, row in enumerate(shape):
        for x, cell in enumerate(row):
            if cell and (
                x + off_x < 0 or
                x + off_x >= GRID_WIDTH or
                y + off_y >= GRID_HEIGHT or
                grid[y + off_y][x + off_x]):
                return True
    return False

# Merge the tetromino with the grid
def merge_shape(grid, shape, offset):
    off_x, off_y = offset
    for y, row in enumerate(shape):
        for x, cell in enumerate(row):
            if cell:
                grid[y + off_y][x + off_x] = cell

# Rotate a shape
def rotate_shape(shape):
    return [list(row)[::-1] for row in zip(*shape)]

# Check for full rows and remove them
def remove_full_rows(grid):
    new_grid = [row for row in grid if any(cell == 0 for cell in row)]
    while len(new_grid) < GRID_HEIGHT:
        new_grid.insert(0, [0] * GRID_WIDTH)
    return new_grid

# Main game loop
def tetris():
    grid = create_grid()
    current_shape = random.choice(SHAPES)
    current_color = SHAPES.index(current_shape) + 1
    shape_pos = [GRID_WIDTH // 2 - len(current_shape[0]) // 2, 0]
    fall_time = 0.5
    fall_speed = time.monotonic()
    running = True

    while running:
        clear_screen()
        draw_grid(grid)

        # Move the shape down
        if time.monotonic() - fall_speed > fall_time:
            shape_pos[1] += 1
            if check_collision(grid, current_shape, shape_pos):
                shape_pos[1] -= 1
                merge_shape(grid, current_shape, shape_pos)
                grid = remove_full_rows(grid)
                current_shape = random.choice(SHAPES)
                current_color = SHAPES.index(current_shape) + 1
                shape_pos = [GRID_WIDTH // 2 - len(current_shape[0]) // 2, 0]
                if check_collision(grid, current_shape, shape_pos):
                    running = False
            fall_speed = time.monotonic()

        # Event handling
        if ion.keypad.is_pressed(ion.keypad.Key.LEFT):
            shape_pos[0] -= 1
            if check_collision(grid, current_shape, shape_pos):
                shape_pos[0] += 1
        elif ion.keypad.is_pressed(ion.keypad.Key.RIGHT):
            shape_pos[0] += 1
            if check_collision(grid, current_shape, shape_pos):
                shape_pos[0] -= 1
        elif ion.keypad.is_pressed(ion.keypad.Key.UP):
            current_shape = rotate_shape(current_shape)
            if check_collision(grid, current_shape, shape_pos):
                current_shape = rotate_shape(rotate_shape(rotate_shape(current_shape)))
        elif ion.keypad.is_pressed(ion.keypad.Key.DOWN):
            shape_pos[1] += 1
            if check_collision(grid, current_shape, shape_pos):
                shape_pos[1] -= 1

        # Draw the current shape
        for y, row in enumerate(current_shape):
            for x, cell in enumerate(row):
                if cell:
                    draw_cell(shape_pos[0] + x, shape_pos[1] + y, COLORS[current_color - 1])

        # Delay for the game loop
        time.sleep(0.1)

    kd.fill_rect(0, 0, SCREEN_WIDTH, SCREEN_HEIGHT, kd.white)
    kd.draw_string("Game Over!", 100, 100, kd.red)

# Start the game
tetris()