import pygame
import random
import math

# Initialize Pygame
pygame.init()

# Screen dimensions
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Tower Defense Game")

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
DARK_RED = (139, 0, 0)
TOWER_COLOR = (50, 50, 255)
ENEMY_COLOR = (255, 0, 0)
PLAYER_COLOR = (0, 255, 255)

# FPS
FPS = 60
clock = pygame.time.Clock()

# Player attributes
player_health = 1000
player_money = 100
tower_price = 200  # Increased spawn cost for towers

# Fonts
font = pygame.font.SysFont('Arial', 30)
button_font = pygame.font.SysFont('Arial', 25)

# Game loop flag
running = True

# Global lists for towers and enemies
towers = []
enemies = []

# Player position (center)
player_x, player_y = SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2
player_width, player_height = 50, 50

# Bullet class
class Bullet:
    def __init__(self, x, y, target_x, target_y, damage):
        self.x = x
        self.y = y
        self.width = 10
        self.height = 10
        self.damage = damage
        angle = math.atan2(target_y - y, target_x - x)
        self.speed = 15
        self.dx = math.cos(angle) * self.speed
        self.dy = math.sin(angle) * self.speed

    def move(self):
        self.x += self.dx
        self.y += self.dy

    def draw(self):
        pygame.draw.rect(screen, WHITE, (self.x, self.y, self.width, self.height))

    def collides_with(self, enemy):
        return (self.x + self.width > enemy.x and self.x < enemy.x + 30 and
                self.y + self.height > enemy.y and self.y < enemy.y + 30)

# Player (center tower)
class Player:
    def __init__(self):
        self.x = player_x
        self.y = player_y
        self.width = player_width
        self.height = player_height
        self.health = 1000
        self.attack_range = 300
        self.attack_cooldown = 200
        self.last_attack_time = 0

    def draw(self):
        pygame.draw.rect(screen, PLAYER_COLOR, (self.x - self.width // 2, self.y - self.height // 2, self.width, self.height))
        self.draw_health_bar(self.x - self.width // 2, self.y - self.height // 2 - 10)
        pygame.draw.circle(screen, (100, 255, 255), (self.x, self.y), self.attack_range, 1)

    def draw_health_bar(self, x, y):
        pygame.draw.rect(screen, DARK_RED, (x, y, self.width, 6))
        pygame.draw.rect(screen, GREEN, (x, y, (self.health / 1000) * self.width, 6))

    def shoot(self, enemies, bullets, current_time):
        if current_time - self.last_attack_time >= self.attack_cooldown:
            for enemy in enemies:
                dist = math.hypot(enemy.x - self.x, enemy.y - self.y)
                if dist <= self.attack_range:
                    bullet = Bullet(self.x, self.y, enemy.x, enemy.y, 50)
                    bullets.append(bullet)
                    self.last_attack_time = current_time
                    break

# Tower class (weaker tower at the start)
class Tower:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.radius = 100
        self.damage = 5  # Start with weaker towers
        self.cooldown = 60
        self.last_shot = 0
        self.health = 500  # Increased tower health significantly
        self.attack_speed = 1  # Initially normal attack speed

    def draw(self):
        # Draw the tower
        pygame.draw.circle(screen, TOWER_COLOR, (self.x, self.y), 40)
        pygame.draw.circle(screen, WHITE, (self.x, self.y), self.radius, 2)

        # Draw tower health bar
        self.draw_health_bar(self.x - 40, self.y - 50)

    def draw_health_bar(self, x, y):
        # Draw health bar above the tower
        pygame.draw.rect(screen, DARK_RED, (x, y, 80, 10))  # Dark red background
        pygame.draw.rect(screen, GREEN, (x, y, (self.health / 500) * 80, 10))  # Green foreground

    def shoot(self, enemies, current_time):
        if current_time - self.last_shot >= self.cooldown:
            for enemy in enemies:
                if math.hypot(self.x - enemy.x, self.y - enemy.y) <= self.radius:
                    enemy.take_damage(self.damage)
                    self.last_shot = current_time
                    break

    # Upgrade tower's stats
    def upgrade(self, stat_type):
        global player_money
        if stat_type == "damage" and player_money >= 20:
            self.damage += 5
            player_money -= 20
        elif stat_type == "cooldown" and player_money >= 20:
            self.cooldown = max(30, self.cooldown - 10)
            player_money -= 20
        elif stat_type == "range" and player_money >= 20:
            self.radius += 20
            player_money -= 20
        elif stat_type == "speed" and player_money >= 20:
            self.attack_speed = min(0.5, self.attack_speed - 0.1)
            player_money -= 20

    # Take damage
    def take_damage(self, damage):
        self.health -= damage
        return self.health <= 0

# Enemy class
class Enemy:
    def __init__(self, wave):
        # Initial enemy stats based on the wave number
        self.x = random.randint(0, SCREEN_WIDTH - 30)
        self.y = random.randint(0, SCREEN_HEIGHT - 30)
        
        # Scale enemy health and damage based on the wave
        self.health = 10 + (wave * 2)  # Start with weaker enemies
        self.attack_damage = 1 + (wave // 2)  # Attack damage increases slightly
        self.speed = 1 + (wave // 3)  # Speed increases a little as waves progress

    def move(self):
        angle = math.atan2(player_y - self.y, player_x - self.x)
        self.x += math.cos(angle) * self.speed
        self.y += math.sin(angle) * self.speed

    def draw(self):
        pygame.draw.rect(screen, ENEMY_COLOR, (self.x, self.y, 30, 30))

    def take_damage(self, damage):
        self.health -= damage
        return self.health <= 0

    def attack(self):
        global player_health
        if math.hypot(self.x - player_x, self.y - player_y) < 30:
            player_health -= self.attack_damage

# Draw HUD
def draw_player_stats(wave):
    draw_text(f"Health: {player_health}", 10, 10)
    draw_text(f"Money: ${player_money}", SCREEN_WIDTH - 150, 10)
    draw_text(f"Wave: {wave}", SCREEN_WIDTH // 2 - 40, 10)

# Draw text
def draw_text(text, x, y, color=WHITE, font_size=30):
    label = font.render(text, True, color)
    screen.blit(label, (x, y))

# Draw buttons
def draw_button(text, x, y, width, height, color, font, action=None):
    pygame.draw.rect(screen, color, (x, y, width, height))
    label = font.render(text, True, BLACK)
    screen.blit(label, (x + (width - label.get_width()) // 2, y + (height - label.get_height()) // 2))
    
    # Check if button is clicked
    mouse_x, mouse_y = pygame.mouse.get_pos()
    mouse_pressed = pygame.mouse.get_pressed()[0]
    if (x < mouse_x < x + width and y < mouse_y < y + height) and mouse_pressed:
        if action:
            action()

# Handle placing towers
def handle_tower_placement(mouse_x, mouse_y):
    global player_money
    if player_money >= tower_price:
        towers.append(Tower(mouse_x, mouse_y))
        player_money -= tower_price

# Handle upgrades (based on user input)
def handle_upgrade(tower, upgrade_type):
    tower.upgrade(upgrade_type)

# Shop Menu
shop_active = False
def toggle_shop():
    global shop_active
    shop_active = not shop_active

def draw_shop():
    pygame.draw.rect(screen, (0, 0, 0, 180), (SCREEN_WIDTH - 200, 0, 200, SCREEN_HEIGHT))  # Semi-transparent background
    draw_text("Shop", SCREEN_WIDTH - 180, 50, color=GREEN)
    draw_button(f"Buy Tower - ${tower_price}", SCREEN_WIDTH - 180, 100, 160, 50, GREEN, button_font, handle_shop_purchase)

def handle_shop_purchase():
    global player_money
    if player_money >= tower_price:
        player_money -= tower_price
        towers.append(Tower(random.randint(50, SCREEN_WIDTH - 50), random.randint(50, SCREEN_HEIGHT - 50)))  # Place a new tower randomly

# Ensure enemies spawn outside tower view
def spawn_enemy(wave):
    max_attempts = 100
    for _ in range(max_attempts):
        enemy_x = random.randint(0, SCREEN_WIDTH - 30)
        enemy_y = random.randint(0, SCREEN_HEIGHT - 30)
        enemy_spawn_valid = True

        # Ensure the spawn point is outside the radius of all towers
        for tower in towers:
            if math.hypot(tower.x - enemy_x, tower.y - enemy_y) <= tower.radius:
                enemy_spawn_valid = False
                break

        if enemy_spawn_valid:
            return Enemy(wave)

    return Enemy(wave)  # If no valid spawn point found, just spawn at random position (fallback)

# Main loop
def main():
    global running, player_health, player_money, upgrade_menu_active, selected_tower, shop_active

    player = Player()
    bullets = []
    wave = 1
    enemies_to_spawn = 2  # Start small
    for _ in range(enemies_to_spawn):
        enemies.append(spawn_enemy(wave))

    upgrade_menu_active = False
    selected_tower = None

    while running:
        screen.fill(BLACK)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 3:  # Right-click = place tower
                    mouse_x, mouse_y = pygame.mouse.get_pos()
                    handle_tower_placement(mouse_x, mouse_y)
                elif event.button == 1:  # Left-click to select tower for upgrades
                    mouse_x, mouse_y = pygame.mouse.get_pos()
                    for tower in towers:
                        if math.hypot(tower.x - mouse_x, tower.y - mouse_y) <= 40:
                            selected_tower = tower
                            upgrade_menu_active = True
                            break

            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_u:  # Press "U" to toggle shop
                    toggle_shop()

        for enemy in enemies:
            enemy.move()
            enemy.attack()

        current_time = pygame.time.get_ticks()

        # Towers shoot
        for tower in towers:
            tower.shoot(enemies, current_time)

        # Player shoots
        player.shoot(enemies, bullets, current_time)

        # Update bullets
        for bullet in bullets[:]:
            bullet.move()
            bullet.draw()
            for enemy in enemies[:]:
                if bullet.collides_with(enemy):
                    if enemy.take_damage(bullet.damage):
                        enemies.remove(enemy)
                        player_money += 10
                    if bullet in bullets:
                        bullets.remove(bullet)
                    break

        # Draw game
        draw_player_stats(wave)
        player.draw()
        for tower in towers:
            tower.draw()
        for enemy in enemies:
            enemy.draw()

        # Draw shop if active
        if shop_active:
            draw_shop()

        # Draw upgrade menu with buttons on the right side
        if upgrade_menu_active:
            pygame.draw.rect(screen, (0, 0, 0, 180), (SCREEN_WIDTH - 200, 0, 200, SCREEN_HEIGHT))  # Semi-transparent background
            draw_text("Upgrade Menu", SCREEN_WIDTH - 180, 50, color=GREEN)

            if selected_tower:
                # Draw buttons for upgrade
                draw_button("Upgrade Damage (+5)", SCREEN_WIDTH - 180, 100, 160, 50, GREEN, button_font, lambda: handle_upgrade(selected_tower, "damage"))
                draw_button("Upgrade Cooldown (-10)", SCREEN_WIDTH - 180, 160, 160, 50, GREEN, button_font, lambda: handle_upgrade(selected_tower, "cooldown"))
                draw_button("Upgrade Range (+20)", SCREEN_WIDTH - 180, 220, 160, 50, GREEN, button_font, lambda: handle_upgrade(selected_tower, "range"))
                draw_button("Upgrade Speed (Faster)", SCREEN_WIDTH - 180, 280, 160, 50, GREEN, button_font, lambda: handle_upgrade(selected_tower, "speed"))
            
            # Close Menu Button
            def close_menu():
                global upgrade_menu_active
                upgrade_menu_active = False

            draw_button("Close Menu", SCREEN_WIDTH - 180, 340, 160, 50, RED, button_font, close_menu)

        # Check for wave progression
        if len(enemies) == 0:
            wave += 1
            enemies_to_spawn += 2
            for _ in range(enemies_to_spawn):
                enemies.append(spawn_enemy(wave))

        pygame.display.update()
        clock.tick(FPS)

# Run the game
main()
pygame.quit()