from flask import Flask, render_template, request
from flask_socketio import SocketIO, emit, join_room, disconnect
import uuid
import copy

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your_secret_key_here'
socketio = SocketIO(app, cors_allowed_origins="*")

# 内存存储
rooms = {}
users = {}
username_to_sid = {}

# ---------- 棋盘初始化 ----------
def init_board():
    return [
        ['r', 'h', 'e', 'a', 'k', 'a', 'e', 'h', 'r'],
        ['.', '.', '.', '.', '.', '.', '.', '.', '.'],
        ['.', 'c', '.', '.', '.', '.', '.', 'c', '.'],
        ['p', '.', 'p', '.', 'p', '.', 'p', '.', 'p'],
        ['.', '.', '.', '.', '.', '.', '.', '.', '.'],
        ['.', '.', '.', '.', '.', '.', '.', '.', '.'],
        ['P', '.', 'P', '.', 'P', '.', 'P', '.', 'P'],
        ['.', 'C', '.', '.', '.', '.', '.', 'C', '.'],
        ['.', '.', '.', '.', '.', '.', '.', '.', '.'],
        ['R', 'H', 'E', 'A', 'K', 'A', 'E', 'H', 'R']
    ]

def board_to_string(board):
    return '\n'.join(''.join(row) for row in board)

# ---------- 规则验证 ----------
def is_valid_move(board, fr, fc, tr, tc, turn):
    """检查移动是否合法"""
    # 边界检查
    if not (0 <= fr < 10 and 0 <= fc < 9 and 0 <= tr < 10 and 0 <= tc < 9):
        return False

    piece = board[fr][fc]
    target = board[tr][tc]

    # 起点不能为空
    if piece == '.':
        return False

    # 检查是否属于当前方
    if turn == 'red' and not piece.isupper():
        return False
    if turn == 'black' and not piece.islower():
        return False

    # 目标格不能有己方棋子
    if target != '.':
        if (turn == 'red' and target.isupper()) or (turn == 'black' and target.islower()):
            return False

    p = piece.lower()
    dr = tr - fr
    dc = tc - fc

    # ---------- 车 ----------
    if p == 'r':
        if fr != tr and fc != tc:
            return False
        step_r = 1 if dr > 0 else -1 if dr < 0 else 0
        step_c = 1 if dc > 0 else -1 if dc < 0 else 0
        r, c = fr + step_r, fc + step_c
        while (r, c) != (tr, tc):
            if board[r][c] != '.':
                return False
            r += step_r
            c += step_c
        return True

    # ---------- 马 ----------
    elif p == 'h':
        if (abs(dr), abs(dc)) not in [(2, 1), (1, 2)]:
            return False
        # 绊马腿
        if abs(dr) == 2:
            block_r, block_c = fr + (1 if dr > 0 else -1), fc
        else:
            block_r, block_c = fr, fc + (1 if dc > 0 else -1)
        if board[block_r][block_c] != '.':
            return False
        return True

    # ---------- 相/象 ----------
    elif p == 'e':
        if abs(dr) != 2 or abs(dc) != 2:
            return False
        # 不能过河
        if turn == 'red' and tr < 5:    # 红相不能过河到对方半场
            return False
        if turn == 'black' and tr > 4:  # 黑象不能过河
            return False
        # 象眼
        block_r, block_c = fr + (1 if dr > 0 else -1), fc + (1 if dc > 0 else -1)
        if board[block_r][block_c] != '.':
            return False
        return True

    # ---------- 士 ----------
    elif p == 'a':
        if abs(dr) != 1 or abs(dc) != 1:
            return False
        # 九宫
        if turn == 'red':
            if tr < 7 or tc < 3 or tc > 5:
                return False
        else:
            if tr > 2 or tc < 3 or tc > 5:
                return False
        return True

    # ---------- 将/帅 ----------
    elif p == 'k':
        if abs(dr) + abs(dc) != 1:
            return False
        # 九宫
        if turn == 'red':
            if tr < 7 or tc < 3 or tc > 5:
                return False
        else:
            if tr > 2 or tc < 3 or tc > 5:
                return False
        return True

    # ---------- 炮 ----------
    elif p == 'c':
        if fr != tr and fc != tc:
            return False
        step_r = 1 if dr > 0 else -1 if dr < 0 else 0
        step_c = 1 if dc > 0 else -1 if dc < 0 else 0
        r, c = fr + step_r, fc + step_c
        cnt = 0
        while (r, c) != (tr, tc):
            if board[r][c] != '.':
                cnt += 1
            r += step_r
            c += step_c
        # 吃子：需要恰好一个炮架
        if target != '.':
            return cnt == 1
        # 不吃子：路径必须为空
        else:
            return cnt == 0

    # ---------- 兵/卒 ----------
    elif p == 'p':
        if turn == 'red':
            if dr > 0:  # 不能后退
                return False
            # 未过河（行 > 4）只能前进
            if fr > 4:
                return dc == 0 and dr == -1
            else:  # 过河
                return (dr == -1 and dc == 0) or (dr == 0 and abs(dc) == 1)
        else:  # 黑方
            if dr < 0:
                return False
            if fr < 5:  # 未过河
                return dc == 0 and dr == 1
            else:       # 过河
                return (dr == 1 and dc == 0) or (dr == 0 and abs(dc) == 1)

    return False

def is_check(board, turn):
    """检查 turn 方是否被将"""
    # 找出将帅位置
    red_king = black_king = None
    for i in range(10):
        for j in range(9):
            if board[i][j] == 'K':
                red_king = (i, j)
            elif board[i][j] == 'k':
                black_king = (i, j)

    if turn == 'red':
        if not red_king:
            return True  # 红方无帅，输
        # 检查黑方所有棋子能否吃到红帅
        for i in range(10):
            for j in range(9):
                piece = board[i][j]
                if piece.islower() and piece != '.':
                    if is_valid_move(board, i, j, red_king[0], red_king[1], 'black'):
                        return True
        # 将帅对面
        if black_king and red_king[1] == black_king[1]:
            c = red_king[1]
            r1, r2 = sorted([red_king[0], black_king[0]])
            for r in range(r1 + 1, r2):
                if board[r][c] != '.':
                    break
            else:
                return True  # 对面
        return False
    else:
        if not black_king:
            return True
        for i in range(10):
            for j in range(9):
                piece = board[i][j]
                if piece.isupper() and piece != '.':
                    if is_valid_move(board, i, j, black_king[0], black_king[1], 'red'):
                        return True
        if red_king and red_king[1] == black_king[1]:
            c = red_king[1]
            r1, r2 = sorted([red_king[0], black_king[0]])
            for r in range(r1 + 1, r2):
                if board[r][c] != '.':
                    break
            else:
                return True
        return False

def make_move(board, fr, fc, tr, tc, turn):
    """执行移动，如果合法返回新棋盘，否则返回 None"""
    if not is_valid_move(board, fr, fc, tr, tc, turn):
        return None
    new_board = copy.deepcopy(board)
    piece = new_board[fr][fc]
    new_board[fr][fc] = '.'
    new_board[tr][tc] = piece
    # 检查移动后是否导致己方被将
    if is_check(new_board, turn):
        return None
    return new_board

# ---------- 路由 ----------
@app.route('/')
def lobby():
    return render_template('lobby.html')

@app.route('/room/<room_id>')
def room_page(room_id):
    if room_id not in rooms:
        return "房间不存在", 404
    return render_template('room.html', room_id=room_id)

# ---------- SocketIO 事件 ----------
@socketio.on('connect')
def handle_connect():
    print('Client connected', request.sid)

@socketio.on('disconnect')
def handle_disconnect():
    sid = request.sid
    username = users.pop(sid, None)
    if username:
        if username_to_sid.get(username) == sid:
            username_to_sid.pop(username, None)
        for rid, room in rooms.items():
            if room['players']['red'] == username:
                room['players']['red'] = None
                room['players_sid']['red'] = None
                room['user_color'].pop(username, None)
                emit('player_left', {'color': 'red'}, room=rid)
            if room['players']['black'] == username:
                room['players']['black'] = None
                room['players_sid']['black'] = None
                room['user_color'].pop(username, None)
                emit('player_left', {'color': 'black'}, room=rid)
            if sid in room['watchers']:
                del room['watchers'][sid]
        broadcast_rooms()

@socketio.on('login')
def handle_login(data):
    username = data['username']
    sid = request.sid
    if username in username_to_sid:
        old_sid = username_to_sid[username]
        emit('force_disconnect', {'msg': '您的账号在另一处登录'}, room=old_sid)
        disconnect(old_sid)
    users[sid] = username
    username_to_sid[username] = sid
    emit('login_success', {'username': username})

@socketio.on('get_rooms')
def handle_get_rooms():
    room_list = []
    for rid, room in rooms.items():
        room_list.append({
            'id': rid,
            'name': room['name'],
            'public': room['public'],
            'has_password': bool(room['password']),
            'players': sum(1 for p in room['players'].values() if p is not None),
            'watchers': len(room['watchers'])
        })
    emit('rooms_list', room_list)

def broadcast_rooms():
    room_list = []
    for rid, room in rooms.items():
        room_list.append({
            'id': rid,
            'name': room['name'],
            'public': room['public'],
            'has_password': bool(room['password']),
            'players': sum(1 for p in room['players'].values() if p is not None),
            'watchers': len(room['watchers'])
        })
    emit('rooms_updated', room_list, broadcast=True)

@socketio.on('create_room')
def handle_create_room(data):
    username = users.get(request.sid)
    if not username:
        emit('error', {'msg': '未登录'})
        return
    room_name = data['name']
    public = data['public']
    password = data.get('password', '') if not public else ''
    room_id = str(uuid.uuid4())[:8]

    rooms[room_id] = {
        'name': room_name,
        'password': password,
        'public': public,
        'board': init_board(),
        'turn': 'red',
        'players': {'red': None, 'black': None},
        'players_sid': {'red': None, 'black': None},
        'watchers': {},
        'user_color': {}
    }

    join_room(room_id)
    rooms[room_id]['players']['red'] = username
    rooms[room_id]['players_sid']['red'] = request.sid
    rooms[room_id]['user_color'][username] = 'red'

    emit('room_created', {'room_id': room_id}, room=request.sid)
    broadcast_rooms()

@socketio.on('join_room')
def handle_join_room(data):
    username = users.get(request.sid)
    if not username:
        emit('error', {'msg': '未登录'})
        return

    room_id = data['room_id']
    password = data.get('password', '')

    if room_id not in rooms:
        emit('error', {'msg': '房间不存在'})
        return

    room = rooms[room_id]

    if room['password'] and room['password'] != password:
        emit('error', {'msg': '密码错误'})
        return

    # 重连处理
    if username in room['user_color']:
        color = room['user_color'][username]
        old_sid = room['players_sid'][color]
        if old_sid and old_sid != request.sid:
            emit('force_disconnect', {'msg': '您在另一个窗口加入了房间'}, room=old_sid)
            disconnect(old_sid)

        room['players_sid'][color] = request.sid
        join_room(room_id)
        emit('room_joined', {
            'board': board_to_string(room['board']),
            'turn': room['turn'],
            'my_color': color,
            'players': room['players'],
            'watchers': len(room['watchers'])
        }, room=request.sid)
        return

    # 分配颜色
    if room['players']['red'] is None:
        color = 'red'
    elif room['players']['black'] is None:
        color = 'black'
    else:
        color = None

    if color:
        room['players'][color] = username
        room['players_sid'][color] = request.sid
        room['user_color'][username] = color
        join_room(room_id)
        emit('room_joined', {
            'board': board_to_string(room['board']),
            'turn': room['turn'],
            'my_color': color,
            'players': room['players'],
            'watchers': len(room['watchers'])
        }, room=request.sid)
        emit('player_joined', {'color': color, 'username': username}, room=room_id, include_self=False)
        print(f"玩家 {username} 加入房间 {room_id} 作为 {color}")
    else:
        room['watchers'][request.sid] = username
        join_room(room_id)
        emit('room_joined', {
            'board': board_to_string(room['board']),
            'turn': room['turn'],
            'my_color': None,
            'players': room['players'],
            'watchers': len(room['watchers'])
        }, room=request.sid)
        emit('watcher_joined', {'username': username}, room=room_id, include_self=False)
        print(f"玩家 {username} 加入房间 {room_id} 作为旁观者")

    broadcast_rooms()

@socketio.on('move')
def handle_move(data):
    sid = request.sid
    username = users.get(sid)
    if not username:
        emit('error', {'msg': '未登录'})
        return

    room_id = data['room_id']
    if room_id not in rooms:
        emit('error', {'msg': '房间不存在'})
        return

    room = rooms[room_id]

    # 回合检查
    if room['turn'] == 'red' and room['players_sid']['red'] != sid:
        emit('move_result', {'success': False, 'message': '不是你的回合'})
        return
    if room['turn'] == 'black' and room['players_sid']['black'] != sid:
        emit('move_result', {'success': False, 'message': '不是你的回合'})
        return

    fr, fc = data['from']['row'], data['from']['col']
    tr, tc = data['to']['row'], data['to']['col']

    print(f"玩家 {username} 尝试移动: ({fr},{fc}) -> ({tr},{tc})")
    print(f"当前棋盘:")
    for row in room['board']:
        print(' '.join(row))

    new_board = make_move(room['board'], fr, fc, tr, tc, room['turn'])
    if new_board is None:
        emit('move_result', {'success': False, 'message': '非法移动'})
        return

    room['board'] = new_board
    room['turn'] = 'black' if room['turn'] == 'red' else 'red'

    print(f"移动成功，新棋盘:")
    for row in room['board']:
        print(' '.join(row))
    print(f"下一回合: {room['turn']}")

    emit('board_updated', {
        'board': board_to_string(room['board']),
        'turn': room['turn']
    }, room=room_id)

    emit('move_result', {'success': True}, room=sid)

if __name__ == '__main__':
    socketio.run(app, debug=True,allow_unsafe_werkzeug=True, host='0.0.0.0', port=5000)