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"


# ---------- 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()
    # print("LOGIN:", body)
    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)
            print(f"[{label}] {data}")
    except websockets.exceptions.ConnectionClosedOK:
        print(f"[{label}] WebSocket closed cleanly")
    except websockets.exceptions.ConnectionClosedError as e:
        print(f"[{label}] WebSocket closed with error: {e}")


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 match from P1
    await ws1.send(json.dumps({"match_create": {}}))
    raw = await ws1.recv()
    print("RAW:", raw)
    msg = json.loads(raw)
    match_id = msg["match"]["match_id"]
    print("Match:", match_id)  # keep the dot; it's part of the 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 final state broadcast (OpState = 2)
    await asyncio.sleep(2)

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


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