import pygame
import random

# Initialize Pygame
pygame.init()

# Constants
SCREEN_WIDTH = 300
SCREEN_HEIGHT = 600
GRID_SIZE = 30
GRID_WIDTH = SCREEN_WIDTH // GRID_SIZE
GRID_HEIGHT = SCREEN_HEIGHT // GRID_SIZE
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# Initialize the screen
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Tetris")

# Tetromino shapes
SHAPES = [
    [[1, 1, 1, 1]],
    [[1, 1], [1, 1]],
    [[1, 1, 1], [0, 1, 0]],
    [[1, 1, 1], [1, 0, 0]],
    [[1, 1, 1], [0, 0, 1]],
    [[1, 1, 1], [0, 1, 0]],
    [[1, 1, 1], [1, 0, 0]],
]

# Tetromino colors
COLORS = [
    (255, 0, 0),
    (0, 255, 0),
    (0, 0, 255),
    (255, 255, 0),
    (255, 0, 255),
    (0, 255, 255),
    (128, 128, 128),
]

# Initialize the grid
grid = [[0] * GRID_WIDTH for _ in range(GRID_HEIGHT)]

# Tetromino class
class Tetromino:
    def __init__(self):
        self.shape = random.choice(SHAPES)
        self.color = random.choice(COLORS)
        self.x = GRID_WIDTH // 2 - len(self.shape[0]) // 2
        self.y = 0

    def rotate(self):
        self.shape = list(zip(*self.shape[::-1]))

    def move_down(self):
        self.y += 1

    def move_left(self):
        self.x -= 1

    def move_right(self):
        self.x += 1

    def can_move(self):
        for row, line in enumerate(self.shape):
            for col, cell in enumerate(line):
                if (
                    cell
                    and (self.y + row >= GRID_HEIGHT or self.x + col < 0 or self.x + col >= GRID_WIDTH)
                ):
                    return False
                if (
                    cell
                    and grid[self.y + row][self.x + col]
                ):
                    return False
        return True

    def place(self):
        for row, line in enumerate(self.shape):
            for col, cell in enumerate(line):
                if cell:
                    grid[self.y + row][self.x + col] = self.color

    def draw(self):
        for row, line in enumerate(self.shape):
            for col, cell in enumerate(line):
                if cell:
                    pygame.draw.rect(
                        screen,
                        self.color,
                        pygame.Rect((self.x + col) * GRID_SIZE, (self.y + row) * GRID_SIZE, GRID_SIZE, GRID_SIZE),
                    )


def draw_grid():
    for row in range(GRID_HEIGHT):
        for col in range(GRID_WIDTH):
            if grid[row][col]:
                pygame.draw.rect(screen, grid[row][col], pygame.Rect(col * GRID_SIZE, row * GRID_SIZE, GRID_SIZE, GRID_SIZE), 0)
            pygame.draw.rect(screen, BLACK, pygame.Rect(col * GRID_SIZE, row * GRID_SIZE, GRID_SIZE, GRID_SIZE), 1)

def clear_rows():
    rows_to_clear = []
    for row in range(GRID_HEIGHT):
        if all(grid[row]):
            rows_to_clear.append(row)

    for row in rows_to_clear:
        del grid[row]
        grid.insert(0, [0] * GRID_WIDTH)

# Create the Tetromino
current_piece = Tetromino()

# Game loop
clock = pygame.time.Clock()
game_over = False

while not game_over:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            game_over = True
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_DOWN:
                current_piece.move_down()
                if not current_piece.can_move():
                    current_piece.move_up()
            elif event.key == pygame.K_LEFT:
                current_piece.move_left()
                if not current_piece.can_move():
                    current_piece.move_right()
            elif event.key == pygame.K_RIGHT:
                current_piece.move_right()
                if not current_piece.can_move():
                    current_piece.move_left()
            elif event.key == pygame.K_UP:
                current_piece.rotate()
                if not current_piece.can_move():
                    current_piece.rotate()

    current_piece.move_down()
    if not current_piece.can_move():
        current_piece.move_up()
        current_piece.place()
        clear_rows()
        current_piece = Tetromino()
        if not current_piece.can_move():
            game_over = True

    screen.fill(WHITE)
    current_piece.draw()
    draw_grid()
    pygame.display.flip()
    clock.tick(5)

pygame.quit()