import pygame
import sys

# ---------------- 基础设置 ----------------
WIDTH, HEIGHT = 640, 720
FPS = 60

BOARD_X = 40
BOARD_Y = 40
CELL = 70
COLS = 9
ROWS = 10

pygame.init()

screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("中国象棋 Chinese Chess")

clock = pygame.time.Clock()
font = pygame.font.SysFont("simhei", 36)
small_font = pygame.font.SysFont("simhei", 24)

WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BOARD_COLOR = (245, 222, 179)
LINE_COLOR = (80, 50, 20)
RED = (200, 40, 40)
BLACK_PIECE = (30, 30, 30)
SELECT_COLOR = (60, 180, 60)
MOVE_COLOR = (60, 120, 200)


# ---------------- 棋盘绘制 ----------------
def draw_board():
    screen.fill(BOARD_COLOR)

    # 画横线
    for r in range(ROWS):
        y = BOARD_Y + r * CELL
        pygame.draw.line(screen, LINE_COLOR, (BOARD_X, y), (BOARD_X + (COLS - 1) * CELL, y), 2)

    # 画竖线（注意楚河汉界中间断开）
    for c in range(COLS):
        if c == 0 or c == COLS - 1:
            pygame.draw.line(screen, LINE_COLOR, (BOARD_X + c * CELL, BOARD_Y), (BOARD_X + c * CELL, BOARD_Y + (ROWS - 1) * CELL), 2)
        else:
            pygame.draw.line(screen, LINE_COLOR, (BOARD_X + c * CELL, BOARD_Y), (BOARD_X + c * CELL, BOARD_Y + 4 * CELL), 2)
            pygame.draw.line(screen, LINE_COLOR, (BOARD_X + c * CELL, BOARD_Y + 5 * CELL), (BOARD_X + c * CELL, BOARD_Y + (ROWS - 1) * CELL), 2)

    # 九宫格斜线
    # 上方九宫
    pygame.draw.line(screen, LINE_COLOR, (BOARD_X + 3 * CELL, BOARD_Y), (BOARD_X + 5 * CELL, BOARD_Y + 2 * CELL), 2)
    pygame.draw.line(screen, LINE_COLOR, (BOARD_X + 5 * CELL, BOARD_Y), (BOARD_X + 3 * CELL, BOARD_Y + 2 * CELL), 2)
    # 下方九宫
    pygame.draw.line(screen, LINE_COLOR, (BOARD_X + 3 * CELL, BOARD_Y + 7 * CELL), (BOARD_X + 5 * CELL, BOARD_Y + 9 * CELL), 2)
    pygame.draw.line(screen, LINE_COLOR, (BOARD_X + 5 * CELL, BOARD_Y + 7 * CELL), (BOARD_X + 3 * CELL, BOARD_Y + 9 * CELL), 2)

    # 楚河汉界
    text = font.render("楚 河          汉 界", True, LINE_COLOR)
    screen.blit(text, (BOARD_X + 40, BOARD_Y + 4 * CELL + 15))


def board_to_screen(col, row):
    return BOARD_X + col * CELL, BOARD_Y + row * CELL


def screen_to_board(x, y):
    col = round((x - BOARD_X) / CELL)
    row = round((y - BOARD_Y) / CELL)
    if 0 <= col < COLS and 0 <= row < ROWS:
        return col, row
    return None, None


# ---------------- 棋子 ----------------
class Piece:
    def __init__(self, name, color, col, row):
        self.name = name      # "车" "马" "炮" "兵" "仕" "相" "帅" 等
        self.color = color    # "red" or "black"
        self.col = col
        self.row = row

    def draw(self, selected=False):
        x, y = board_to_screen(self.col, self.row)

        # 选中高亮
        if selected:
            pygame.draw.circle(screen, SELECT_COLOR, (x, y), 34, 3)

        # 棋子底色
        pygame.draw.circle(screen, (255, 245, 220), (x, y), 30)
        pygame.draw.circle(screen, LINE_COLOR, (x, y), 30, 2)

        # 棋子文字
        color = RED if self.color == "red" else BLACK_PIECE
        text = font.render(self.name, True, color)
        screen.blit(text, (x - text.get_width() // 2, y - text.get_height() // 2))


# ---------------- 初始化棋子 ----------------
def init_pieces():
    pieces = []

    # 黑方（上方，row 0~4）
    back_row = ["车", "马", "相", "仕", "将", "仕", "相", "马", "车"]
    for c, name in enumerate(back_row):
        pieces.append(Piece(name, "black", c, 0))
    pieces.append(Piece("炮", "black", 1, 2))
    pieces.append(Piece("炮", "black", 7, 2))
    for c in [0, 2, 4, 6, 8]:
        pieces.append(Piece("卒", "black", c, 3))

    # 红方（下方，row 5~9）
    back_row = ["车", "马", "相", "仕", "帅", "仕", "相", "马", "车"]
    for c, name in enumerate(back_row):
        pieces.append(Piece(name, "red", c, 9))
    pieces.append(Piece("炮", "red", 1, 7))
    pieces.append(Piece("炮", "red", 7, 7))
    for c in [0, 2, 4, 6, 8]:
        pieces.append(Piece("兵", "red", c, 6))

    return pieces


def get_piece_at(pieces, col, row):
    for p in pieces:
        if p.col == col and p.row == row:
            return p
    return None


# ---------------- 走法生成 ----------------
def in_board(col, row):
    return 0 <= col < COLS and 0 <= row < ROWS


def in_palace(col, row, color):
    if not (3 <= col <= 5):
        return False
    if color == "red":
        return 7 <= row <= 9
    else:
        return 0 <= row <= 2


def get_legal_moves(piece, pieces):
    moves = []
    name = piece.name
    color = piece.color
    col, row = piece.col, piece.row

    if name == "车":
        for dc, dr in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
            nc, nr = col + dc, row + dr
            while in_board(nc, nr):
                target = get_piece_at(pieces, nc, nr)
                if target is None:
                    moves.append((nc, nr))
                else:
                    if target.color != color:
                        moves.append((nc, nr))
                    break
                nc += dc
                nr += dr

    elif name == "马":
        for dc, dr, bc, br in [
            (1, 2, 1, 0), (-1, 2, -1, 0),
            (1, -2, 1, 0), (-1, -2, -1, 0),
            (2, 1, 0, 1), (2, -1, 0, -1),
            (-2, 1, 0, 1), (-2, -1, 0, -1),
        ]:
            nc, nr = col + dc, row + dr
            block_c, block_r = col + bc, row + br
            if in_board(nc, nr) and get_piece_at(pieces, block_c, block_r) is None:
                target = get_piece_at(pieces, nc, nr)
                if target is None or target.color != color:
                    moves.append((nc, nr))

    elif name == "相" or name == "象":
        for dc, dr in [(2, 2), (2, -2), (-2, 2), (-2, -2)]:
            nc, nr = col + dc, row + dr
            ec, er = col + dc // 2, row + dr // 2
            if not in_board(nc, nr):
                continue
            # 不能过河
            if color == "red" and nr < 5:
                continue
            if color == "black" and nr > 4:
                continue
            # 塞象眼
            if get_piece_at(pieces, ec, er) is not None:
                continue
            target = get_piece_at(pieces, nc, nr)
            if target is None or target.color != color:
                moves.append((nc, nr))

    elif name == "仕" or name == "士":
        for dc, dr in [(1, 1), (1, -1), (-1, 1), (-1, -1)]:
            nc, nr = col + dc, row + dr
            if in_palace(nc, nr, color):
                target = get_piece_at(pieces, nc, nr)
                if target is None or target.color != color:
                    moves.append((nc, nr))

    elif name == "帅" or name == "将":
        for dc, dr in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
            nc, nr = col + dc, row + dr
            if in_palace(nc, nr, color):
                target = get_piece_at(pieces, nc, nr)
                if target is None or target.color != color:
                    moves.append((nc, nr))

    elif name == "兵":
        # 红兵向上
        nc, nr = col, row - 1
        if in_board(nc, nr):
            target = get_piece_at(pieces, nc, nr)
            if target is None or target.color != color:
                moves.append((nc, nr))
        # 过河后可左右
        if row <= 4:
            for dc in [-1, 1]:
                nc, nr = col + dc, row
                if in_board(nc, nr):
                    target = get_piece_at(pieces, nc, nr)
                    if target is None or target.color != color:
                        moves.append((nc, nr))

    elif name == "卒":
        # 黑卒向下
        nc, nr = col, row + 1
        if in_board(nc, nr):
            target = get_piece_at(pieces, nc, nr)
            if target is None or target.color != color:
                moves.append((nc, nr))
        # 过河后可左右
        if row >= 5:
            for dc in [-1, 1]:
                nc, nr = col + dc, row
                if in_board(nc, nr):
                    target = get_piece_at(pieces, nc, nr)
                    if target is None or target.color != color:
                        moves.append((nc, nr))

    elif name == "炮":
        for dc, dr in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
            nc, nr = col + dc, row + dr
            jumped = False
            while in_board(nc, nr):
                target = get_piece_at(pieces, nc, nr)
                if not jumped:
                    if target is None:
                        moves.append((nc, nr))
                    else:
                        jumped = True
                else:
                    if target is not None:
                        if target.color != color:
                            moves.append((nc, nr))
                        break
                nc += dc
                nr += dr

    return moves


def kings_facing(pieces):
    """检查将帅是否照面"""
    red_king = black_king = None
    for p in pieces:
        if p.name == "帅":
            red_king = p
        elif p.name == "将":
            black_king = p

    if red_king is None or black_king is None:
        return False

    if red_king.col != black_king.col:
        return False

    for p in pieces:
        if p.col == red_king.col:
            if (p.row > red_king.row and p.row < black_king.row) or \
               (p.row > black_king.row and p.row < red_king.row):
                return False
    return True


def is_in_check(color, pieces):
    """检查某方是否被将军"""
    king_name = "帅" if color == "red" else "将"
    king = None
    for p in pieces:
        if p.name == king_name:
            king = p
            break

    if king is None:
        return True

    enemy_color = "black" if color == "red" else "red"
    for p in pieces:
        if p.color == enemy_color:
            moves = get_legal_moves(p, pieces)
            if (king.col, king.row) in moves:
                return True
    return False


def make_move(piece, to_col, to_row, pieces):
    """执行走子，返回被吃的棋子（如果有）"""
    captured = get_piece_at(pieces, to_col, to_row)
    if captured:
        pieces.remove(captured)
    piece.col = to_col
    piece.row = to_row
    return captured


def undo_move(piece, from_col, from_row, captured, pieces):
    """撤销走子"""
    piece.col = from_col
    piece.row = from_row
    if captured:
        pieces.append(captured)


def get_legal_moves_filtered(piece, pieces):
    """过滤掉会导致自己被将军的走法"""
    raw_moves = get_legal_moves(piece, pieces)
    legal = []
    for nc, nr in raw_moves:
        from_col, from_row = piece.col, piece.row
        captured = make_move(piece, nc, nr, pieces)
        if not is_in_check(piece.color, pieces):
            legal.append((nc, nr))
        undo_move(piece, from_col, from_row, captured, pieces)
    return legal


# ---------------- 游戏主循环 ----------------
def main():
    pieces = init_pieces()
    current_turn = "red"
    selected_piece = None
    legal_moves = []
    game_over = False
    winner = None

    running = True
    while running:
        clock.tick(FPS)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r:
                    pieces = init_pieces()
                    current_turn = "red"
                    selected_piece = None
                    legal_moves = []
                    game_over = False
                    winner = None

            if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1 and not game_over:
                mx, my = event.pos
                col, row = screen_to_board(mx, my)
                if col is None:
                    continue

                clicked_piece = get_piece_at(pieces, col, row)

                if selected_piece is None:
                    # 选择己方棋子
                    if clicked_piece and clicked_piece.color == current_turn:
                        selected_piece = clicked_piece
                        legal_moves = get_legal_moves_filtered(selected_piece, pieces)
                else:
                    # 尝试走子
                    if (col, row) in legal_moves:
                        make_move(selected_piece, col, row, pieces)

                        # 检查是否吃掉对方将/帅
                        enemy_king = "将" if current_turn == "red" else "帅"
                        if clicked_piece and clicked_piece.name == enemy_king:
                            game_over = True
                            winner = current_turn

                        # 检查是否照面
                        if kings_facing(pieces):
                            # 照面时走子无效，撤销
                            undo_move(selected_piece, selected_piece.col, selected_piece.row, clicked_piece, pieces)
                        else:
                            # 切换回合
                            current_turn = "black" if current_turn == "red" else "red"

                            # 检查对方是否被将死
                            if is_in_check(current_turn, pieces):
                                has_legal = False
                                for p in pieces:
                                    if p.color == current_turn:
                                        if get_legal_moves_filtered(p, pieces):
                                            has_legal = True
                                            break
                                if not has_legal:
                                    game_over = True
                                    winner = "black" if current_turn == "red" else "red"

                        selected_piece = None
                        legal_moves = []
                    else:
                        # 点击了其他己方棋子，切换选择
                        if clicked_piece and clicked_piece.color == current_turn:
                            selected_piece = clicked_piece
                            legal_moves = get_legal_moves_filtered(selected_piece, pieces)
                        else:
                            selected_piece = None
                            legal_moves = []

        # ---------------- 绘制 ----------------
        draw_board()

        # 绘制合法走法提示
        for mc, mr in legal_moves:
            x, y = board_to_screen(mc, mr)
            target = get_piece_at(pieces, mc, mr)
            if target:
                pygame.draw.circle(screen, MOVE_COLOR, (x, y), 32, 3)
            else:
                pygame.draw.circle(screen, MOVE_COLOR, (x, y), 8)

        # 绘制棋子
        for p in pieces:
            p.draw(selected=(p is selected_piece))

        # 状态栏
        turn_text = "红方走棋" if current_turn == "red" else "黑方走棋"
        turn_color = RED if current_turn == "red" else BLACK_PIECE
        text = small_font.render(turn_text, True, turn_color)
        screen.blit(text, (BOARD_X, HEIGHT - 40))

        hint = small_font.render("按 R 重新开始", True, LINE_COLOR)
        screen.blit(hint, (WIDTH - 180, HEIGHT - 40))

        if game_over:
            overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
            overlay.fill((0, 0, 0, 140))
            screen.blit(overlay, (0, 0))

            win_text = "红方胜利！" if winner == "red" else "黑方胜利！"
            win_color = RED if winner == "red" else WHITE
            text = font.render(win_text, True, win_color)
            screen.blit(text, (WIDTH // 2 - text.get_width() // 2, HEIGHT // 2 - 30))

            restart_text = small_font.render("按 R 重新开始", True, WHITE)
            screen.blit(restart_text, (WIDTH // 2 - restart_text.get_width() // 2, HEIGHT // 2 + 20))

        pygame.display.flip()

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()