import pygame
import random

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

# Tetrominoes
tetrominoes = [
    [[1, 1, 1, 1]],               # I
    [[1, 1, 1], [0, 1, 0]],      # T
    [[1, 1, 1], [1, 0, 0]],      # L
    [[1, 1, 1], [0, 0, 1]],      # J
    [[1, 1], [1, 1]],            # O
    [[0, 1, 1], [1, 1, 0]],      # S
    [[1, 1, 0], [0, 1, 1]]       # Z
]

# Initialize Pygame
pygame.init()

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

# Fonts
font = pygame.font.Font(None, 36)

# Function to draw a block
def draw_block(x, y, color):
    pygame.draw.rect(screen, color, (x * BLOCK_SIZE, y * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE))

# Function to draw the grid
def draw_grid():
    for x in range(GRID_WIDTH):
        for y in range(GRID_HEIGHT):
            draw_block(x, y, WHITE)
            pygame.draw.rect(screen, BLACK, (x * BLOCK_SIZE, y * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE), 1)

# Function to draw a tetromino
def draw_tetromino(x, y, shape, color):
    for i in range(len(shape)):
        for j in range(len(shape[i])):
            if shape[i][j]:
                draw_block(x + j, y + i, color)

# Function to check if a position is valid for the current tetromino
def is_valid_position(x, y, shape):
    for i in range(len(shape)):
        for j in range(len(shape[i])):
            if shape[i][j]:
                if x + j < 0 or x + j >= GRID_WIDTH or y + i >= GRID_HEIGHT:
                    return False
    return True

# Function to rotate a tetromino
def rotate_tetromino(shape):
    return [[shape[j][i] for j in range(len(shape))] for i in range(len(shape[0]) - 1, -1, -1)]

# Function to clear full rows
def clear_rows(grid):
    full_rows = []
    for i in range(len(grid)):
        if all(grid[i]):
            full_rows.append(i)
    for row in full_rows:
        del grid[row]
        grid.insert(0, [0] * GRID_WIDTH)

# Main function
def main():
    grid = [[0] * GRID_WIDTH for _ in range(GRID_HEIGHT)]
    current_tetromino = random.choice(tetrominoes)
    tetromino_x = GRID_WIDTH // 2 - len(current_tetromino[0]) // 2
    tetromino_y = 0
    score = 0
    clock = pygame.time.Clock()
    game_over = False

    while not game_over:
        screen.fill(BLACK)
        draw_grid()
        draw_tetromino(tetromino_x, tetromino_y, current_tetromino, WHITE)

        # Handle user input
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                game_over = True
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_LEFT:
                    if is_valid_position(tetromino_x - 1, tetromino_y, current_tetromino):
                        tetromino_x -= 1
                elif event.key == pygame.K_RIGHT:
                    if is_valid_position(tetromino_x + 1, tetromino_y, current_tetromino):
                        tetromino_x += 1
                elif event.key == pygame.K_DOWN:
                    if is_valid_position(tetromino_x, tetromino_y + 1, current_tetromino):
                        tetromino_y += 1
                elif event.key == pygame.K_UP:
                    rotated_tetromino = rotate_tetromino(current_tetromino)
                    if is_valid_position(tetromino_x, tetromino_y, rotated_tetromino):
                        current_tetromino = rotated_tetromino

        # Check for collision or landing
        if not is_valid_position(tetromino_x, tetromino_y + 1, current_tetromino):
            for i in range(len(current_tetromino)):
                for j in range(len(current_tetromino[i])):
                    if current_tetromino[i][j]:
                        grid[tetromino_y + i][tetromino_x + j] = 1
            clear_rows(grid)
            current_tetromino = random.choice(tetrominoes)
            tetromino_x = GRID_WIDTH // 2 - len(current_tetromino[0]) // 2
            tetromino_y = 0
            if not is_valid_position(tetromino_x, tetromino_y, current_tetromino):
                game_over = True
        else:
            tetromino_y += 1

        # Draw grid
        for i in range(GRID_HEIGHT):
            for j in range(GRID_WIDTH):
                if grid[i][j]:
                    draw_block(j, i, WHITE)

        # Draw score
        score_text = font.render(f"Score: {score}", True, WHITE)
        screen.blit(score_text, (10, 10))

        pygame.display.flip()
        clock.tick(5)

    pygame.quit()

if __name__ == "__main__":
    main()