import pygame
import sys
import math
import random

# -----------------------------
# Pygame and world initialization
# -----------------------------
pygame.init()
SCREEN_WIDTH, SCREEN_HEIGHT = 800, 600
BLOCK_SIZE = 40
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("2D Voxel Game with Collision Pixel Lighting (No Bouncing)")
clock = pygame.time.Clock()
font = pygame.font.Font(None, 20)

# -----------------------------
# Global Options
# -----------------------------
debug_mode = False         # Toggle full debug overlay (F3)
num_rays = 360 * 2             # Number of primary rays to cast
max_distance = 1000        # Maximum distance for a ray if nothing is hit
# Global counter for ray intersection tests (resets every frame)
ray_intersect_count = 0

# -----------------------------
# Voxel World Generation (each block gets a random color)
# -----------------------------
def generate_blocks():
    blocks = []
    cols = SCREEN_WIDTH // BLOCK_SIZE
    rows = SCREEN_HEIGHT // BLOCK_SIZE
    for i in range(cols):
        for j in range(rows):
            # 20% chance that this grid cell is a solid block.
            if random.random() < 0.2:
                rect = pygame.Rect(i * BLOCK_SIZE, j * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE)
                color = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))
                # Store block as tuple: (pygame.Rect, color)
                blocks.append((rect, color))
    return blocks

blocks = generate_blocks()

# -----------------------------
# BVH Data Structures and Build
# -----------------------------
class BVHNode:
    def __init__(self, bbox, left=None, right=None, block=None):
        self.bbox = bbox      # (min_x, min_y, max_x, max_y)
        self.left = left
        self.right = right
        self.block = block    # For leaf nodes, store the block (tuple: (pygame.Rect, color))

def rect_to_bbox(block):
    rect = block[0] if isinstance(block, tuple) else block
    return (rect.left, rect.top, rect.right, rect.bottom)

def union_bbox(b1, b2):
    x1 = min(b1[0], b2[0])
    y1 = min(b1[1], b2[1])
    x2 = max(b1[2], b2[2])
    y2 = max(b1[3], b2[3])
    return (x1, y1, x2, y2)

def build_bvh(block_list):
    if not block_list:
        return None
    if len(block_list) == 1:
        return BVHNode(rect_to_bbox(block_list[0]), block=block_list[0])
    # Sort blocks by center x-coordinate
    block_list.sort(key=lambda block: block[0].centerx)
    mid = len(block_list) // 2
    left = build_bvh(block_list[:mid])
    right = build_bvh(block_list[mid:])
    if left and right:
        bbox = union_bbox(left.bbox, right.bbox)
    elif left:
        bbox = left.bbox
    else:
        bbox = right.bbox
    return BVHNode(bbox, left, right)

bvh_root = build_bvh(blocks)

# -----------------------------
# BVH Statistics Functions
# -----------------------------
def get_bvh_stats(node):
    if node is None:
        return (0, 0)
    if node.block is not None:
        return (1, 1)
    left_count, left_depth = get_bvh_stats(node.left)
    right_count, right_depth = get_bvh_stats(node.right)
    total = 1 + left_count + right_count
    depth = 1 + max(left_depth, right_depth)
    return (total, depth)

# -----------------------------
# Ray-AABB Intersection (Slab Method)
# -----------------------------
def ray_intersect_aabb(origin, direction, bbox):
    global ray_intersect_count
    ray_intersect_count += 1

    tmin = -math.inf
    tmax = math.inf
    ox, oy = origin
    dx, dy = direction

    # X slab
    if dx != 0:
        tx1 = (bbox[0] - ox) / dx
        tx2 = (bbox[2] - ox) / dx
        tmin = max(tmin, min(tx1, tx2))
        tmax = min(tmax, max(tx1, tx2))
    else:
        if not (bbox[0] <= ox <= bbox[2]):
            return None

    # Y slab
    if dy != 0:
        ty1 = (bbox[1] - oy) / dy
        ty2 = (bbox[3] - oy) / dy
        tmin = max(tmin, min(ty1, ty2))
        tmax = min(tmax, max(ty1, ty2))
    else:
        if not (bbox[1] <= oy <= bbox[3]):
            return None

    if tmax >= tmin and tmax >= 0:
        return tmin if tmin >= 0 else tmax
    return None

# -----------------------------
# BVH Ray Casting
# -----------------------------
def ray_cast_bvh(node, origin, direction):
    if node is None:
        return None, None
    t_bbox = ray_intersect_aabb(origin, direction, node.bbox)
    if t_bbox is None:
        return None, None

    if node.block is not None:
        t_hit = ray_intersect_aabb(origin, direction, rect_to_bbox(node.block))
        if t_hit is not None:
            return t_hit, node.block
        else:
            return None, None

    t_left, block_left = ray_cast_bvh(node.left, origin, direction)
    t_right, block_right = ray_cast_bvh(node.right, origin, direction)

    if t_left is not None and t_right is not None:
        return (t_left, block_left) if t_left < t_right else (t_right, block_right)
    elif t_left is not None:
        return t_left, block_left
    elif t_right is not None:
        return t_right, block_right
    return None, None

# -----------------------------
# Single Ray Casting (No Bouncing)
# -----------------------------
def cast_ray_single(origin, direction):
    """
    Cast a single ray from origin in the given direction.
    If the ray collides with a block within max_distance, return (hit_point, block_color).
    Otherwise, return None.
    """
    t, hit_block = ray_cast_bvh(bvh_root, origin, direction)
    if t is None or t > max_distance:
        return None
    hit_point = (origin[0] + direction[0] * t,
                 origin[1] + direction[1] * t)
    block_color = hit_block[1] if isinstance(hit_block, tuple) else (255, 255, 255)
    return (hit_point, block_color)

# -----------------------------
# Advanced Debug Menu (No Background)
# -----------------------------
def draw_debug_menu(surface, fps):
    bvh_node_count, bvh_depth = get_bvh_stats(bvh_root)
    debug_lines = [
        "DEBUG MODE ACTIVE",
        f"FPS: {fps:.1f}",
        f"Blocks: {len(blocks)}",
        f"BVH Nodes: {bvh_node_count}",
        f"BVH Depth: {bvh_depth}",
        f"Light Source: {light_source}",
        f"Rays Cast: {num_rays}",
        f"Ray Intersection Tests (this frame): {ray_intersect_count}",
        f"Avg Tests per Ray: {ray_intersect_count / num_rays:.2f}",
        "Toggles: F3 - Debug, R - Regenerate World"
    ]
    y_offset = 10
    for line in debug_lines:
        text_surf = font.render(line, True, (255, 255, 255))
        surface.blit(text_surf, (10, y_offset))
        y_offset += text_surf.get_height() + 2

# -----------------------------
# Main Game Loop
# -----------------------------
light_source = (SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2)

running = True
while running:
    ray_intersect_count = 0  # Reset each frame
    
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_F3:
                debug_mode = not debug_mode
            elif event.key == pygame.K_r:
                blocks = generate_blocks()
                bvh_root = build_bvh(blocks)
        elif event.type == pygame.MOUSEMOTION:
            light_source = event.pos

    screen.fill((30, 30, 30))

    # Draw voxel blocks in their random colors.
    # for block in blocks:
    #     pygame.draw.rect(screen, block[1], block[0])

    # For each ray, cast once and mark the collision point with the block's color.
    for angle in range(0, 360, max(1, 360 // num_rays)):
        rad = math.radians(angle)
        direction = (math.cos(rad), math.sin(rad))
        result = cast_ray_single(light_source, direction)
        if result is not None:
            hit_point, hit_color = result
            # Draw a small circle (pixel) at the collision point.
            pygame.draw.circle(screen, hit_color, (int(hit_point[0]), int(hit_point[1])), 3)

    # Draw the light source.
    pygame.draw.circle(screen, (255, 255, 0), light_source, 5)

    if debug_mode:
        fps = clock.get_fps()
        draw_debug_menu(screen, fps)

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

pygame.quit()
sys.exit()