import pygame
import sys

# ===================== 全局配置 =====================
pygame.init()

# 窗口设置
SCREEN_WIDTH = 960
SCREEN_HEIGHT = 640
CELL_SIZE = 50
MAP_COLS = 16
MAP_ROWS = 12

# 颜色定义
BG_COLOR = (30, 30, 30)
GRID_COLOR = (80, 80, 80)
MOVE_COLOR = (220, 220, 0)
RED = (200, 30, 30)
BLUE = (30, 30, 200)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GRAY = (120, 120, 120)
HP_GREEN = (0, 200, 0)
HP_RED = (200, 0, 0)

# 窗口 & 时钟
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("2D War Board Game")
clock = pygame.time.Clock()
fps = 60

# 修复字体问题：使用 pygame 内置字体，不调用系统字体
font_small = pygame.font.Font(pygame.font.get_default_font(), 20)
font_mid = pygame.font.Font(pygame.font.get_default_font(), 28)
font_big = pygame.font.Font(pygame.font.get_default_font(), 48)

# 兵种属性
UNIT_INFO = {
    "infantry":  {"hp": 100, "atk": 18, "move": 2, "name": "Inf"},
    "cavalry":   {"hp": 80,  "atk": 28, "move": 4, "name": "Cav"},
    "artillery": {"hp": 60,  "atk": 42, "move": 1, "name": "Art"}
}

# ===================== 作战单位类 =====================
class GameUnit:
    def __init__(self, team, unit_type, col, row):
        self.team = team
        self.type = unit_type
        self.col = col
        self.row = row
        cfg = UNIT_INFO[unit_type]

        self.max_hp = cfg["hp"]
        self.hp = cfg["hp"]
        self.atk = cfg["atk"]
        self.move_range = cfg["move"]
        self.name = cfg["name"]

        self.has_moved = False
        self.has_attacked = False

    def draw(self, surf):
        cx = self.col * CELL_SIZE + CELL_SIZE // 2
        cy = self.row * CELL_SIZE + CELL_SIZE // 2
        color = RED if self.team == "red" else BLUE

        # 绘制单位本体
        pygame.draw.circle(surf, color, (cx, cy), CELL_SIZE // 2 - 5)
        pygame.draw.circle(surf, WHITE, (cx, cy), CELL_SIZE // 2 - 5, 2)

        # 单位名称
        txt = font_small.render(self.name, True, WHITE)
        surf.blit(txt, (cx - txt.get_width() // 2, cy - 12))

        # 血条
        bar_w = 36
        bar_h = 6
        hp_rate = self.hp / self.max_hp
        pygame.draw.rect(surf, GRAY, (cx - bar_w//2, cy + 14, bar_w, bar_h))
        hp_color = HP_GREEN if hp_rate > 0.3 else HP_RED
        pygame.draw.rect(surf, hp_color, (cx - bar_w//2, cy + 14, bar_w * hp_rate, bar_h))

# ===================== 游戏主逻辑类 =====================
class WarGame:
    def __init__(self):
        self.units = []
        self.current_team = "red"
        self.selected_unit = None
        self.move_area = []
        self.spawn_units()

    def spawn_units(self):
        # 红方部队
        self.units.append(GameUnit("red", "infantry", 2, 3))
        self.units.append(GameUnit("red", "infantry", 2, 7))
        self.units.append(GameUnit("red", "cavalry", 3, 5))
        self.units.append(GameUnit("red", "artillery", 1, 9))

        # 蓝方部队
        self.units.append(GameUnit("blue", "infantry", 13, 3))
        self.units.append(GameUnit("blue", "infantry", 13, 7))
        self.units.append(GameUnit("blue", "cavalry", 12, 5))
        self.units.append(GameUnit("blue", "artillery", 14, 9))

    def get_unit_at(self, col, row):
        for u in self.units:
            if u.col == col and u.row == row:
                return u
        return None

    def calc_move_area(self, unit):
        self.move_area.clear()
        rng = unit.move_range
        for c in range(MAP_COLS):
            for r in range(MAP_ROWS):
                dist = abs(c - unit.col) + abs(r - unit.row)
                if 0 < dist <= rng and not self.get_unit_at(c, r):
                    self.move_area.append((c, r))

    def switch_turn(self):
        self.current_team = "blue" if self.current_team == "red" else "red"
        for u in self.units:
            u.has_moved = False
            u.has_attacked = False
        self.selected_unit = None
        self.move_area.clear()

    def check_win(self):
        red_alive = any(u.team == "red" for u in self.units)
        blue_alive = any(u.team == "blue" for u in self.units)
        if not red_alive:
            return "Blue Win"
        if not blue_alive:
            return "Red Win"
        return None

    def handle_click(self, mx, my):
        col = mx // CELL_SIZE
        row = my // CELL_SIZE

        # 移动单位
        if self.selected_unit and (col, row) in self.move_area:
            self.selected_unit.col = col
            self.selected_unit.row = row
            self.selected_unit.has_moved = True
            self.selected_unit = None
            self.move_area.clear()
            return

        click_unit = self.get_unit_at(col, row)
        if click_unit:
            # 攻击敌方
            if click_unit.team != self.current_team:
                if self.selected_unit and not self.selected_unit.has_attacked:
                    dist = abs(self.selected_unit.col - col) + abs(self.selected_unit.row - row)
                    if dist <= 2:
                        click_unit.hp -= self.selected_unit.atk
                        self.selected_unit.has_attacked = True
                        if click_unit.hp <= 0:
                            self.units.remove(click_unit)
                        self.selected_unit = None
                        self.move_area.clear()
                return
            # 选中己方单位
            if not click_unit.has_moved:
                self.selected_unit = click_unit
                self.calc_move_area(click_unit)
                return

        # 点击空白取消选中
        self.selected_unit = None
        self.move_area.clear()

    def draw(self):
        screen.fill(BG_COLOR)

        # 绘制网格线
        for c in range(MAP_COLS + 1):
            pygame.draw.line(screen, GRID_COLOR, (c*CELL_SIZE, 0), (c*CELL_SIZE, MAP_ROWS*CELL_SIZE))
        for r in range(MAP_ROWS + 1):
            pygame.draw.line(screen, GRID_COLOR, (0, r*CELL_SIZE), (MAP_COLS*CELL_SIZE, r*CELL_SIZE))

        # 绘制可移动区域
        for (c, r) in self.move_area:
            rect = (c*CELL_SIZE+2, r*CELL_SIZE+2, CELL_SIZE-4, CELL_SIZE-4)
            pygame.draw.rect(screen, MOVE_COLOR, rect, 2)

        # 绘制所有单位
        for u in self.units:
            u.draw(screen)

        # 右侧UI文字
        turn_text = f"Turn: {'Red Team' if self.current_team == 'red' else 'Blue Team'}"
        t1 = font_mid.render(turn_text, True, WHITE)
        screen.blit(t1, (MAP_COLS*CELL_SIZE + 10, 20))

        hint1 = font_small.render("Click unit -> Move (Yellow)", True, WHITE)
        hint2 = font_small.render("Click enemy to attack", True, WHITE)
        hint3 = font_small.render("SPACE = End Turn", True, WHITE)
        screen.blit(hint1, (MAP_COLS*CELL_SIZE + 10, 80))
        screen.blit(hint2, (MAP_COLS*CELL_SIZE + 10, 110))
        screen.blit(hint3, (MAP_COLS*CELL_SIZE + 10, 140))

        # 胜负提示
        win_state = self.check_win()
        if win_state:
            if win_state == "Red Win":
                win_txt = font_big.render("RED TEAM WIN!", True, RED)
            else:
                win_txt = font_big.render("BLUE TEAM WIN!", True, BLUE)
            screen.blit(win_txt, (SCREEN_WIDTH//2 - win_txt.get_width()//2, SCREEN_HEIGHT//2 - 30))

# ===================== 主循环 =====================
def main():
    game = WarGame()
    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 event.button == 1:
                mx, my = pygame.mouse.get_pos()
                game.handle_click(mx, my)
            # 空格切换回合
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    game.switch_turn()

        game.draw()
        pygame.display.flip()

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()