import pygame
import sys
from typing import List, Tuple, Optional

# ==================== 可配置参数区域 ====================
# 【核心参数】修改这个值可以整体缩放棋盘和棋子
# 原始值约为 60,缩小可尝试设置为 40 或 50
CELL_SIZE = 45 

# 派生参数(自动计算,无需手动修改)
BOARD_COLS = 9
BOARD_ROWS = 10
# 棋子半径为格子大小的 0.4 倍
PIECE_RADIUS = int(CELL_SIZE * 0.4) 
# 边距为格子大小的 0.8 倍
MARGIN = int(CELL_SIZE * 0.8) 
# 底部状态栏高度 
BOTTOM_BAR_HEIGHT = int(CELL_SIZE * 1.2) 

# 窗口大小计算
WINDOW_WIDTH = (BOARD_COLS - 1) * CELL_SIZE + 2 * MARGIN
WINDOW_HEIGHT = (BOARD_ROWS - 1) * CELL_SIZE + 2 * MARGIN + BOTTOM_BAR_HEIGHT
# =======================================================

# 颜色定义
BOARD_COLOR = (255, 206, 158)
LINE_COLOR = (139, 69, 19)
RED_COLOR = (200, 30, 30)
BLACK_COLOR = (30, 30, 30)
BG_COLOR = (245, 222, 179)
TEXT_COLOR = (80, 40, 0)

# 棋子类型
PIECE_TYPES = {
 'r_king': '帥', 'r_advisor': '仕', 'r_elephant': '相',
 'r_horse': '傌', 'r_chariot': '俥', 'r_cannon': '炮', 'r_soldier': '兵',
 'b_king': '將', 'b_advisor': '士', 'b_elephant': '象',
 'b_horse': '馬', 'b_chariot': '車', 'b_cannon': '砲', 'b_soldier': '卒'
}

# ==================== 棋子类 ====================
class Piece:
 def __init__(self, piece_type: str, is_red: bool, col: int, row: int):
 self.piece_type = piece_type
 self.is_red = is_red
 self.col = col
 self.row = row

 def get_char(self) -> str:
 return PIECE_TYPES[self.piece_type]

 def get_name(self) -> str:
 return self.piece_type.split('_')[1]

# ==================== 游戏逻辑类 ====================
class ChineseChess:
 def __init__(self):
 self.board: List[List[Optional[Piece]]] = [[None] * BOARD_COLS for _ in range(BOARD_ROWS)]
 self.pieces: List[Piece] = []
 self.is_red_turn = True
 self.selected_piece: Optional[Piece] = None
 self.valid_moves: List[Tuple[int, int]] = []
 self.game_over = False
 self.winner: Optional[str] = None
 self.message = "红方先行"
 self.init_board()

 def init_board(self):
 self.board = [[None] * BOARD_COLS for _ in range(BOARD_ROWS)]
 self.pieces = []
 initial_black = [
 ('chariot', 0, 0), ('horse', 1, 0), ('elephant', 2, 0), ('advisor', 3, 0), ('king', 4, 0),
 ('advisor', 5, 0), ('elephant', 6, 0), ('horse', 7, 0), ('chariot', 8, 0),
 ('cannon', 1, 2), ('cannon', 7, 2),
 ('soldier', 0, 3), ('soldier', 2, 3), ('soldier', 4, 3), ('soldier', 6, 3), ('soldier', 8, 3)
 ]
 initial_red = [
 ('chariot', 0, 9), ('horse', 1, 9), ('elephant', 2, 9), ('advisor', 3, 9), ('king', 4, 9),
 ('advisor', 5, 9), ('elephant', 6, 9), ('horse', 7, 9), ('chariot', 8, 9),
 ('cannon', 1, 7), ('cannon', 7, 7),
 ('soldier', 0, 6), ('soldier', 2, 6), ('soldier', 4, 6), ('soldier', 6, 6), ('soldier', 8, 6)
 ]
 for pt, c, r in initial_black:
 p = Piece('b_' + pt, False, c, r)
 self.board[r][c] = p
 self.pieces.append(p)
 for pt, c, r in initial_red:
 p = Piece('r_' + pt, True, c, r)
 self.board[r][c] = p
 self.pieces.append(p)

 def get_piece(self, col: int, row: int) -> Optional[Piece]:
 if 0 <= col < BOARD_COLS and 0 <= row < BOARD_ROWS:
 return self.board[row][col]
 return None

 def is_in_palace(self, col: int, row: int, is_red: bool) -> bool:
 if col < 3 or col > 5: return False
 return (7 <= row <= 9) if is_red else (0 <= row <= 2)

 def get_valid_moves(self, piece: Piece) -> List[Tuple[int, int]]:
 moves = []
 name = piece.get_name()
 if name == 'king': moves = self._get_king_moves(piece)
 elif name == 'advisor': moves = self._get_advisor_moves(piece)
 elif name == 'elephant': moves = self._get_elephant_moves(piece)
 elif name == 'horse': moves = self._get_horse_moves(piece)
 elif name == 'chariot': moves = self._get_chariot_moves(piece)
 elif name == 'cannon': moves = self._get_cannon_moves(piece)
 elif name == 'soldier': moves = self._get_soldier_moves(piece)
 
 # 过滤不安全的走法
 valid_moves = []
 for col, row in moves:
 if self._is_safe_move(piece, col, row):
 valid_moves.append((col, row))
 return valid_moves

 def _get_king_moves(self, piece: Piece):
 moves = []
 for dc, dr in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
 nc, nr = piece.col + dc, piece.row + dr
 if self.is_in_palace(nc, nr, piece.is_red):
 target = self.get_piece(nc, nr)
 if target is None or target.is_red != piece.is_red: moves.append((nc, nr))
 # 飞将
 ok_col, ok_row = self._find_king(not piece.is_red)
 if ok_col == piece.col:
 blocked = False
 min_r, max_r = min(piece.row, ok_row), max(piece.row, ok_row)
 for r in range(min_r + 1, max_r):
 if self.board[r][piece.col]: blocked = True
 if not blocked: moves.append((ok_col, ok_row))
 return moves

 def _get_advisor_moves(self, piece: Piece):
 moves = []
 for dc, dr in [(1, 1), (1, -1), (-1, 1), (-1, -1)]:
 nc, nr = piece.col + dc, piece.row + dr
 if self.is_in_palace(nc, nr, piece.is_red):
 target = self.get_piece(nc, nr)
 if target is None or target.is_red != piece.is_red: moves.append((nc, nr))
 return moves

 def _get_elephant_moves(self, piece: Piece):
 moves = []
 dirs = [(2, 2, 1, 1), (2, -2, 1, -1), (-2, 2, -1, 1), (-2, -2, -1, -1)]
 for dc, dr, bc, br in dirs:
 nc, nr = piece.col + dc, piece.row + dr
 # 不能过河
 if piece.is_red and nr < 5: continue
 if not piece.is_red and nr > 4: continue
 if 0 <= nc < 9 and 0 <= nr < 10:
 if self.get_piece(piece.col + bc, piece.row + br) is None:
 target = self.get_piece(nc, nr)
 if target is None or target.is_red != piece.is_red: moves.append((nc, nr))
 return moves

 def _get_horse_moves(self, piece: Piece):
 moves = []
 horse_legs = [(1,2,0,1),(1,-2,0,-1),(-1,2,0,1),(-1,-2,0,-1),
 (2,1,1,0),(2,-1,1,0),(-2,1,-1,0),(-2,-1,-1,0)]
 for dc, dr, bc, br in horse_legs:
 nc, nr = piece.col + dc, piece.row + dr
 if 0 <= nc < 9 and 0 <= nr < 10:
 if self.get_piece(piece.col + bc, piece.row + br) is None:
 target = self.get_piece(nc, nr)
 if target is None or target.is_red != piece.is_red: moves.append((nc, nr))
 return moves

 def _get_chariot_moves(self, piece: Piece):
 moves = []
 for dc, dr in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
 nc, nr = piece.col + dc, piece.row + dr
 while 0 <= nc < 9 and 0 <= nr < 10:
 target = self.get_piece(nc, nr)
 if target is None: moves.append((nc, nr))
 else:
 if target.is_red != piece.is_red: moves.append((nc, nr))
 break
 nc += dc; nr += dr
 return moves

 def _get_cannon_moves(self, piece: Piece):
 moves = []
 for dc, dr in [(0, 1), (0, -1), (1, 0), (-1, 0)]:
 nc, nr = piece.col + dc, piece.row + dr
 jumped = False
 while 0 <= nc < 9 and 0 <= nr < 10:
 target = self.get_piece(nc, nr)
 if not jumped:
 if target is None: moves.append((nc, nr))
 else: jumped = True
 else:
 if target:
 if target.is_red != piece.is_red: moves.append((nc, nr))
 break
 nc += dc; nr += dr
 return moves

 def _get_soldier_moves(self, piece: Piece):
 moves = []
 # 修正方向:红方向上(row减小), 黑方向下(row增大)
 forward = -1 if piece.is_red else 1
 crossed = (piece.is_red and piece.row <= 4) or (not piece.is_red and piece.row >= 5)
 
 nr = piece.row + forward
 if 0 <= nr < 10:
 target = self.get_piece(piece.col, nr)
 if target is None or target.is_red != piece.is_red: moves.append((piece.col, nr))
 
 if crossed:
 for dc in [-1, 1]:
 nc = piece.col + dc
 if 0 <= nc < 9:
 target = self.get_piece(nc, piece.row)
 if target is None or target.is_red != piece.is_red: moves.append((nc, piece.row))
 return moves

 def _find_king(self, is_red: bool):
 prefix = 'r_' if is_red else 'b_'
 for p in self.pieces:
 if p.piece_type == prefix + 'king': return (p.col, p.row)
 return (-1, -1)

 def _is_king_in_check(self, is_red: bool):
 k_col, k_row = self._find_king(is_red)
 for p in self.pieces:
 if p.is_red != is_red:
 # 简单检测:是否能走到将的位置(需排除飞将本身以免递归)
 raw_moves = []
 if p.get_name() == 'king': raw_moves = self._get_king_moves(p)
 elif p.get_name() == 'chariot': raw_moves = self._get_chariot_moves(p)
 elif p.get_name() == 'cannon': raw_moves = self._get_cannon_moves(p)
 elif p.get_name() == 'horse': raw_moves = self._get_horse_moves(p)
 elif p.get_name() == 'soldier': raw_moves = self._get_soldier_moves(p)
 elif p.get_name() == 'advisor': raw_moves = self._get_advisor_moves(p)
 elif p.get_name() == 'elephant': raw_moves = self._get_elephant_moves(p)
 
 if (k_col, k_row) in raw_moves: return True
 return False

 def _is_safe_move(self, piece, to_col, to_row):
 # 模拟走棋
 from_col, from_row = piece.col, piece.row
 captured = self.board[to_row][to_col]
 
 self.board[from_row][from_col] = None
 self.board[to_row][to_col] = piece
 piece.col, piece.row = to_col, to_row
 if captured: self.pieces.remove(captured)

 in_check = self._is_king_in_check(piece.is_red)

 # 恢复
 self.board[from_row][from_col] = piece
 self.board[to_row][to_col] = captured
 piece.col, piece.row = from_col, from_row
 if captured: self.pieces.append(captured)
 return not in_check

 def make_move(self, piece, to_col, to_row):
 captured = self.board[to_row][to_col]
 if captured: 
 self.pieces.remove(captured)
 if captured.get_name() == 'king':
 self.game_over = True
 self.winner = "红方" if piece.is_red else "黑方"
 
 self.board[piece.row][piece.col] = None
 self.board[to_row][to_col] = piece
 piece.col, piece.row = to_col, to_row
 self.is_red_turn = not self.is_red_turn
 
 if not self.game_over:
 self.message = ("红方" if self.is_red_turn else "黑方") + "走棋"
 if self._is_king_in_check(self.is_red_turn): self.message = "将军!"

 def select_piece(self, col, row):
 if self.game_over: return False
 piece = self.get_piece(col, row)
 if piece and piece.is_red == self.is_red_turn:
 self.selected_piece = piece
 self.valid_moves = self.get_valid_moves(piece)
 return True
 elif self.selected_piece and (col, row) in self.valid_moves:
 self.make_move(self.selected_piece, col, row)
 self.selected_piece = None
 self.valid_moves = []
 return True
 return False

# ==================== 图形界面类 ====================
class ChessGUI:
 def __init__(self):
 self.screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
 pygame.display.set_caption(f"中国象棋 (尺寸: {CELL_SIZE})")
 self.clock = pygame.time.Clock()
 self.game = ChineseChess()
 
 # 字体大小随棋盘缩放
 # 稍微减小字体比例,防止棋子太小时字体溢出
 font_size_large = max(16, int(CELL_SIZE * 0.6))
 font_size_medium = max(14, int(CELL_SIZE * 0.5))
 font_size_small = max(12, int(CELL_SIZE * 0.4))
 
 self.font_large = pygame.font.SysFont('SimHei', font_size_large)
 self.font_medium = pygame.font.SysFont('SimHei', font_size_medium)
 self.font_small = pygame.font.SysFont('SimHei', font_size_small)
 
 self.restart_btn = pygame.Rect(WINDOW_WIDTH//2 - 40, WINDOW_HEIGHT - BOTTOM_BAR_HEIGHT + 10, 80, 30)

 def board_to_screen(self, col, row):
 return (MARGIN + col * CELL_SIZE, MARGIN + row * CELL_SIZE)

 def screen_to_board(self, x, y):
 return (round((x - MARGIN) / CELL_SIZE), round((y - MARGIN) / CELL_SIZE))

 def draw_board(self):
 self.screen.fill(BG_COLOR)
 # 棋盘矩形背景
 rect = pygame.Rect(MARGIN - CELL_SIZE//2, MARGIN - CELL_SIZE//2, 
 (BOARD_COLS-1)*CELL_SIZE + CELL_SIZE, (BOARD_ROWS-1)*CELL_SIZE + CELL_SIZE)
 pygame.draw.rect(self.screen, BOARD_COLOR, rect)
 
 # 画线
 for i in range(BOARD_ROWS):
 y = MARGIN + i * CELL_SIZE
 # 楚河汉界断开
 if i == 4 or i == 5:
 pygame.draw.line(self.screen, LINE_COLOR, (MARGIN, y), (MARGIN + 3*CELL_SIZE, y), 2)
 pygame.draw.line(self.screen, LINE_COLOR, (MARGIN + 5*CELL_SIZE, y), (MARGIN + 8*CELL_SIZE, y), 2)
 else:
 pygame.draw.line(self.screen, LINE_COLOR, (MARGIN, y), (MARGIN + 8*CELL_SIZE, y), 2)
 
 for i in range(BOARD_COLS):
 x = MARGIN + i * CELL_SIZE
 pygame.draw.line(self.screen, LINE_COLOR, (x, MARGIN), (x, MARGIN + 4*CELL_SIZE), 2)
 pygame.draw.line(self.screen, LINE_COLOR, (x, MARGIN + 5*CELL_SIZE), (x, MARGIN + 9*CELL_SIZE), 2)
 
 # 九宫斜线
 for x1,y1,x2,y2 in [(3,0,5,2), (3,7,5,9)]:
 s1 = self.board_to_screen(x1,y1); s2 = self.board_to_screen(x2,y2)
 pygame.draw.line(self.screen, LINE_COLOR, s1, s2, 2)
 pygame.draw.line(self.screen, LINE_COLOR, (s2[0], s1[1]), (s1[0], s2[1]), 2)
 
 # 楚河汉界文字
 if CELL_SIZE > 30: # 太小就不显示文字了
 txt = self.font_large.render("楚河 漢界", True, LINE_COLOR)
 self.screen.blit(txt, (MARGIN, MARGIN + 4*CELL_SIZE + (CELL_SIZE - txt.get_height())//2))

 def draw_pieces(self):
 for p in self.game.pieces:
 x, y = self.board_to_screen(p.col, p.row)
 color = (255, 230, 200) if p.is_red else (255, 250, 240)
 border = RED_COLOR if p.is_red else BLACK_COLOR
 
 pygame.draw.circle(self.screen, color, (x,y), PIECE_RADIUS)
 pygame.draw.circle(self.screen, border, (x,y), PIECE_RADIUS, 2)
 
 txt_color = RED_COLOR if p.is_red else BLACK_COLOR
 # 只有棋子足够大时才显示文字
 if PIECE_RADIUS > 10:
 txt = self.font_medium.render(p.get_char(), True, txt_color)
 self.screen.blit(txt, (x - txt.get_width()//2, y - txt.get_height()//2))
 
 # 选中状态
 if self.game.selected_piece:
 x,y = self.board_to_screen(self.game.selected_piece.col, self.game.selected_piece.row)
 pygame.draw.circle(self.screen, (0,255,0), (x,y), PIECE_RADIUS + 3, 3)
 for c,r in self.game.valid_moves:
 mx,my = self.board_to_screen(c,r)
 pygame.draw.circle(self.screen, (0,200,0), (mx,my), 5)

 def draw_ui(self):
 # 底部栏
 pygame.draw.rect(self.screen, (200, 180, 150), (0, WINDOW_HEIGHT - BOTTOM_BAR_HEIGHT, WINDOW_WIDTH, BOTTOM_BAR_HEIGHT))
 msg = f"{self.game.winner}获胜!" if self.game.game_over else self.game.message
 clr = RED_COLOR if self.game.is_red_turn else BLACK_COLOR
 text = self.font_small.render(msg, True, clr)
 self.screen.blit(text, (20, WINDOW_HEIGHT - BOTTOM_BAR_HEIGHT + 15))
 
 # 重新开始按钮
 pygame.draw.rect(self.screen, (180, 140, 100), self.restart_btn, border_radius=3)
 btn_txt = self.font_small.render("重开", True, TEXT_COLOR)
 self.screen.blit(btn_txt, (self.restart_btn.centerx - btn_txt.get_width()//2, 
 self.restart_btn.centery - btn_txt.get_height()//2))

 def run(self):
 while True:
 for e in pygame.event.get():
 if e.type == pygame.QUIT: pygame.quit(); sys.exit()
 if e.type == pygame.MOUSEBUTTONDOWN:
 if self.restart_btn.collidepoint(e.pos): self.game = ChineseChess()
 else:
 c, r = self.screen_to_board(*e.pos)
 self.game.select_piece(c, r)
 
 self.draw_board()
 self.draw_pieces()
 self.draw_ui()
 pygame.display.flip()
 self.clock.tick(60)

if __name__ == '__main__':
 gui = ChessGUI()
 gui.run()