import pygame
import sys
import random
import math
from pygame.math import Vector2

# -------------------------------
# Component System and Base Classes
# -------------------------------

class Component:
    def __init__(self, owner):
        self.owner = owner
    def update(self, dt):
        pass

# RigidBody provides linear and angular motion
class RigidBody(Component):
    def __init__(self, owner, mass=1.0):
        super().__init__(owner)
        self.mass = mass
        self.inv_mass = 1.0 / mass if mass != 0 else 0
        self.velocity = Vector2(0, 0)
        self.force = Vector2(0, 0)
        self.angle = 0.0           # in radians
        self.angular_velocity = 0.0
        self.torque = 0.0
        self.inertia = 1.0         # moment of inertia (set later)
        self.inv_inertia = 1.0
    def update(self, dt):
        # --- Linear integration ---
        acceleration = self.force * self.inv_mass
        self.velocity += acceleration * dt
        self.owner.pos += self.velocity * dt
        self.force = Vector2(0, 0)
        # --- Angular integration ---
        angular_acc = self.torque * self.inv_inertia
        self.angular_velocity += angular_acc * dt
        self.angle += self.angular_velocity * dt
        self.torque = 0.0
    def apply_force(self, force, point=None):
        self.force += force
        if point is not None:
            # Torque = r x force (scalar in 2D: r.x * force.y - r.y * force.x)
            r = point - self.owner.pos
            self.torque += r.cross(force)

# -------------------------------
# Collision Components
# -------------------------------

# Base CollisionComponent (extends Component)
class CollisionComponent(Component):
    def __init__(self, owner):
        super().__init__(owner)
    # A draw method is provided for debugging outlines.
    def draw(self, surface):
        pass

# CircleCollider extends CollisionComponent
class CircleCollider(CollisionComponent):
    def __init__(self, owner, radius):
        super().__init__(owner)
        self.radius = radius
    def get_position(self):
        return self.owner.pos
    def draw(self, surface):
        pygame.draw.circle(surface, (255,255,255), 
                           (int(self.owner.pos.x), int(self.owner.pos.y)), 
                           self.radius, 1)

# SquareCollider extends CollisionComponent
class SquareCollider(CollisionComponent):
    def __init__(self, owner, width, height):
        super().__init__(owner)
        self.width = width
        self.height = height
    def get_corners(self):
        hw = self.width / 2
        hh = self.height / 2
        # Retrieve the current rotation (in radians) from the RigidBody if present.
        rb = self.owner.get_component(RigidBody)
        angle = rb.angle if rb else 0
        # Define corners in local space.
        corners = [Vector2(-hw, -hh), Vector2(hw, -hh),
                   Vector2(hw, hh), Vector2(-hw, hh)]
        # Pygame’s rotate() expects degrees.
        rotated = [corner.rotate(math.degrees(angle)) + self.owner.pos for corner in corners]
        return rotated
    def draw(self, surface):
        corners = self.get_corners()
        points = [(p.x, p.y) for p in corners]
        pygame.draw.polygon(surface, (255,255,255), points, 1)

# PolygonCollider extends CollisionComponent
class PolygonCollider(CollisionComponent):
    def __init__(self, owner, points):
        super().__init__(owner)
        # points is a list of Vector2 objects, defined in local space.
        self.points = points
    def get_transformed_points(self):
        rb = self.owner.get_component(RigidBody)
        angle = rb.angle if rb else 0
        transformed = [p.rotate(math.degrees(angle)) + self.owner.pos for p in self.points]
        return transformed
    def draw(self, surface):
        points = self.get_transformed_points()
        pts = [(p.x, p.y) for p in points]
        pygame.draw.polygon(surface, (255,255,255), pts, 1)

# -------------------------------
# GameObject: Entities in the Scene
# -------------------------------

class GameObject:
    def __init__(self, pos):
        self.pos = Vector2(pos)
        self.components = []
    def add_component(self, comp):
        self.components.append(comp)
    def get_component(self, comp_type):
        for comp in self.components:
            if isinstance(comp, comp_type):
                return comp
        return None
    def update(self, dt):
        for comp in self.components:
            comp.update(dt)


    def draw(self, surface):
        # Draw any collision shapes (for visualization)
        for comp in self.components:
            if isinstance(comp, CollisionComponent):
                comp.draw(surface)

# -------------------------------
# Collision Detection Functions
# -------------------------------

def circle_vs_circle(circleA, circleB):
    posA = circleA.get_position()
    posB = circleB.get_position()
    rA = circleA.radius
    rB = circleB.radius
    diff = posB - posA
    dist = diff.length()
    if dist < (rA + rB):
        penetration = (rA + rB) - dist
        normal = diff.normalize() if dist != 0 else Vector2(1, 0)
        return normal, penetration
    return None, None

def circle_vs_square(circle, square):
    center = circle.get_position()
    # Obtain the square’s rotation angle (in radians)
    rb = square.owner.get_component(RigidBody)
    angle = rb.angle if rb else 0
    # Transform the circle center into the square's local space.
    local_center = (center - square.owner.pos).rotate(-math.degrees(angle))
    hw = square.width / 2
    hh = square.height / 2
    closest_x = max(-hw, min(local_center.x, hw))
    closest_y = max(-hh, min(local_center.y, hh))
    closest = Vector2(closest_x, closest_y)
    diff = local_center - closest
    dist = diff.length()
    if dist < circle.radius:
        penetration = circle.radius - dist
        local_normal = diff.normalize() if dist != 0 else Vector2(1, 0)
        # Rotate the normal back to world space.
        normal = local_normal.rotate(math.degrees(angle))
        return normal, penetration
    return None, None

def square_vs_square(sq1, sq2):
    corners1 = sq1.get_corners()
    corners2 = sq2.get_corners()
    axes = []
    def get_axes(corners):
        axes = []
        for i in range(len(corners)):
            edge = corners[(i+1) % len(corners)] - corners[i]
            if edge.length() == 0:
                continue
            normal = Vector2(-edge.y, edge.x).normalize()
            axes.append(normal)
        return axes
    axes.extend(get_axes(corners1))
    axes.extend(get_axes(corners2))
    min_overlap = float('inf')
    collision_normal = None
    for axis in axes:
        proj1 = [p.dot(axis) for p in corners1]
        proj2 = [p.dot(axis) for p in corners2]
        min1, max1 = min(proj1), max(proj1)
        min2, max2 = min(proj2), max(proj2)
        if max1 < min2 or max2 < min1:
            return None, None  # Separating axis found; no collision.
        overlap = min(max1, max2) - max(min1, min2)
        if overlap < min_overlap:
            min_overlap = overlap
            collision_normal = axis
    # Ensure the normal points from sq1 to sq2.
    if (sq2.owner.pos - sq1.owner.pos).dot(collision_normal) < 0:
        collision_normal = -collision_normal
    return collision_normal, min_overlap

# Universal collision check selects routine based on collider types.
def check_collision(objA, objB):
    colA = objA.get_component(CollisionComponent)
    colB = objB.get_component(CollisionComponent)
    if colA is None or colB is None:
        return None, None
    # Circle vs Circle.
    if isinstance(colA, CircleCollider) and isinstance(colB, CircleCollider):
        return circle_vs_circle(colA, colB)
    # Circle vs Square.
    if isinstance(colA, CircleCollider) and isinstance(colB, SquareCollider):
        return circle_vs_square(colA, colB)
    if isinstance(colA, SquareCollider) and isinstance(colB, CircleCollider):
        normal, penetration = circle_vs_square(colB, colA)
        if normal is not None:
            return -normal, penetration
    # Square vs Square.
    if isinstance(colA, SquareCollider) and isinstance(colB, SquareCollider):
        return square_vs_square(colA, colB)
    # (PolygonCollider collisions can be implemented similarly using SAT.)
    return None, None

# -------------------------------
# Collision Resolution
# -------------------------------

def resolve_collision(objA, objB, normal, penetration):
    rbA = objA.get_component(RigidBody)
    rbB = objB.get_component(RigidBody)
    if rbA is None or rbB is None:
        return
    total_inv_mass = rbA.inv_mass + rbB.inv_mass
    if total_inv_mass == 0:
        return
    # Correct positions (simple penetration resolution).
    correction = normal * (penetration / total_inv_mass * 0.5)
    objA.pos -= correction * rbA.inv_mass
    objB.pos += correction * rbB.inv_mass

    # Relative velocity along collision normal.
    relative_vel = rbB.velocity - rbA.velocity
    vel_along_normal = relative_vel.dot(normal)
    if vel_along_normal > 0:
        return  # They are separating.
    restitution = 0.5  # Bounce factor.
    impulse_scalar = -(1 + restitution) * vel_along_normal / total_inv_mass
    impulse = normal * impulse_scalar

    rbA.velocity -= impulse * rbA.inv_mass
    rbB.velocity += impulse * rbB.inv_mass

    # --- Angular response ---
    # As an approximation, we use the average of the object centers as the contact point.
    contact_point = (objA.pos + objB.pos) * 0.5
    rA = contact_point - objA.pos
    rB = contact_point - objB.pos
    # In 2D, the scalar "cross" is computed with Vector2.cross().
    rbA.torque -= rA.cross(impulse)
    rbB.torque += rB.cross(impulse)

# -------------------------------
# Utility: Set the Moment of Inertia Based on the Collider Type
# -------------------------------

def set_inertia(obj):
    rb = obj.get_component(RigidBody)
    if rb is None:
        return
    col = obj.get_component(CollisionComponent)
    if isinstance(col, CircleCollider):
        rb.inertia = 0.5 * rb.mass * (col.radius ** 2)
    elif isinstance(col, SquareCollider):
        rb.inertia = rb.mass * ((col.width ** 2 + col.height ** 2) / 12)
    else:
        rb.inertia = rb.mass  # Default value.
    rb.inv_inertia = 1.0 / rb.inertia if rb.inertia != 0 else 0

# -------------------------------
# Object Factories
# -------------------------------

def create_circle(pos, radius, mass):
    obj = GameObject(pos)
    rb = RigidBody(obj, mass)
    obj.add_component(rb)
    col = CircleCollider(obj, radius)
    obj.add_component(col)
    set_inertia(obj)
    return obj

def create_square(pos, width, height, mass):
    obj = GameObject(pos)
    rb = RigidBody(obj, mass)
    obj.add_component(rb)
    col = SquareCollider(obj, width, height)
    obj.add_component(col)
    set_inertia(obj)
    return obj

# -------------------------------
# Main Game Loop and Setup
# -------------------------------

pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Physics Engine with Rotating Collisions")
clock = pygame.time.Clock()

# Create a list of objects.
objects = []

# Create some circles.
for _ in range(3):
    pos = (random.randint(100, 700), random.randint(100, 300))
    circle = create_circle(pos, random.randint(15, 30), mass=1)
    # Give each circle an initial random velocity.
    circle.get_component(RigidBody).velocity = Vector2(random.uniform(-100, 100), random.uniform(-100, 100))
    objects.append(circle)

# Create some squares ("cubes").
for _ in range(3):
    pos = (random.randint(100, 700), random.randint(100, 300))
    square = create_square(pos, random.randint(30, 50), random.randint(30, 50), mass=2)
    square.get_component(RigidBody).velocity = Vector2(random.uniform(-100, 100), random.uniform(-100, 100))
    objects.append(square)

# Gravity (pixels per second^2).
GRAVITY = Vector2(0, 200)

# Main loop.
running = True
while running:
    dt = clock.tick(60) / 1000.0  # Delta time in seconds.
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Apply gravity to each object.
    for obj in objects:
        rb = obj.get_component(RigidBody)
        if rb:
            rb.apply_force(GRAVITY * rb.mass)

    # Update all objects.
    for obj in objects:
        obj.update(dt)

    # Check collisions between all pairs.
    for i in range(len(objects)):
        for j in range(i + 1, len(objects)):
            normal, penetration = check_collision(objects[i], objects[j])
            if normal is not None:
                resolve_collision(objects[i], objects[j], normal, penetration)

    # Rendering.
    screen.fill((30, 30, 30))
    for obj in objects:
        obj.draw(screen)
    pygame.display.flip()

pygame.quit()
sys.exit()