import pygame
import sys
import math

# 初始化
pygame.init()
SCREEN_WIDTH = 1000
SCREEN_HEIGHT = 700
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("回合制战争兵棋推演")
clock = pygame.time.Clock()
FPS = 60

# 颜色常量
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (220, 30, 30)
BLUE = (30, 80, 220)
GREEN = (40, 180, 60)
YELLOW = (220, 200, 30)
GRAY = (80, 80, 80)

# 网格设置
GRID_SIZE = 60
GRID_X = 12
GRID_Y = 9

# 字体
font = pygame.font.SysFont("simhei", 20)
big_font = pygame.font.SysFont("simhei",28)

class Unit:
    def __init__(self, x, y, team):
        self.x = x
        self.y = y
        self.team = team # 0蓝方 1红方
        self.hp = 100
        self.max_hp = 100
        self.attack = 25
        self.move_range = 2
        self.attack_range = 2
        self.moved = False
        self.attacked = False

    def draw(self):
        cx = self.x * GRID_SIZE + GRID_SIZE//2
        cy = self.y * GRID_SIZE + GRID_SIZE//2
        r = 22
        if self.team == 0:
            pygame.draw.circle(screen, BLUE, (cx, cy), r)
        else:
            pygame.draw.circle(screen, RED, (cx, cy), r)
        # 血条
        bar_w = 36
        bar_h =5
        hp_ratio = self.hp / self.max_hp
        pygame.draw.rect(screen,GRAY,(cx-bar_w//2, cy - r -8, bar_w, bar_h))
        pygame.draw.rect(screen,GREEN,(cx-bar_w//2, cy - r -8, bar_w*hp_ratio, bar_h))

# 获取两点格子距离
def grid_dist(x1,y1,x2,y2):
    return abs(x1-x2)+abs(y1-y2)

# 初始化单位
units = []
# 蓝方
for i in range(4):
    units.append(Unit(1, 1+i, 0))
# 红方
for i in range(4):
    units.append(Unit(GRID_X-2,1+i,1))

current_team = 0
select_unit = None
move_points = []
attack_points = []

def refresh_select():
    global move_points,attack_points
    move_points.clear()
    attack_points.clear()
    if not select_unit:
        return
    u = select_unit
    # 可移动点位
    if not u.moved:
        for gx in range(GRID_X):
            for gy in range(GRID_Y):
                d = grid_dist(u.x,u.y,gx,gy)
                if 0<d <= u.move_range:
                    # 判断格子是否被占用
                    occupied = False
                    for uu in units:
                        if uu.x == gx and uu.y == gy:
                            occupied=True
                            break
                    if not occupied:
                        move_points.append((gx,gy))
    # 可攻击敌人
    if not u.attacked:
        for enemy in units:
            if enemy.team != u.team:
                d = grid_dist(u.x,u.y,enemy.x,enemy.y)
                if d <= u.attack_range:
                    attack_points.append((enemy.x,enemy.y,enemy))

running = True
while running:
    screen.fill(BLACK)
    # 绘制网格
    for gx in range(GRID_X+1):
        pygame.draw.line(screen,GRAY,(gx*GRID_SIZE,0),(gx*GRID_SIZE,GRID_Y*GRID_SIZE))
    for gy in range(GRID_Y+1):
        pygame.draw.line(screen,GRAY,(0,gy*GRID_SIZE),(GRID_X*GRID_SIZE,gy*GRID_SIZE))

    # 绘制可移动、攻击格子
    for (gx,gy) in move_points:
        rect = pygame.Rect(gx*GRID_SIZE,gy*GRID_SIZE,GRID_SIZE-2,GRID_SIZE-2)
        pygame.draw.rect(screen,(0,120,0),rect)
    for (gx,gy,enemy) in attack_points:
        rect = pygame.Rect(gx*GRID_SIZE,gy*GRID_SIZE,GRID_SIZE-2,GRID_SIZE-2)
        pygame.draw.rect(screen,(120,0,0),rect)

    # 绘制棋子
    for u in units:
        if u.hp>0:
            u.draw()
    if select_unit and select_unit.hp>0:
        sx = select_unit.x*GRID_SIZE
        sy = select_unit.y*GRID_SIZE
        pygame.draw.rect(screen,YELLOW,(sx+2,sy+2,GRID_SIZE-4,GRID_SIZE-4),3)

    # UI文字
    team_name = ["蓝方回合","红方回合"][current_team]
    t1 = big_font.render(f"当前:{team_name} | 空格切换回合",True,WHITE)
    screen.blit(t1,(20,GRID_Y*GRID_SIZE+10))
    t2 = font.render("左键选中棋子、绿格移动、红格进攻",True,WHITE)
    screen.blit(t2,(20,GRID_Y*GRID_SIZE+45))

    # 事件循环
    for e in pygame.event.get():
        if e.type == pygame.QUIT:
            running = False
        if e.type == pygame.KEYDOWN:
            if e.key == pygame.K_SPACE:
                # 切换回合，重置行动标记
                current_team = 1 - current_team
                for u in units:
                    if u.team == current_team:
                        u.moved = False
                        u.attacked = False
                select_unit = None
                refresh_select()
        if e.type == pygame.MOUSEBUTTONDOWN and e.button == 1:
            mx,my = pygame.mouse.get_pos()
            gx = mx // GRID_SIZE
            gy = my // GRID_SIZE
            # 判断点击攻击敌人
            hit_enemy = None
            for (ax,ay,en) in attack_points:
                if ax == gx and ay == gy:
                    hit_enemy = en
                    break
            if hit_enemy is not None and select_unit is not None:
                hit_enemy.hp -= select_unit.attack
                select_unit.attacked = True
                refresh_select()
                if hit_enemy.hp <= 0:
                    units.remove(hit_enemy)
                continue
            # 判断移动
            if (gx,gy) in move_points and select_unit is not None:
                select_unit.x = gx
                select_unit.y = gy
                select_unit.moved = True
                refresh_select()
                continue
            # 选中己方棋子
            new_sel = None
            for u in units:
                if u.hp>0 and u.x == gx and u.y == gy and u.team == current_team:
                    new_sel = u
                    break
            select_unit = new_sel
            refresh_select()

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

pygame.quit()
sys.exit()