import pygame
import random
import math

# ---------------- CONFIG ----------------
WIDTH, HEIGHT = 800, 600
FPS = 60

SQUARE_SIZE = 220
SQUARE_X = (WIDTH - SQUARE_SIZE) // 2
SQUARE_Y = (HEIGHT - SQUARE_SIZE) // 2

BALL_RADIUS = 10
PLAYER_POWER = 7
FRICTION = 0.99
STOP_THRESHOLD = 0.1

TOTAL_ENEMIES = 10

# Gap settings
GAP_SIZE = 50  # opening size

# ---------------- BALL CLASS ----------------
class Ball:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.vx = 0
        self.vy = 0
        self.color = color

    def move(self):
        self.x += self.vx
        self.y += self.vy

        self.vx *= FRICTION
        self.vy *= FRICTION

        if abs(self.vx) < STOP_THRESHOLD:
            self.vx = 0
        if abs(self.vy) < STOP_THRESHOLD:
            self.vy = 0

    def draw(self, screen):
        pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), BALL_RADIUS)

# ---------------- GAP CHECK ----------------
def in_gap(x, y):
    # Top gap
    if (SQUARE_X + SQUARE_SIZE//2 - GAP_SIZE//2 <= x <=
        SQUARE_X + SQUARE_SIZE//2 + GAP_SIZE//2 and
        y <= SQUARE_Y):
        return True

    # Bottom gap
    if (SQUARE_X + SQUARE_SIZE//2 - GAP_SIZE//2 <= x <=
        SQUARE_X + SQUARE_SIZE//2 + GAP_SIZE//2 and
        y >= SQUARE_Y + SQUARE_SIZE):
        return True

    # Left gap
    if (SQUARE_Y + SQUARE_SIZE//2 - GAP_SIZE//2 <= y <=
        SQUARE_Y + SQUARE_SIZE//2 + GAP_SIZE//2 and
        x <= SQUARE_X):
        return True

    # Right gap
    if (SQUARE_Y + SQUARE_SIZE//2 - GAP_SIZE//2 <= y <=
        SQUARE_Y + SQUARE_SIZE//2 + GAP_SIZE//2 and
        x >= SQUARE_X + SQUARE_SIZE):
        return True

    return False

# ---------------- COLLISION ----------------
def ball_collision(b1, b2):
    dx = b1.x - b2.x
    dy = b1.y - b2.y
    dist = math.hypot(dx, dy)

    if dist == 0 or dist > BALL_RADIUS * 2:
        return

    nx = dx / dist
    ny = dy / dist

    p = 2 * (b1.vx * nx + b1.vy * ny - b2.vx * nx - b2.vy * ny) / 2

    b1.vx -= p * nx
    b1.vy -= p * ny
    b2.vx += p * nx
    b2.vy += p * ny

    overlap = BALL_RADIUS * 2 - dist
    b1.x += nx * overlap / 2
    b1.y += ny * overlap / 2
    b2.x -= nx * overlap / 2
    b2.y -= ny * overlap / 2

# ---------------- MAIN GAME ----------------
def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Knock-Out Ball Game")
    clock = pygame.time.Clock()
    font = pygame.font.Font(None, 36)

    enemies = []
    for _ in range(TOTAL_ENEMIES):
        x = random.randint(SQUARE_X + BALL_RADIUS, SQUARE_X + SQUARE_SIZE - BALL_RADIUS)
        y = random.randint(SQUARE_Y + BALL_RADIUS, SQUARE_Y + SQUARE_SIZE - BALL_RADIUS)
        enemies.append(Ball(x, y, (255, 0, 0)))

    player = Ball(WIDTH // 2, HEIGHT - 40, (0, 0, 255))

    aiming = False
    start_pos = (0, 0)
    score = 0
    game_over = False

    running = True
    while running:
        clock.tick(FPS)

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

            if event.type == pygame.MOUSEBUTTONDOWN and not aiming and not game_over:
                aiming = True
                start_pos = (player.x, player.y)

            if event.type == pygame.MOUSEBUTTONUP and aiming:
                aiming = False
                mx, my = pygame.mouse.get_pos()
                dx = start_pos[0] - mx
                dy = start_pos[1] - my
                dist = math.hypot(dx, dy)
                if dist > 0:
                    player.vx = (dx / dist) * PLAYER_POWER
                    player.vy = (dy / dist) * PLAYER_POWER

        # ---------------- UPDATE ----------------
        player.move()
        for e in enemies:
            e.move()

        for e in enemies:
            ball_collision(player, e)

        for i in range(len(enemies)):
            for j in range(i + 1, len(enemies)):
                ball_collision(enemies[i], enemies[j])

        # Bounce enemies on walls EXCEPT gaps
        #LEFT WALL
        for e in enemies:
            if e.x - BALL_RADIUS <= SQUARE_X:
                        if not (SQUARE_Y + SQUARE_SIZE//2 - GAP_SIZE//2 <= e.y <= SQUARE_Y + SQUARE_SIZE//2 + GAP_SIZE//2):
                             e.x = SQUARE_X + BALL_RADIUS
                             e.vx *= -1

        #RIGHT WALL
              
            if e.x + BALL_RADIUS >= SQUARE_X + SQUARE_SIZE:

               if not (SQUARE_Y + SQUARE_SIZE//2 - GAP_SIZE//2 <= e.y <=
                SQUARE_Y + SQUARE_SIZE//2 + GAP_SIZE//2):
                       e.x = SQUARE_X + SQUARE_SIZE - BALL_RADIUS
                       e.vx *= -1

        # TOP WALL
            if e.y - BALL_RADIUS <= SQUARE_Y:
               if not (SQUARE_X + SQUARE_SIZE//2 - GAP_SIZE//2 <= e.x <=
                SQUARE_X + SQUARE_SIZE//2 + GAP_SIZE//2):
                   e.y = SQUARE_Y + BALL_RADIUS
                   e.vy *= -1

        # BOTTOM WALL
            # BOTTOM WALL
            if e.y + BALL_RADIUS >= SQUARE_Y + SQUARE_SIZE:
               if not (SQUARE_X + SQUARE_SIZE//2 - GAP_SIZE//2 <= e.x <=
                SQUARE_X + SQUARE_SIZE//2 + GAP_SIZE//2):
                   e.y = SQUARE_Y + SQUARE_SIZE - BALL_RADIUS
                   e.vy *= -1



        # Remove balls that escaped through gaps
        escaped = 0
        for e in enemies[:]:
            if (e.x < SQUARE_X - BALL_RADIUS or
                e.x > SQUARE_X + SQUARE_SIZE + BALL_RADIUS or
                e.y < SQUARE_Y - BALL_RADIUS or
                e.y > SQUARE_Y + SQUARE_SIZE + BALL_RADIUS):
                enemies.remove(e)
                escaped += 1

        score += escaped

        # Reset player when stopped
        if player.vx == 0 and player.vy == 0 and not aiming:
            player.x = WIDTH // 2
            player.y = HEIGHT - 40

        if len(enemies) == 0:
            game_over = True

        # ---------------- DRAW ----------------
        screen.fill((240, 240, 240))

        # Draw square with gaps
        pygame.draw.rect(screen, (0, 0, 0),
                         (SQUARE_X, SQUARE_Y, SQUARE_SIZE, SQUARE_SIZE), 2)

        # Draw gaps
        pygame.draw.rect(screen, (240, 240, 240),
            (SQUARE_X + SQUARE_SIZE//2 - GAP_SIZE//2, SQUARE_Y - 2, GAP_SIZE, 4))
        pygame.draw.rect(screen, (240, 240, 240),
            (SQUARE_X + SQUARE_SIZE//2 - GAP_SIZE//2, SQUARE_Y + SQUARE_SIZE - 2, GAP_SIZE, 4))
        pygame.draw.rect(screen, (240, 240, 240),
            (SQUARE_X - 2, SQUARE_Y + SQUARE_SIZE//2 - GAP_SIZE//2, 4, GAP_SIZE))
        pygame.draw.rect(screen, (240, 240, 240),
            (SQUARE_X + SQUARE_SIZE - 2, SQUARE_Y + SQUARE_SIZE//2 - GAP_SIZE//2, 4, GAP_SIZE))

        for e in enemies:
            e.draw(screen)

        player.draw(screen)

        if aiming:
            mx, my = pygame.mouse.get_pos()
            pygame.draw.line(screen, (0, 0, 255), start_pos, (mx, my), 2)

        text = font.render(f"Score: {score}", True, (0, 0, 0))
        screen.blit(text, (10, 10))

        if game_over:
            msg = font.render("YOU CLEARED THE BOX!", True, (0, 150, 0))
            screen.blit(msg, (WIDTH // 2 - 140, HEIGHT // 2))

        pygame.display.flip()

    pygame.quit()

if __name__ == "__main__":
    main()
