import pygame
import random
import math

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

SQUARE_SIZE = 300
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

# ---------------- 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

        # Bounce inside square
        if self.color == (255, 0, 0):
            if self.x - BALL_RADIUS <= SQUARE_X or self.x + BALL_RADIUS >= SQUARE_X + SQUARE_SIZE:
                self.vx *= -1
            if self.y - BALL_RADIUS <= SQUARE_Y or self.y + BALL_RADIUS >= SQUARE_Y + SQUARE_SIZE:
                self.vy *= -1

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

# ---------------- 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)

    # Enemy balls
    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
    player = Ball(WIDTH // 2, HEIGHT - 40, (0, 0, 255))

    aiming = False
    start_pos = (0, 0)
    game_over = False
    win = 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])

        # Remove knocked-out balls
        for e in enemies[:]:
            if (e.x < SQUARE_X or e.x > SQUARE_X + SQUARE_SIZE or
                e.y < SQUARE_Y or e.y > SQUARE_Y + SQUARE_SIZE):
                enemies.remove(e)

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

        # Win condition
        if len(enemies) == 0:
            win = True
            game_over = True

        # ---------------- DRAW ----------------
        screen.fill((240, 240, 240))
        pygame.draw.rect(screen, (0, 0, 0),
                         (SQUARE_X, SQUARE_Y, SQUARE_SIZE, SQUARE_SIZE), 2)

        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"Remaining balls: {len(enemies)}", True, (0, 0, 0))
        screen.blit(text, (10, 10))

        if game_over:
            msg = "YOU WIN!" if win else "GAME OVER"
            msg_text = font.render(msg, True, (0, 150, 0))
            screen.blit(msg_text, (WIDTH // 2 - 70, HEIGHT // 2))

        pygame.display.flip()

    pygame.quit()

if __name__ == "__main__":
    main()
