#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
四方洲行动 · 联机中继服务器（Python 零依赖版）
用法： python server.py [端口]      例如： python server.py 8787

功能与 Node 版一致：房间制（路径 /:roomCode）、邀请码、状态/聊天转发、局域网 UDP 发现。
仅用 Python 标准库（Python 3.7+）。
"""
import sys, os, json, random, string, struct, base64, hashlib, socket, threading, uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from socketserver import ThreadingMixIn

PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 3210
GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

rooms = {}  # code -> set of conn
rooms_lock = threading.Lock()

def gen_code(n=6):
    chars = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
    return "".join(random.choice(chars) for _ in range(n))

def lan_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80)); return s.getsockname()[0]
    except Exception:
        return "127.0.0.1"
    finally:
        s.close()

def get_room(code):
    with rooms_lock:
        if code not in rooms:
            rooms[code] = set()
        return rooms[code]

# ---------- WebSocket 帧 ----------
def ws_accept(key):
    return base64.b64encode(hashlib.sha1((key + GUID).encode()).digest()).decode()

def send_frame(sock, data: bytes, opcode=0x1):
    length = len(data)
    if length < 126:
        header = bytes([0x80 | opcode, length])
    elif length < 65536:
        header = bytes([0x80 | opcode, 126]) + struct.pack(">H", length)
    else:
        header = bytes([0x80 | opcode, 127]) + struct.pack(">Q", length)
    try:
        sock.sendall(header + data)
    except Exception:
        pass

def recv_frame(sock):
    """读取一个完整帧，返回 (opcode, payload bytes) 或 None（断开）"""
    try:
        hdr = sock.recv(2)
    except Exception:
        return None
    if len(hdr) < 2:
        return None
    opcode = hdr[0] & 0x0F
    masked = (hdr[1] & 0x80) != 0
    length = hdr[1] & 0x7F
    idx = 2
    if length == 126:
        length = struct.unpack(">H", sock.recv(2))[0]; idx += 2
    elif length == 127:
        length = struct.unpack(">Q", sock.recv(8))[0]; idx += 8
    mask = None
    if masked:
        mask = sock.recv(4)
    payload = b""
    while len(payload) < length:
        chunk = sock.recv(length - len(payload))
        if not chunk:
            return None
        payload += chunk
    if masked and mask:
        payload = bytes(payload[i] ^ mask[i % 4] for i in range(len(payload)))
    return opcode, payload

def handle_conn(conn_sock, addr):
    try:
        handshake = b""
        while b"\r\n\r\n" not in handshake:
            handshake += conn_sock.recv(4096)
        lines = handshake.split(b"\r\n")
        key = ""
        path = "/LOBBY"
        for ln in lines:
            if ln.lower().startswith(b"sec-websocket-key:"):
                key = ln.split(b":", 1)[1].strip().decode()
            if ln.startswith(b"GET "):
                path = ln.split(b" ", 2)[1].decode()
        code = path.split("?")[0].lstrip("/").upper() or "LOBBY"
        code = code[:16]
        accept = ws_accept(key)
        conn_sock.sendall(
            b"HTTP/1.1 101 Switching Protocols\r\n"
            b"Upgrade: websocket\r\nConnection: Upgrade\r\n"
            b"Sec-WebSocket-Accept: " + accept.encode() + b"\r\n\r\n"
        )
        cid = uuid.uuid4().hex[:8]
        rset = get_room(code)
        conn = {"sock": conn_sock, "id": cid}
        with rooms_lock:
            rset.add(conn)
        print(f"[+] 玩家 {cid} 进入房间 {code}（当前 {len(rset)} 人）")
        while True:
            f = recv_frame(conn_sock)
            if f is None:
                break
            opcode, payload = f
            if opcode == 0x8:
                break
            if opcode == 0x9:
                send_frame(conn_sock, payload, 0xA)
                continue
            if opcode in (0x1, 0x0):
                try:
                    msg = json.loads(payload.decode("utf-8", "ignore"))
                except Exception:
                    msg = {"raw": payload.decode("utf-8", "ignore")}
                msg["_from"] = cid
                out = json.dumps(msg).encode("utf-8")
                with rooms_lock:
                    for c in list(rset):
                        if c is not conn:
                            send_frame(c["sock"], out, 0x1)
    except Exception:
        pass
    finally:
        with rooms_lock:
            rset.discard(conn)
        try:
            conn_sock.close()
        except Exception:
            pass
        print(f"[-] 玩家 {cid} 离开房间 {code}")

# ---------- HTTP 状态页 ----------
class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        code = gen_code()
        invite = f"ws://{lan_ip()}:{PORT}/{code}"
        html = f"""<!doctype html><meta charset=utf-8><title>四方洲服务器</title>
<body style="font-family:sans-serif;background:#0a0e17;color:#eee;padding:40px">
<h1>四方洲联机服务器运行中 ✅</h1>
<p>端口：{PORT}</p>
<p>局域网地址：<code>ws://{lan_ip()}:{PORT}</code></p>
<p>朋友连这个地址进同一个房间即可联机。每位腐竹可在游戏联机框填：</p>
<pre style="background:#111827;padding:14px;border-radius:10px">ws://{lan_ip()}:{PORT}/{code}  ← 自动生成的房间邀请码：{code}</pre>
<p style="color:#9ca3af">提示：异地联机需把本机映射到公网（内网穿透/端口转发），再把公网地址发给朋友。</p>
</body>"""
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.end_headers()
        self.wfile.write(html.encode("utf-8"))
    def log_message(self, *a):
        pass

class Server(ThreadingMixIn, ThreadingHTTPServer):
    daemon_threads = True

# ---------- 局域网 UDP 信标 ----------
def beacon_loop():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    ip = lan_ip()
    bc = json.dumps({"game": "sifangzhou", "port": PORT, "ip": ip}).encode()
    while True:
        try:
            s.sendto(bc, ("255.255.255.255", 8788))
        except Exception:
            pass
        threading.Event().wait(2)

if __name__ == "__main__":
    code = gen_code()
    print("=" * 40)
    print("  四方洲联机服务器已启动（Python）")
    print("  本机端口   :", PORT)
    print("  局域网地址 :", f"ws://{lan_ip()}:{PORT}")
    print("  自动邀请码 :", code, " → ", f"ws://{lan_ip()}:{PORT}/{code}")
    print("  朋友在游戏联机框填上面地址即可进同一房间")
    print("=" * 40)
    threading.Thread(target=beacon_loop, daemon=True).start()
    httpd = Server(("0.0.0.0", PORT), Handler)
    httpd.serve_forever()