import pygame
import math
import random

# Setup
pygame.init()
WIDTH, HEIGHT = 1200, 800
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Galaxy Creator 🌌")
clock = pygame.time.Clock()

# Constants
G = 1.5
CENTER_MASS = 50000
CENTER = (WIDTH // 2, HEIGHT // 2)
TIME_STEP = 0.1

# Colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
CENTER_COLOR = (255, 180, 180)

STAR_COLORS = [(255, 255, 180), (255, 200, 200), (180, 220, 255)]
PLANET_COLORS = [(0, 200, 255), (0, 255, 100), (255, 100, 0)]
ASTEROID_COLORS = [(169, 169, 169), (105, 105, 105), (139, 69, 19)]

# Fonts
font = pygame.font.SysFont("Arial", 20)

# Object Types
class SpaceObjectType:
    def __init__(self, name, radius, colors, speed_factor, glow):
        self.name = name
        self.radius = radius
        self.colors = colors
        self.speed_factor = speed_factor
        self.glow = glow

OBJECT_TYPES = {
    pygame.K_1: SpaceObjectType("Star", 2, STAR_COLORS, 1.0, True),
    pygame.K_2: SpaceObjectType("Planet", 6, PLANET_COLORS, 0.6, False),
    pygame.K_3: SpaceObjectType("Asteroid", 1, ASTEROID_COLORS, 1.2, False),
    # Add more types here!
}

selected_type = OBJECT_TYPES[pygame.K_1]

# SpaceObject class
class SpaceObject:
    def __init__(self, x, y, obj_type: SpaceObjectType):
        self.x = x
        self.y = y
        dx = x - CENTER[0]
        dy = y - CENTER[1]
        distance = math.hypot(dx, dy)
        norm = max(distance, 0.01)

        speed = math.sqrt(G * CENTER_MASS / norm) * obj_type.speed_factor
        self.vx = -dy / norm * speed
        self.vy = dx / norm * speed

        self.radius = obj_type.radius
        self.color = random.choice(obj_type.colors)
        self.glow = obj_type.glow

    def update(self):
        dx = CENTER[0] - self.x
        dy = CENTER[1] - self.y
        distance = math.hypot(dx, dy)
        if distance < 5:
            return
        force = G * CENTER_MASS / (distance ** 2)
        ax = dx / distance * force
        ay = dy / distance * force
        self.vx += ax * TIME_STEP
        self.vy += ay * TIME_STEP
        self.x += self.vx * TIME_STEP
        self.y += self.vy * TIME_STEP

    def draw(self, surface):
        flicker = random.randint(-1, 1) if self.radius == 2 else 0
        pygame.draw.circle(surface, self.color, (int(self.x), int(self.y)), max(1, self.radius + flicker))
        if self.glow:
            glow_color = (*self.color[:3], 20)
            glow_surface = pygame.Surface((8, 8), pygame.SRCALPHA)
            pygame.draw.circle(glow_surface, glow_color, (4, 4), 4)
            surface.blit(glow_surface, (int(self.x) - 4, int(self.y) - 4))

# UI Helpers
def draw_button(surface, rect, text, hover=False):
    color = (80, 80, 80) if hover else (50, 50, 50)
    pygame.draw.rect(surface, color, rect)
    pygame.draw.rect(surface, WHITE, rect, 2)
    label = font.render(text, True, WHITE)
    label_rect = label.get_rect(center=rect.center)
    surface.blit(label, label_rect)

def draw_info(surface, selected_type):
    label = font.render(f"Current: {selected_type.name}", True, WHITE)
    surface.blit(label, (140, 30))

# Game State
objects = []
restart_rect = pygame.Rect(20, 20, 100, 40)
running = True

# Main Loop
while running:
    screen.fill(BLACK)
    pygame.draw.circle(screen, CENTER_COLOR, CENTER, 6)

    mx, my = pygame.mouse.get_pos()
    click = pygame.mouse.get_pressed()

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

        elif event.type == pygame.MOUSEBUTTONDOWN:
            if restart_rect.collidepoint(mx, my):
                objects.clear()

        elif event.type == pygame.KEYDOWN:
            if event.key in OBJECT_TYPES:
                selected_type = OBJECT_TYPES[event.key]
            elif event.key == pygame.K_s:
                for i in range(300):
                    angle = i * 0.2
                    radius = 10 + i * 1.2
                    x = CENTER[0] + math.cos(angle) * radius
                    y = CENTER[1] + math.sin(angle) * radius
                    objects.append(SpaceObject(x, y, OBJECT_TYPES[pygame.K_1]))

    # Hold mouse to continuously add objects, except on restart button
    if click[0] and not restart_rect.collidepoint(mx, my):
        objects.append(SpaceObject(mx, my, selected_type))

    draw_button(screen, restart_rect, "Restart", hover=restart_rect.collidepoint(mx, my))
    draw_info(screen, selected_type)

    for obj in objects:
        obj.update()
        obj.draw(screen)

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

pygame.quit()