import streamlit as st
import pygame
import random
import time

# Initialize Pygame
pygame.init()

# Game Constants
WIDTH, HEIGHT = 500, 500
CELL_SIZE = 20
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLACK = (0, 0, 0)

# Snake and Food
snake = [(100, 100), (90, 100), (80, 100)]
direction = "RIGHT"
food = (random.randrange(0, WIDTH, CELL_SIZE), random.randrange(0, HEIGHT, CELL_SIZE))
score = 0

def move_snake():
    global snake, food, score, direction
    head_x, head_y = snake[0]
    
    if direction == "UP":
        head_y -= CELL_SIZE
    elif direction == "DOWN":
        head_y += CELL_SIZE
    elif direction == "LEFT":
        head_x -= CELL_SIZE
    elif direction == "RIGHT":
        head_x += CELL_SIZE
    
    new_head = (head_x, head_y)
    
    if new_head in snake or head_x < 0 or head_x >= WIDTH or head_y < 0 or head_y >= HEIGHT:
        return False
    
    snake.insert(0, new_head)
    
    if new_head == food:
        score += 1
        food = (random.randrange(0, WIDTH, CELL_SIZE), random.randrange(0, HEIGHT, CELL_SIZE))
    else:
        snake.pop()
    
    return True

def draw_game(screen):
    screen.fill(BLACK)
    
    for segment in snake:
        pygame.draw.rect(screen, GREEN, (*segment, CELL_SIZE, CELL_SIZE))
    
    pygame.draw.rect(screen, RED, (*food, CELL_SIZE, CELL_SIZE))
    
    pygame.display.flip()

def game_loop():
    global direction
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    clock = pygame.time.Clock()
    running = True
    
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP and direction != "DOWN":
                    direction = "UP"
                elif event.key == pygame.K_DOWN and direction != "UP":
                    direction = "DOWN"
                elif event.key == pygame.K_LEFT and direction != "RIGHT":
                    direction = "LEFT"
                elif event.key == pygame.K_RIGHT and direction != "LEFT":
                    direction = "RIGHT"
        
        if not move_snake():
            break
        
        draw_game(screen)
        clock.tick(10)
    
    pygame.quit()

def main():
    st.title("Snake Game in Streamlit")
    st.write("Use arrow keys to move the snake.")
    
    if st.button("Start Game"):
        game_loop()

if __name__ == "__main__":
    main()