import pygame
import random

# ---------- 基础 ----------
pygame.init()
TILE = 64
COLS, ROWS = 12, 9
WIDTH = COLS * TILE
HEIGHT = ROWS * TILE + 60
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("科幻兵棋 · 完整版（含 AI）")
clock = pygame.time.Clock()
font = pygame.font.SysFont(None, 22)

# ---------- 单位 ----------
UNITS = {
    "INF":    {"hp": 80,  "atk": 25, "defense": 8,  "ap": 4, "rng": 2},
    "MECH":   {"hp": 160, "atk": 40, "defense": 18, "ap": 2, "rng": 1},
    "SNIPER": {"hp": 70,  "atk": 55, "defense": 5,  "ap": 3, "rng": 4},
}

# ---------- 地形 ----------
TERRAIN = {
    0: ("平原", 1, (40, 80, 40)),
    1: ("掩体", 2, (120, 120, 120)),
    2: ("废墟", 2, (90, 90, 90)),
    3: ("真空", -1, (10, 10, 30)),
}

# ---------- 地图 ----------
game_map = [[0 for _ in range(COLS)] for _ in range(ROWS)]
game_map[4][4] = 1
game_map[5][5] = 2
game_map[6][3] = 3

# ---------- 单位类 ----------
class Unit:
    def __init__(self, x, y, team, utype):
        self.x, self.y = x, y
        self.team = team
        self.type = utype
        d = UNITS[utype]
        self.hp = self.max_hp = d["hp"]
        self.atk = d["atk"]
        self.defense = d["defense"]
        self.ap = d["ap"]
        self.rng = d["rng"]

    def draw(self):
        cx = self.x * TILE + TILE // 2
        cy = self.y * TILE + TILE // 2
        color = (0, 180, 255) if self.team == 0 else (255, 60, 60)
        pygame.draw.rect(screen, color, (cx - 22, cy - 22, 44, 44))
        pygame.draw.rect(screen, (0, 0, 0), (cx - 22, cy - 22, 44, 44), 2)

        hp_w = int(40 * self.hp / self.max_hp)
        pygame.draw.rect(screen, (0, 255, 0), (cx - 20, cy - 30, hp_w, 5))

# ---------- 工具 ----------
def unit_at(x, y):
    return next((u for u in units if u.x == x and u.y == y), None)

def can_move(u, x, y):
    if not (0 <= x < COLS and 0 <= y < ROWS):
        return False
    if TERRAIN[game_map[y][x]][1] == -1:
        return False
    if unit_at(x, y):
        return False
    return abs(u.x - x) + abs(u.y - y) <= u.ap

def hit_chance(a, d):
    dist = abs(a.x - d.x) + abs(a.y - d.y)
    return max(10, 85 - dist * 5 - TERRAIN[game_map[d.y][d.x]][1])

def combat(a, d):
    if random.randint(1, 100) <= hit_chance(a, d):
        dmg = max(1, a.atk - d.defense)
        d.hp -= dmg
        print("命中！伤害", dmg)
    else:
        print("未命中")

# ---------- 敌方 AI ----------
def enemy_ai(u):
    targets = [t for t in units if t.team == 0]
    if not targets:
        return

    target = min(targets, key=lambda t: abs(t.x - u.x) + abs(t.y - u.y))
    dist = abs(target.x - u.x) + abs(target.y - u.y)

    # 狙击手拉开
    if u.type == "SNIPER" and dist <= 2:
        for dx, dy in [(1,0),(-1,0),(0,1),(0,-1)]:
            nx, ny = u.x + dx, u.y + dy
            if 0 <= nx < COLS and 0 <= ny < ROWS:
                if TERRAIN[game_map[ny][nx]][1] != -1 and not unit_at(nx, ny):
                    u.x, u.y = nx, ny
                    return

    if dist <= u.rng:
        combat(u, target)
        if target.hp <= 0:
            units.remove(target)
        return

    # 移动逼近
    for dx, dy in [(1,0),(-1,0),(0,1),(0,-1)]:
        nx, ny = u.x + dx, u.y + dy
        if 0 <= nx < COLS and 0 <= ny < ROWS:
            if TERRAIN[game_map[ny][nx]][1] != -1 and not unit_at(nx, ny):
                if abs(target.x - nx) + abs(target.y - ny) < dist:
                    u.x, u.y = nx, ny
                    return

# ---------- 初始单位 ----------
units = [
    Unit(1, 1, 0, "INF"),
    Unit(2, 1, 0, "SNIPER"),
    Unit(9, 7, 1, "MECH"),
    Unit(10, 6, 1, "INF"),
    Unit(8, 6, 1, "SNIPER"),
]

turn = 0
selected = None
mode = "move"

# ---------- 主循环 ----------
running = True
while running:
    screen.fill((10, 10, 20))

    # 地图
    for y in range(ROWS):
        for x in range(COLS):
            r = pygame.Rect(x * TILE, y * TILE, TILE, TILE)
            pygame.draw.rect(screen, TERRAIN[game_map[y][x]][2], r)
            pygame.draw.rect(screen, (60, 60, 80), r, 1)

    # 单位
    for u in units:
        u.draw()

    # 移动高亮
    if selected and mode == "move":
        for y in range(ROWS):
            for x in range(COLS):
                if can_move(selected, x, y):
                    r = pygame.Rect(x * TILE, y * TILE, TILE, TILE)
                    pygame.draw.rect(screen, (0, 255, 255), r, 2)

    # UI
    screen.blit(font.render(
        f"回合: {'玩家' if turn==0 else '敌方'} | 模式: {mode.upper()} | E=结束",
        True, (220, 220, 220)), (10, HEIGHT - 50))

    # 事件
    for e in pygame.event.get():
        if e.type == pygame.QUIT:
            running = False

        elif e.type == pygame.MOUSEBUTTONDOWN and turn == 0:
            mx, my = pygame.mouse.get_pos()
            gx, gy = mx // TILE, my // TILE

            if mode == "move":
                u = unit_at(gx, gy)
                if u and u.team == 0:
                    selected = u
                elif selected and can_move(selected, gx, gy):
                    selected.x, selected.y = gx, gy
                    selected = None

            elif mode == "attack":
                t = unit_at(gx, gy)
                if selected and t and t.team != selected.team:
                    combat(selected, t)
                    if t.hp <= 0:
                        units.remove(t)
                    selected = None

        elif e.type == pygame.KEYDOWN and turn == 0:
            if e.key == pygame.K_m:
                mode = "move"
            elif e.key == pygame.K_a:
                mode = "attack"
            elif e.key == pygame.K_e:
                turn = 1

    # AI
    if turn == 1:
        for u in [u for u in units if u.team == 1]:
            enemy_ai(u)
        turn = 0

    # 胜负
    if not any(u.team == 0 for u in units):
        print("💀 玩家战败")
        running = False
    if not any(u.team == 1 for u in units):
        print("🎉 玩家胜利")
        running = False

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

pygame.quit()