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")


class PlayerWebSocketHandler(object):
    def __init__(self, custom_id: str, label: str):
        self.custom_id = custom_id
        self.label = label
        self.token = None
        self.ws = None
        self.listener_task = None

    # ---------- Auth & Connect ----------

    def login(self):
        """Authenticate via custom ID and store token."""
        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": self.custom_id},
        )
        r.raise_for_status()
        self.token = r.json()["token"]
        return self.token

    async def connect(self):
        """Open Nakama WebSocket using stored auth token."""
        if not self.token:
            self.login()

        url = f"{WS}/ws?token={self.token}"
        self.ws = await websockets.connect(url)

    # ---------- Listener ----------

    async def listener(self):
        """Continuously receive events + decode match data."""
        try:
            while True:
                raw = await self.ws.recv()
                msg = json.loads(raw)

                if "match_data" not in msg:
                    print(f"[{self.label}] {msg}")
                    continue

                md = msg["match_data"]
                op = int(md.get("op_code", 0))
                payload = json.loads(base64.b64decode(md["data"]).decode())

                if op == 1:
                    print(f"[{self.label}] MOVE -> {payload}")
                elif op == 2:
                    print(f"\n[{self.label}] GAME STATE")
                    print_board(payload["board"])
                    print(f"Turn={payload['turn']}  Winner={payload['winner']}\n")
                else:
                    print(f"[{self.label}] UNKNOWN OPCODE {op}: {payload}")

        except websockets.exceptions.ConnectionClosedOK:
            print(f"[{self.label}] WebSocket closed gracefully")
        except websockets.exceptions.ConnectionClosedError as e:
            print(f"[{self.label}] WebSocket closed with error: {e}")

    def start_listener(self):
        """Spawn background task for async message stream."""
        if self.ws:
            self.listener_task = asyncio.create_task(self.listener())

    # ---------- Match Helpers ----------

    async def create_match(self) -> str:
        await self.ws.send(json.dumps({"match_create": {}}))
        msg = json.loads(await self.ws.recv())

        match_id = msg["match"]["match_id"]
        print(f"[{self.label}] Created match: {match_id}")
        return match_id

    async def join_match(self, match_id: str):
        await self.ws.send(json.dumps({"match_join": {"match_id": match_id}}))
        print(f"[{self.label}] Joined match: {match_id}")

    # ---------- Gameplay ----------

    async def send_move(self, match_id: str, row: int, col: int):
        payload = {"row": row, "col": col}
        encoded = base64.b64encode(json.dumps(payload).encode()).decode()

        await self.ws.send(json.dumps({
            "match_data_send": {
                "match_id": match_id,
                "op_code": 1,
                "data": encoded,
            }
        }))
        print(f"[{self.label}] Sent move: {payload}")

    # ---------- Cleanup ----------

    async def close(self):
        if self.listener_task:
            self.listener_task.cancel()
        if self.ws:
            await self.ws.close()
            print(f"[{self.label}] Connection closed")


async def main():
    # Create players
    p1 = PlayerWebSocketHandler("player_one_123456", "P1")
    p2 = PlayerWebSocketHandler("player_two_123456", "P2")

    # Connect
    await p1.connect()
    await p2.connect()

    # Match create + join
    match_id = await p1.create_match()
    await p2.join_match(match_id)

    # Start listeners
    p1.start_listener()
    p2.start_listener()

    await asyncio.sleep(1)

    # Play moves
    await p1.send_move(match_id, 0, 0)
    await asyncio.sleep(0.3)
    await p2.send_move(match_id, 1, 1)
    await asyncio.sleep(0.3)
    await p1.send_move(match_id, 0, 1)
    await asyncio.sleep(0.3)
    await p2.send_move(match_id, 2, 2)
    await asyncio.sleep(0.3)
    await p1.send_move(match_id, 0, 2)

    await asyncio.sleep(2)

    await p1.close()
    await p2.close()


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