import pygame
import json
import os
import random
import numpy as np

POOL_FILE = "qrng_pool.json"

def load_entropy_pool():
    if not os.path.exists(POOL_FILE):
        print("⚠️ No entropy pool found!")
        return []
    with open(POOL_FILE, 'r') as f:
        data = json.load(f)
        return data.get('pool', [])

class Firework:
    def __init__(self, x, y, color, num_particles):
        self.particles = []
        for _ in range(num_particles):
            angle = random.uniform(0, 2 * np.pi)
            speed = random.uniform(2, 5)
            vx = speed * np.cos(angle)
            vy = speed * np.sin(angle)
            lifespan = random.randint(30, 90)
            self.particles.append({
                'x': x, 'y': y, 'vx': vx, 'vy': vy,
                'color': color, 'lifespan': lifespan, 'age': 0
            })

    def update(self):
        for p in self.particles:
            p['x'] += p['vx']
            p['y'] += p['vy']
            p['vy'] += 0.05  # gravity
            p['vx'] *= 0.99  # damping
            p['age'] += 1

    def draw(self, screen):
        for p in self.particles:
            if p['age'] < p['lifespan']:
                alpha = 255 * (1 - p['age'] / p['lifespan'])
                color = (*p['color'], int(alpha))
                surface = pygame.Surface((4, 4), pygame.SRCALPHA)
                pygame.draw.circle(surface, color, (2, 2), 2)
                screen.blit(surface, (p['x'], p['y']))

    def is_dead(self):
        return all(p['age'] >= p['lifespan'] for p in self.particles)

def run():
    pool = load_entropy_pool()
    if not pool:
        return

    pygame.init()
    WIDTH, HEIGHT = 800, 600
    screen = pygame.display.set_mode((WIDTH, HEIGHT), pygame.RESIZABLE)
    pygame.display.set_caption("Quantum Fireworks")

    clock = pygame.time.Clock()
    fireworks = []
    pool_index = 0

    running = True
    while running:
        screen.fill((0, 0, 0))
        WIDTH, HEIGHT = screen.get_size()

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

        # Spawn fireworks across the entire screen
        if pool_index + 3 < len(pool):
            # Spread the fireworks across the screen
            x = random.randint(0, WIDTH)
            y = random.randint(0, HEIGHT)

            # Get color from entropy pool
            hue = pool[(pool_index + 2) % len(pool)] / 255.0
            color_obj = pygame.Color(0)
            color_obj.hsva = (hue * 360, 100, 100, 100)
            color = color_obj.r, color_obj.g, color_obj.b

            # Number of particles per firework
            num_particles = 20 + (pool[(pool_index + 3) % len(pool)] % 30)
            fireworks.append(Firework(x, y, color, num_particles))
            pool_index = (pool_index + 4) % len(pool)

        for fw in fireworks[:]:
            fw.update()
            fw.draw(screen)
            if fw.is_dead():
                fireworks.remove(fw)

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

    pygame.quit()
