import pygame
import sys
import random
from collections import deque

# 初始化Pygame
pygame.init()

# 屏幕设置
WIDTH, HEIGHT = 900, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("战争兵棋推演")

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GRAY = (128, 128, 128)
LIGHT_GRAY = (200, 200, 200)
BLUE = (0, 0, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
YELLOW = (255, 255, 0)
BROWN = (139, 69, 19)
DARK_GREEN = (0, 100, 0)

# 修复中文显示：使用系统中文字体
def get_chinese_font(size):
    font_names = [
        "SimHei", "Microsoft YaHei", "SimSun",
        "PingFang SC", "Heiti SC",
        "WenQuanYi Micro Hei", "Noto Sans CJK SC"
    ]
    
    for font_name in font_names:
        try:
            font = pygame.font.SysFont(font_name, size)
            test_text = font.render("测试", True, BLACK)
            if test_text.get_width() > 10:
                return font
        except:
            continue
    
    print("警告：未找到支持中文的字体，中文可能显示为方块")
    return pygame.font.Font(None, size)

# 字体设置
font_large = get_chinese_font(36)
font_medium = get_chinese_font(24)
font_small = get_chinese_font(18)

# 地形类型
TERRAIN_PLAIN = 0  # 平原
TERRAIN_FOREST = 1 # 森林
TERRAIN_MOUNTAIN = 2 # 山地

# 单位类型
UNIT_INFANTRY = 0  # 步兵
UNIT_TANK = 1      # 坦克
UNIT_ARTILLERY = 2 # 炮兵

# 单位属性定义
unit_stats = {
    UNIT_INFANTRY: {"name": "步兵", "hp": 100, "attack": 30, "move": 3, "range": 1, "color": (0, 0, 200)},
    UNIT_TANK: {"name": "坦克", "hp": 150, "attack": 50, "move": 4, "range": 1, "color": (0, 100, 200)},
    UNIT_ARTILLERY: {"name": "炮兵", "hp": 80, "attack": 70, "move": 2, "range": 3, "color": (100, 0, 200)}
}

# 地形属性定义
terrain_stats = {
    TERRAIN_PLAIN: {"name": "平原", "move_cost": 1, "defense": 0, "color": (144, 238, 144)},
    TERRAIN_FOREST: {"name": "森林", "move_cost": 2, "defense": 20, "color": (34, 139, 34)},
    TERRAIN_MOUNTAIN: {"name": "山地", "move_cost": 3, "defense": 40, "color": (169, 169, 169)}
}

# 单位类
class Unit:
    def __init__(self, x, y, unit_type, player):
        self.x = x
        self.y = y
        self.unit_type = unit_type
        self.player = player
        self.stats = unit_stats[unit_type].copy()
        self.max_hp = self.stats["hp"]
        self.has_moved = False
        self.has_attacked = False
    
    def draw(self, screen, cell_size):
        x = self.x * cell_size + 50
        y = self.y * cell_size + 50
        radius = cell_size // 3
        
        # 绘制单位底色
        color = self.stats["color"] if self.player == 0 else (200, 0, 0)
        pygame.draw.circle(screen, color, (x + cell_size//2, y + cell_size//2), radius)
        pygame.draw.circle(screen, BLACK, (x + cell_size//2, y + cell_size//2), radius, 2)
        
        # 绘制生命值条
        hp_width = cell_size - 10
        hp_height = 5
        hp_x = x + 5
        hp_y = y + cell_size - 10
        hp_ratio = self.stats["hp"] / self.max_hp
        
        pygame.draw.rect(screen, RED, (hp_x, hp_y, hp_width, hp_height))
        pygame.draw.rect(screen, GREEN, (hp_x, hp_y, int(hp_width * hp_ratio), hp_height))
        
        # 绘制单位类型缩写
        text = font_small.render(self.stats["name"][0], True, WHITE)
        screen.blit(text, (x + cell_size//2 - text.get_width()//2, y + cell_size//2 - text.get_height()//2))

# 地图类
class GameMap:
    def __init__(self, width, height, cell_size):
        self.width = width
        self.height = height
        self.cell_size = cell_size
        self.terrain = [[TERRAIN_PLAIN for _ in range(width)] for _ in range(height)]
        self.generate_terrain()
    
    def generate_terrain(self):
        # 随机生成地形
        for y in range(self.height):
            for x in range(self.width):
                r = random.random()
                if r < 0.1:
                    self.terrain[y][x] = TERRAIN_FOREST
                elif r < 0.15:
                    self.terrain[y][x] = TERRAIN_MOUNTAIN
    
    def draw(self, screen):
        for y in range(self.height):
            for x in range(self.width):
                terrain = self.terrain[y][x]
                color = terrain_stats[terrain]["color"]
                rect = pygame.Rect(x*self.cell_size+50, y*self.cell_size+50, self.cell_size, self.cell_size)
                pygame.draw.rect(screen, color, rect)
                pygame.draw.rect(screen, BLACK, rect, 1)
    
    def get_move_range(self, unit):
        # 使用BFS计算可移动范围
        move_range = {}
        queue = deque()
        queue.append((unit.x, unit.y, unit.stats["move"]))
        move_range[(unit.x, unit.y)] = 0
        
        while queue:
            x, y, remaining = queue.popleft()
            
            for dx, dy in [(-1,0), (1,0), (0,-1), (0,1)]:
                nx, ny = x + dx, y + dy
                if 0 <= nx < self.width and 0 <= ny < self.height:
                    terrain = self.terrain[ny][nx]
                    cost = terrain_stats[terrain]["move_cost"]
                    if remaining >= cost and (nx, ny) not in move_range:
                        move_range[(nx, ny)] = unit.stats["move"] - remaining + cost
                        queue.append((nx, ny, remaining - cost))
        
        return move_range
    
    def get_attack_range(self, unit):
        # 计算攻击范围
        attack_range = set()
        range_val = unit.stats["range"]
        
        for dy in range(-range_val, range_val+1):
            for dx in range(-range_val, range_val+1):
                if abs(dx) + abs(dy) <= range_val and (dx != 0 or dy != 0):
                    nx, ny = unit.x + dx, unit.y + dy
                    if 0 <= nx < self.width and 0 <= ny < self.height:
                        attack_range.add((nx, ny))
        
        return attack_range

# 按钮类
class Button:
    def __init__(self, x, y, w, h, text, color=LIGHT_GRAY, hover_color=GRAY):
        self.rect = pygame.Rect(x, y, w, h)
        self.text = text
        self.color = color
        self.hover_color = hover_color
        self.is_hovered = False

    def draw(self, screen):
        current_color = self.hover_color if self.is_hovered else self.color
        pygame.draw.rect(screen, current_color, self.rect)
        pygame.draw.rect(screen, BLACK, self.rect, 2)
        text_surface = font_small.render(self.text, True, BLACK)
        screen.blit(text_surface, (self.rect.x + self.rect.width//2 - text_surface.get_width()//2,
                                  self.rect.y + self.rect.height//2 - text_surface.get_height()//2))

    def handle_event(self, event):
        if event.type == pygame.MOUSEMOTION:
            self.is_hovered = self.rect.collidepoint(event.pos)
        if event.type == pygame.MOUSEBUTTONDOWN:
            if self.rect.collidepoint(event.pos):
                return True
        return False

# 主游戏类
class Game:
    def __init__(self):
        self.cell_size = 50
        self.map_width = 12
        self.map_height = 10
        self.game_map = GameMap(self.map_width, self.map_height, self.cell_size)
        self.units = []
        self.selected_unit = None
        self.move_range = {}
        self.attack_range = set()
        self.current_player = 0  # 0: 玩家, 1: AI
        self.turn = 1
        self.game_over = False
        self.winner = -1
        
        # 初始化单位
        self.init_units()
        
        # 创建按钮
        self.end_turn_button = Button(700, 100, 150, 50, "结束回合")
        self.restart_button = Button(700, 200, 150, 50, "重新开始")
    
    def init_units(self):
        # 玩家单位（左侧）
        self.units.append(Unit(1, 2, UNIT_INFANTRY, 0))
        self.units.append(Unit(1, 4, UNIT_TANK, 0))
        self.units.append(Unit(1, 6, UNIT_INFANTRY, 0))
        self.units.append(Unit(2, 3, UNIT_ARTILLERY, 0))
        self.units.append(Unit(2, 5, UNIT_INFANTRY, 0))
        
        # AI单位（右侧）
        self.units.append(Unit(10, 2, UNIT_INFANTRY, 1))
        self.units.append(Unit(10, 4, UNIT_TANK, 1))
        self.units.append(Unit(10, 6, UNIT_INFANTRY, 1))
        self.units.append(Unit(9, 3, UNIT_ARTILLERY, 1))
        self.units.append(Unit(9, 5, UNIT_INFANTRY, 1))
    
    def get_unit_at(self, x, y):
        for unit in self.units:
            if unit.x == x and unit.y == y:
                return unit
        return None
    
    def select_unit(self, x, y):
        unit = self.get_unit_at(x, y)
        if unit and unit.player == self.current_player and not unit.has_moved:
            self.selected_unit = unit
            self.move_range = self.game_map.get_move_range(unit)
            self.attack_range = self.game_map.get_attack_range(unit)
        else:
            self.selected_unit = None
            self.move_range = {}
            self.attack_range = set()
    
    def move_unit(self, x, y):
        if self.selected_unit and (x, y) in self.move_range:
            # 检查目标位置是否有单位
            if not self.get_unit_at(x, y):
                self.selected_unit.x = x
                self.selected_unit.y = y
                self.selected_unit.has_moved = True
                self.move_range = {}
                # 更新攻击范围
                self.attack_range = self.game_map.get_attack_range(self.selected_unit)
    
    def attack_unit(self, target):
        if self.selected_unit and not self.selected_unit.has_attacked:
            if (target.x, target.y) in self.attack_range:
                # 计算伤害
                terrain = self.game_map.terrain[target.y][target.x]
                defense = terrain_stats[terrain]["defense"]
                damage = max(0, self.selected_unit.stats["attack"] - defense)
                target.stats["hp"] -= damage
                
                self.selected_unit.has_attacked = True
                self.attack_range = set()
                
                # 检查单位是否死亡
                if target.stats["hp"] <= 0:
                    self.units.remove(target)
                    self.check_game_over()
    
    def check_game_over(self):
        player_units = [u for u in self.units if u.player == 0]
        ai_units = [u for u in self.units if u.player == 1]
        
        if not player_units:
            self.game_over = True
            self.winner = 1
        elif not ai_units:
            self.game_over = True
            self.winner = 0
    
    def ai_turn(self):
        # 简单AI：每个单位移动到最近的敌方单位并攻击
        ai_units = [u for u in self.units if u.player == 1]
        player_units = [u for u in self.units if u.player == 0]
        
        for unit in ai_units:
            if not player_units:
                break
                
            # 找到最近的敌方单位
            closest = None
            min_dist = float('inf')
            for target in player_units:
                dist = abs(unit.x - target.x) + abs(unit.y - target.y)
                if dist < min_dist:
                    min_dist = dist
                    closest = target
            
            if closest:
                # 移动到最近的位置
                move_range = self.game_map.get_move_range(unit)
                best_pos = (unit.x, unit.y)
                best_dist = min_dist
                
                for (x, y) in move_range:
                    if not self.get_unit_at(x, y):
                        dist = abs(x - closest.x) + abs(y - closest.y)
                        if dist < best_dist:
                            best_dist = dist
                            best_pos = (x, y)
                
                unit.x, unit.y = best_pos
                unit.has_moved = True
                
                # 攻击如果在范围内
                attack_range = self.game_map.get_attack_range(unit)
                if (closest.x, closest.y) in attack_range and not unit.has_attacked:
                    terrain = self.game_map.terrain[closest.y][closest.x]
                    defense = terrain_stats[terrain]["defense"]
                    damage = max(0, unit.stats["attack"] - defense)
                    closest.stats["hp"] -= damage
                    unit.has_attacked = True
                    
                    if closest.stats["hp"] <= 0:
                        self.units.remove(closest)
                        player_units.remove(closest)
                        self.check_game_over()
    
    def end_turn(self):
        # 重置所有单位的移动和攻击状态
        for unit in self.units:
            unit.has_moved = False
            unit.has_attacked = False
        
        self.selected_unit = None
        self.move_range = {}
        self.attack_range = set()
        
        # 切换玩家
        self.current_player = 1 - self.current_player
        if self.current_player == 0:
            self.turn += 1
        else:
            # AI回合
            self.ai_turn()
            # AI回合结束后切换回玩家
            self.end_turn()
    
    def restart(self):
        self.__init__()
    
    def draw(self, screen):
        screen.fill(WHITE)
        
        # 绘制地图
        self.game_map.draw(screen)
        
        # 绘制可移动范围
        for (x, y) in self.move_range:
            if not self.get_unit_at(x, y):
                rect = pygame.Rect(x*self.cell_size+50, y*self.cell_size+50, self.cell_size, self.cell_size)
                pygame.draw.rect(screen, (173, 216, 230, 100), rect, 3)
        
        # 绘制可攻击范围
        for (x, y) in self.attack_range:
            if self.get_unit_at(x, y) and self.get_unit_at(x, y).player != self.current_player:
                rect = pygame.Rect(x*self.cell_size+50, y*self.cell_size+50, self.cell_size, self.cell_size)
                pygame.draw.rect(screen, (255, 100, 100, 100), rect, 3)
        
        # 绘制单位
        for unit in self.units:
            unit.draw(screen, self.cell_size)
        
        # 绘制选中的单位
        if self.selected_unit:
            x = self.selected_unit.x * self.cell_size + 50
            y = self.selected_unit.y * self.cell_size + 50
            rect = pygame.Rect(x, y, self.cell_size, self.cell_size)
            pygame.draw.rect(screen, YELLOW, rect, 3)
        
        # 绘制信息面板
        info_x = 700
        info_y = 300
        
        # 回合信息
        turn_text = font_medium.render(f"回合: {self.turn}", True, BLACK)
        screen.blit(turn_text, (info_x, info_y))
        
        player_text = font_medium.render("玩家: " + ("你" if self.current_player == 0 else "AI"), True, BLACK)
        screen.blit(player_text, (info_x, info_y + 40))
        
        # 选中单位信息
        if self.selected_unit:
            unit = self.selected_unit
            name_text = font_medium.render(f"单位: {unit.stats['name']}", True, BLACK)
            hp_text = font_small.render(f"生命: {unit.stats['hp']}/{unit.max_hp}", True, BLACK)
            attack_text = font_small.render(f"攻击: {unit.stats['attack']}", True, BLACK)
            move_text = font_small.render(f"移动: {unit.stats['move']}", True, BLACK)
            range_text = font_small.render(f"射程: {unit.stats['range']}", True, BLACK)
            
            screen.blit(name_text, (info_x, info_y + 100))
            screen.blit(hp_text, (info_x, info_y + 140))
            screen.blit(attack_text, (info_x, info_y + 170))
            screen.blit(move_text, (info_x, info_y + 200))
            screen.blit(range_text, (info_x, info_y + 230))
        
        # 绘制按钮
        self.end_turn_button.draw(screen)
        self.restart_button.draw(screen)
        
        # 绘制游戏结束信息
        if self.game_over:
            overlay = pygame.Surface((WIDTH, HEIGHT))
            overlay.set_alpha(180)
            overlay.fill(BLACK)
            screen.blit(overlay, (0, 0))
            
            result_text = font_large.render("游戏结束", True, WHITE)
            winner_text = font_large.render("胜利者: " + ("你" if self.winner == 0 else "AI"), True, WHITE)
            
            screen.blit(result_text, (WIDTH//2 - result_text.get_width()//2, HEIGHT//2 - 50))
            screen.blit(winner_text, (WIDTH//2 - winner_text.get_width()//2, HEIGHT//2 + 20))
            
            hint_text = font_medium.render("点击重新开始按钮继续", True, WHITE)
            screen.blit(hint_text, (WIDTH//2 - hint_text.get_width()//2, HEIGHT//2 + 100))
        
        # 绘制操作提示
        hint_text = font_small.render("点击己方单位选择，点击蓝色区域移动，红色区域攻击", True, GRAY)
        screen.blit(hint_text, (50, HEIGHT - 30))
    
    def handle_event(self, event):
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                pygame.quit()
                sys.exit()
        
        if self.game_over:
            if self.restart_button.handle_event(event):
                self.restart()
            return
        
        if self.end_turn_button.handle_event(event):
            self.end_turn()
            return
        
        if self.restart_button.handle_event(event):
            self.restart()
            return
        
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            x, y = event.pos
            # 转换为地图坐标
            map_x = (x - 50) // self.cell_size
            map_y = (y - 50) // self.cell_size
            
            if 0 <= map_x < self.map_width and 0 <= map_y < self.map_height:
                # 检查是否点击了敌方单位（攻击）
                target = self.get_unit_at(map_x, map_y)
                if target and target.player != self.current_player and self.selected_unit:
                    self.attack_unit(target)
                # 检查是否点击了可移动区域
                elif self.selected_unit and (map_x, map_y) in self.move_range:
                    self.move_unit(map_x, map_y)
                # 否则选择单位
                else:
                    self.select_unit(map_x, map_y)

# 主程序
def main():
    game = Game()
    clock = pygame.time.Clock()
    
    while True:
        for event in pygame.event.get():
            game.handle_event(event)
        
        game.draw(screen)
        pygame.display.flip()
        clock.tick(60)

if __name__ == "__main__":
    main()