import asyncio
import base64
import math
import sys
import threading
import time

import cv2
import numpy as np
from websockets.asyncio.client import connect

playing = True


def display_base64_image(b64_string):
    # Remove the data URI prefix if it exists (e.g., "data:image/png;base64,")
    if "," in b64_string:
        b64_string = b64_string.split(",")[1]

    # Convert base64 to bytes and decode into a NumPy array
    img_bytes = base64.b64decode(b64_string)
    np_arr = np.frombuffer(img_bytes, np.uint8)

    # Decode array into an OpenCV image
    img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

    if img is not None:
        cv2.imshow(window_name, img)
    else:
        print("Error: Invalid image data")


def show_img():
    global playing

    try:
        while playing:
            display_base64_image(base64_png)
            key = cv2.waitKey(30) & 0xFF
            if key == ord("q"):
                playing = False
                break
    finally:
        cv2.destroyAllWindows()


def run_async():
    asyncio.run(hello(id, seed))


async def hello(id, seed):
    print(f"ID:*{id}*")
    uri = "ws://ai-challenger.net:5555/traffic/client?ID=" + id + "&ARG0=" + seed
    pair_sems = dict()
    odd_sems = dict()
    counter = 0

    async with connect(uri) as websocket:
        cur_x = 0
        cur_y = 0
        next_x = 0
        next_y = 0
        cur_angle = 0
        while playing:
            response = await websocket.recv()
            counter = counter + 1
            lines = response.split("\n")
            # print(f"Received: {response}")
            for line in lines:
                if line.startswith("moving"):
                    print(f"Received: {line}")
                if line.startswith("sem"):
                    # print(f"Received: {line}")
                    tokens = line.removeprefix("sem:")
                    id = int(tokens.split(",")[0])
                    if id % 2 == 0:
                        pair_sems[id] = int(tokens.split(",")[1])
                    else:
                        odd_sems[id] = int(tokens.split(",")[1])
                if line.startswith("snap:data:image/png;base64,"):
                    global base64_png
                    base64_png = line.removeprefix("snap:data:image/png;base64,")
            if counter % 150 == 0:
                for key, value in pair_sems.items():
                    pair_sems[key] = int(counter / 150) % 2
                    await websocket.send("sem:" + str(key) + ", " + str(value))
                for key, value in odd_sems.items():
                    odd_sems[key] = int(counter / 150) % 2 + 1
                    await websocket.send("sem:" + str(key) + ", " + str(value))


id = sys.argv[1]
seed = sys.argv[2]
base64_png = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
window_name = "Game capture"


thread1 = threading.Thread(target=run_async)
thread1.start()
show_img()