import asyncio
import base64
import json

import requests
import websockets

HOST = "http://127.0.0.1:7350"
WS   = "ws://127.0.0.1:7350"
SERVER_KEY = "defaultkey"


def print_board(board):
    # board = [["X","",""], ["","O",""], ["","",""]]
    symbols = lambda c: c if c else " "
    rows = [" | ".join(symbols(c) for c in row) for row in board]
    separator = "\n---------\n"
    print("\n" + separator.join(rows) + "\n")


# ---------- Auth ----------

def login(custom_id: str) -> str:
    """Authenticate via custom ID and return Nakama session token (JWT)."""
    basic = base64.b64encode(f"{SERVER_KEY}:".encode()).decode()
    r = requests.post(
        f"{HOST}/v2/account/authenticate/custom?create=true",
        headers={"Authorization": f"Basic {basic}"},
        json={"id": custom_id},
    )
    r.raise_for_status()
    body = r.json()
    return body["token"]


# ---------- WebSocket helpers ----------

async def connect(token: str) -> websockets.ClientConnection:
    url = f"{WS}/ws?token={token}"
    ws = await websockets.connect(url)
    return ws


async def listener(ws, label: str):
    """Log all messages for a given socket."""
    try:
        while True:
            raw = await ws.recv()
            data = json.loads(raw)
            if "match_data" not in data:
                print(f"[{label}] {data}")
                continue
            md = data["match_data"]
            parse_match_data(md, label)
    except websockets.exceptions.ConnectionClosedOK:
        print(f"[{label}] WebSocket closed cleanly")
        return
    except websockets.exceptions.ConnectionClosedError as e:
        print(f"[{label}] WebSocket closed with error: {e}")
        return

def parse_match_data(md, label: str):
    op = int(md.get("op_code", 0))

    # decode base64 payload
    payload_json = base64.b64decode(md["data"]).decode()
    payload = json.loads(payload_json)

    # OpMove → just log move
    if op == 1:
        print(f"[{label}] MOVE -> {payload}")
        return

    # OpState → pretty-print board
    if op == 2:
        print(f"\n[{label}] GAME STATE")
        print_board(payload["board"])
        print(f"Turn={payload['turn']}  Winner={payload['winner']}\n")
        return

    # other opcodes
    raise RuntimeError(f"[{label}] UNKNOWN OPCODE {op}: {payload}")


async def send_move(ws, match_id: str, row: int, col: int):
    """Send a TicTacToe move using OpMove = 1 with base64-encoded data."""
    payload = {"row": row, "col": col}
    # Nakama expects `data` as bytes -> base64 string in JSON
    data_bytes = json.dumps(payload).encode("utf-8")
    data_b64 = base64.b64encode(data_bytes).decode("ascii")

    msg = {
        "match_data_send": {
            "match_id": match_id,
            "op_code": 1,      # OpMove
            "data": data_b64,
        }
    }
    await ws.send(json.dumps(msg))


# ---------- Main flow ----------

async def main():
    # 1) Login 2 players
    token1 = login("player_one_123456")
    token2 = login("player_two_123456")

    # 2) Connect sockets
    ws1 = await connect(token1)
    ws2 = await connect(token2)

    # 3) Create a match from P1
    await ws1.send(json.dumps({"match_create": {}}))
    raw = await ws1.recv()
    msg = json.loads(raw)
    match_id = msg["match"]["match_id"]
    print("Match:", match_id)

    # 4) Only P2 explicitly joins (creator is auto-joined)
    await ws2.send(json.dumps({"match_join": {"match_id": match_id}}))

    # 5) Start listeners
    asyncio.create_task(listener(ws1, "P1"))
    asyncio.create_task(listener(ws2, "P2"))

    # Give server time to process joins and initial state
    await asyncio.sleep(1)

    # 6) Play a quick winning game for P1 (X)
    # P1: (0,0)
    await send_move(ws1, match_id, 0, 0)
    await asyncio.sleep(0.3)

    # P2: (1,1)
    await send_move(ws2, match_id, 1, 1)
    await asyncio.sleep(0.3)

    # P1: (0,1)
    await send_move(ws1, match_id, 0, 1)
    await asyncio.sleep(0.3)

    # P2: (2,2)
    await send_move(ws2, match_id, 2, 2)
    await asyncio.sleep(0.3)

    # P1: (0,2) -> X wins by top row
    await send_move(ws1, match_id, 0, 2)

    # Wait to receive the final state broadcast (OpState = 2)
    await asyncio.sleep(2)

    await ws1.close()
    await ws2.close()


if __name__ == "__main__":
    asyncio.run(main())