import sys
import pygame
from game.config import (
    WINDOW_WIDTH, WINDOW_HEIGHT, FPS,
    COLOR_BG, COLOR_WHITE, INITIAL_GOLD, INITIAL_LIVES,
    UI_TOP_HEIGHT, UI_BOTTOM_HEIGHT, GAME_HEIGHT, TOWER_DATA,
    LEVELS,
)
from game.map import GameMap
from game.enemy import Enemy
from game.tower import Tower
from game.projectile import Projectile
from game.wave import WaveManager
from game.ui import UI


class Game:
    def __init__(self):
        pygame.init()
        self.screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
        pygame.display.set_caption("塔防游戏")
        self.clock = pygame.time.Clock()
        self.font = pygame.font.SysFont("microsoftyahei", 20)
        self.small_font = pygame.font.SysFont("microsoftyahei", 16)
        self.big_font = pygame.font.SysFont("microsoftyahei", 48)
        self.title_font = pygame.font.SysFont("microsoftyahei", 36)
        self.state = "select"
        self.current_level_idx = 0
        self.ui = UI()

    def _start_level(self, level_idx):
        level = LEVELS[level_idx]
        self.current_level_idx = level_idx
        self.game_map = GameMap(level["waypoints"], level["color"])
        self.enemies = []
        self.towers = []
        self.projectiles = []
        self.wave_mgr = WaveManager(level["waves"], level["waypoints"])
        self.ui = UI()
        self.gold = level["start_gold"]
        self.lives = level["start_lives"]
        self.game_over = False
        self.won = False
        self.state = "playing"

    def run(self):
        while True:
            dt = self.clock.tick(FPS) / 1000.0
            dt = min(dt, 0.05)

            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()
                self._handle_event(event)

            if self.state == "playing" and not self.game_over:
                self._update(dt)

            self._draw()

    def _handle_event(self, event):
        if self.state == "select":
            self._handle_select_event(event)
            return

        if self.state == "playing":
            self._handle_playing_event(event)
            return

        if self.state == "game_over":
            self._handle_game_over_event(event)

    def _handle_select_event(self, event):
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            pos = event.pos
            for i in range(len(LEVELS)):
                btn_rect = self._level_button_rect(i)
                if btn_rect.collidepoint(pos):
                    self._start_level(i)
                    return

    def _handle_playing_event(self, event):
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            pos = event.pos

            action = self.ui.handle_click(pos, self.gold)
            if action == "wave":
                if self.wave_mgr.has_more_waves() and not self.wave_mgr.wave_active:
                    self.wave_mgr.start_next_wave()
                return

            tower_type = self.ui.handle_grid_click(pos)
            if tower_type:
                col, row = self.game_map.pixel_to_grid(*pos)
                if self.game_map.is_buildable(col, row):
                    occupied = any(t.col == col and t.row == row for t in self.towers)
                    price = TOWER_DATA[tower_type]["price"]
                    if not occupied and self.gold >= price:
                        px, py = self.game_map.grid_to_pixel(col, row)
                        self.towers.append(Tower(tower_type, col, row, px, py))
                        self.gold -= price
                        self.ui.selected_tower = None
                return

            self.ui.selected_tower = None

        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                self.ui.selected_tower = None

    def _handle_game_over_event(self, event):
        if event.type == pygame.KEYDOWN and event.key == pygame.K_r:
            self._start_level(self.current_level_idx)
            return

        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            pos = event.pos
            # Restart button
            restart_rect = pygame.Rect(WINDOW_WIDTH // 2 - 80, WINDOW_HEIGHT // 2 + 50, 160, 45)
            if restart_rect.collidepoint(pos):
                self._start_level(self.current_level_idx)
                return

            # Next level button (win only, not last level)
            if self.won and self.current_level_idx < len(LEVELS) - 1:
                next_rect = pygame.Rect(WINDOW_WIDTH // 2 - 80, WINDOW_HEIGHT // 2 + 105, 160, 45)
                if next_rect.collidepoint(pos):
                    self._start_level(self.current_level_idx + 1)
                    return

            # Back to select button
            select_rect = pygame.Rect(WINDOW_WIDTH // 2 - 80, WINDOW_HEIGHT // 2 + 160, 160, 45)
            if select_rect.collidepoint(pos):
                self.state = "select"
                return

    def _level_button_rect(self, idx):
        bw, bh = 300, 100
        col = idx % 3
        total_w = 3 * bw + 2 * 30
        start_x = (WINDOW_WIDTH - total_w) // 2
        start_y = 250
        return pygame.Rect(start_x + col * (bw + 30), start_y, bw, bh)

    def _update(self, dt):
        if self.ui.selected_tower is not None:
            dt *= 0.3

        self.wave_mgr.update(dt, self.enemies)

        for e in self.enemies:
            e.update(dt, self.enemies)
            if e.reached_end:
                self.lives -= 2 if e.type == "boss" else 1
                if self.lives <= 0:
                    self.lives = 0
                    self.game_over = True
                    self.won = False
                    self.state = "game_over"

        for t in self.towers:
            t.update(dt, self.enemies, self.projectiles)

        for p in self.projectiles:
            p.update(dt, self.enemies)

        for e in self.enemies:
            if not e.alive and not e.reached_end:
                self.gold += e.reward
                e.reward = 0

        self.enemies = [e for e in self.enemies if e.alive]
        self.projectiles = [p for p in self.projectiles if p.alive]

        total_waves = self.wave_mgr.total_waves
        if self.wave_mgr.all_waves_done or (self.wave_mgr.current_wave >= total_waves and not self.wave_mgr.wave_active and not self.enemies):
            if not self.game_over:
                self.game_over = True
                self.won = True
                self.state = "game_over"

    def _draw(self):
        if self.state == "select":
            self._draw_select()
            return

        # Playing / game_over
        self.screen.fill(COLOR_BG)
        self.game_map.draw(self.screen)

        for t in self.towers:
            t.draw(self.screen)

        for e in self.enemies:
            e.draw(self.screen)

        for p in self.projectiles:
            p.draw(self.screen)

        mx, my = pygame.mouse.get_pos()
        self.ui.draw_placement_preview(self.screen, mx, my, self.game_map)

        self.ui.draw(
            self.screen, self.gold, self.lives,
            self.wave_mgr.current_wave, self.wave_mgr.wave_active,
            self.wave_mgr.has_more_waves(),
            self.font, self.small_font,
            self.wave_mgr.total_waves,
        )

        self.ui.draw_slow_overlay(self.screen)

        if self.state == "game_over":
            self._draw_game_over()

        pygame.display.flip()

    def _draw_select(self):
        self.screen.fill(COLOR_BG)

        title = self.title_font.render("塔防游戏 - 选择关卡", True, COLOR_WHITE)
        self.screen.blit(title, title.get_rect(center=(WINDOW_WIDTH // 2, 80)))

        subtitle = self.small_font.render("点击关卡开始游戏", True, (180, 180, 180))
        self.screen.blit(subtitle, subtitle.get_rect(center=(WINDOW_WIDTH // 2, 130)))

        level_colors = [(60, 100, 60), (120, 100, 50), (60, 50, 80)]
        for i, level in enumerate(LEVELS):
            btn_rect = self._level_button_rect(i)
            color = level_colors[i % len(level_colors)]
            pygame.draw.rect(self.screen, color, btn_rect, border_radius=10)
            pygame.draw.rect(self.screen, COLOR_WHITE, btn_rect, 2, border_radius=10)

            name_text = self.font.render(level["name"], True, COLOR_WHITE)
            self.screen.blit(name_text, name_text.get_rect(center=(btn_rect.centerx, btn_rect.y + 25)))

            info = f"{len(level['waves'])}波 | {level['start_gold']}金 | {level['start_lives']}命"
            info_text = self.small_font.render(info, True, (200, 200, 200))
            self.screen.blit(info_text, info_text.get_rect(center=(btn_rect.centerx, btn_rect.y + 55)))

            diff = ["简单", "中等", "困难"][i]
            diff_colors = [COLOR_WHITE, (255, 200, 50), (255, 80, 80)][i]
            diff_text = self.small_font.render(diff, True, diff_colors)
            self.screen.blit(diff_text, diff_text.get_rect(center=(btn_rect.centerx, btn_rect.y + 80)))

        pygame.display.flip()

    def _draw_game_over(self):
        overlay = pygame.Surface((WINDOW_WIDTH, WINDOW_HEIGHT), pygame.SRCALPHA)
        overlay.fill((0, 0, 0, 150))
        self.screen.blit(overlay, (0, 0))

        text = "胜利!" if self.won else "失败!"
        color = (255, 215, 0) if self.won else (220, 50, 50)
        rendered = self.big_font.render(text, True, color)
        self.screen.blit(rendered, rendered.get_rect(center=(WINDOW_WIDTH // 2, WINDOW_HEIGHT // 2 - 40)))

        cx = WINDOW_WIDTH // 2

        # Restart button
        restart_rect = pygame.Rect(cx - 80, WINDOW_HEIGHT // 2 + 50, 160, 45)
        pygame.draw.rect(self.screen, (80, 120, 80), restart_rect, border_radius=8)
        pygame.draw.rect(self.screen, COLOR_WHITE, restart_rect, 2, border_radius=8)
        r_text = self.font.render("重新开始", True, COLOR_WHITE)
        self.screen.blit(r_text, r_text.get_rect(center=restart_rect.center))

        # Next level button (only on win, not last level)
        if self.won and self.current_level_idx < len(LEVELS) - 1:
            next_rect = pygame.Rect(cx - 80, WINDOW_HEIGHT // 2 + 105, 160, 45)
            pygame.draw.rect(self.screen, (50, 80, 150), next_rect, border_radius=8)
            pygame.draw.rect(self.screen, COLOR_WHITE, next_rect, 2, border_radius=8)
            n_text = self.font.render("下一关", True, COLOR_WHITE)
            self.screen.blit(n_text, n_text.get_rect(center=next_rect.center))

        # Back to select
        select_rect = pygame.Rect(cx - 80, WINDOW_HEIGHT // 2 + 160, 160, 45)
        pygame.draw.rect(self.screen, (100, 100, 100), select_rect, border_radius=8)
        pygame.draw.rect(self.screen, COLOR_WHITE, select_rect, 2, border_radius=8)
        s_text = self.font.render("选择关卡", True, COLOR_WHITE)
        self.screen.blit(s_text, s_text.get_rect(center=select_rect.center))

        hint = self.small_font.render("或按 R 重新开始", True, (180, 180, 180))
        self.screen.blit(hint, hint.get_rect(center=(cx, WINDOW_HEIGHT // 2 + 225)))


if __name__ == "__main__":
    Game().run()