import gradio as gr
import numpy as np
import time
from PIL import Image, ImageDraw
import io
import base64

# Game state class
class SpaceShooter:
    def __init__(self):
        self.width = 480
        self.height = 640
        self.score = 0
        self.lives = 3
        
        # Player
        self.player = {
            "x": self.width / 2,
            "y": self.height - 60,
            "width": 40,
            "height": 40,
            "speed": 5
        }
        
        self.bullets = []
        self.enemies = []
        self.last_enemy_spawn = time.time()
        self.enemy_spawn_rate = 1.5  # seconds
        self.last_shot = 0
        self.shoot_cooldown = 0.3  # seconds

    def spawn_enemy(self, current_time):
        if current_time - self.last_enemy_spawn > self.enemy_spawn_rate:
            size = np.random.uniform(20, 40)
            self.enemies.append({
                "x": np.random.uniform(size/2, self.width - size/2),
                "y": 0,
                "width": size,
                "height": size,
                "speed": np.random.uniform(1, 3)
            })
            self.last_enemy_spawn = current_time

    def shoot(self, current_time):
        if current_time - self.last_shot > self.shoot_cooldown:
            self.bullets.append({
                "x": self.player["x"],
                "y": self.player["y"] - self.player["height"]/2,
                "width": 4,
                "height": 15,
                "speed": 10
            })
            self.last_shot = current_time

    def update(self):
        current_time = time.time()
        
        # Update bullets
        for bullet in self.bullets[:]:
            bullet["y"] -= bullet["speed"]
            if bullet["y"] < 0:
                self.bullets.remove(bullet)

        # Update enemies
        for enemy in self.enemies[:]:
            enemy["y"] += enemy["speed"]
            if enemy["y"] > self.height:
                self.enemies.remove(enemy)
                self.lives -= 1

        # Check collisions
        for bullet in self.bullets[:]:
            for enemy in self.enemies[:]:
                if (bullet["x"] < enemy["x"] + enemy["width"]/2 and
                    bullet["x"] + bullet["width"] > enemy["x"] - enemy["width"]/2 and
                    bullet["y"] < enemy["y"] + enemy["height"]/2 and
                    bullet["y"] + bullet["height"] > enemy["y"] - enemy["height"]/2):
                    self.bullets.remove(bullet)
                    self.enemies.remove(enemy)
                    self.score += int(enemy["width"])
                    break

        # Spawn new enemies
        self.spawn_enemy(current_time)

    def draw(self):
        # Create image
        img = Image.new('RGB', (self.width, self.height), color='black')
        draw = ImageDraw.Draw(img)

        # Draw player
        draw.polygon([
            (self.player["x"], self.player["y"] - self.player["height"]/2),
            (self.player["x"] - self.player["width"]/2, self.player["y"] + self.player["height"]/2),
            (self.player["x"] + self.player["width"]/2, self.player["y"] + self.player["height"]/2)
        ], fill='#3498db')

        # Draw bullets
        for bullet in self.bullets:
            draw.rectangle([
                (bullet["x"] - bullet["width"]/2, bullet["y"]),
                (bullet["x"] + bullet["width"]/2, bullet["y"] + bullet["height"])
            ], fill='#ff0')

        # Draw enemies
        for enemy in self.enemies:
            draw.ellipse([
                (enemy["x"] - enemy["width"]/2, enemy["y"] - enemy["height"]/2),
                (enemy["x"] + enemy["width"]/2, enemy["y"] + enemy["height"]/2)
            ], fill='#ff5555')

        return img

# Game instance
game = SpaceShooter()

def move_player(direction):
    if game.lives <= 0:
        return game.draw(), f"Game Over! Score: {game.score}"
    
    if direction == "left" and game.player["x"] > game.player["width"]/2:
        game.player["x"] -= game.player["speed"]
    elif direction == "right" and game.player["x"] < game.width - game.player["width"]/2:
        game.player["x"] += game.player["speed"]
    elif direction == "shoot":
        game.shoot(time.time())
    
    game.update()
    status = f"Score: {game.score} | Lives: {game.lives}"
    if game.lives <= 0:
        status = f"Game Over! Score: {game.score}"
    
    return game.draw(), status

def reset_game():
    global game
    game = SpaceShooter()
    return game.draw(), f"Score: {game.score} | Lives: {game.lives}"

# Gradio interface
with gr.Blocks(title="Space Shooter") as demo:
    gr.Markdown("# Space Shooter")
    gr.Markdown("Use the buttons to move left, right, and shoot. Game updates after each action.")
    
    with gr.Row():
        game_output = gr.Image(label="Game", width=480, height=640)
        status_output = gr.Textbox(label="Status")
    
    with gr.Row():
        left_btn = gr.Button("Left")
        right_btn = gr.Button("Right")
        shoot_btn = gr.Button("Shoot")
        reset_btn = gr.Button("Reset")
    
    # Initial state
    game_output.value = game.draw()
    status_output.value = f"Score: {game.score} | Lives: {game.lives}"
    
    # Event handlers
    left_btn.click(move_player, inputs=gr.State(value="left"), outputs=[game_output, status_output])
    right_btn.click(move_player, inputs=gr.State(value="right"), outputs=[game_output, status_output])
    shoot_btn.click(move_player, inputs=gr.State(value="shoot"), outputs=[game_output, status_output])
    reset_btn.click(reset_game, inputs=None, outputs=[game_output, status_output])

demo.launch()