import asyncio
import json
import logging
from websockets.asyncio.server import serve
from websockets.asyncio.server import broadcast
import re
from passlib.hash import scrypt
import db
import uuid
import secrets

addr = "localhost"
port = 3636

logging.basicConfig(level=logging.INFO)

error_contexts = {
    "malformedJson": "The JSON data sent to the server could not be parsed.",
    "lengthInvalid": "A value in the JSON data is longer or shorter than expected.",
    "invalidUsername": "Username is invalid. It may contain characters that are not permitted in usernames.",
    "invalidInvite": "The invite code you are trying to use is invalid or has expired.",
    "usernameTaken": "This username has been taken.",
    "notExists": "The requested value does not exist."
}

ulist = {}
user_clients = {}
clients = []

invite_codes = ["STANLEYYELNATSAB"]
locked = False

class util:
    def error(code, listener):
        if code in error_contexts:
            context = error_contexts[code]
        else:
            context = ""
        return json.dumps({
            "error": True,
            "code": code,
            "form": "helium-util",
            "context": context,
            "listener": listener
        })
    
    def field_check(expects, gets):
        for i in expects:
            if i not in gets:
                return "malformedJson"
            if type(gets[i]) in [str, dict, list]:
                if len(gets[i]) not in expects[i]:
                    return "lengthInvalid"
        return True

async def handler(websocket):
    clients.append(websocket)
    async for message in websocket:
        try:
            r = json.loads(message)
        except:
            await websocket.send(util.error("malformedJson", None))
            continue
        if "listener" not in r:
            r["listener"] = None
        listener = r["listener"]
        if "command" not in r:
            await websocket.send(util.error("malformedJson", listener))
            continue
        if r["command"] == "register":
            fc = util.field_check({"username": range(1,21), "password": range(8,256), "invite_code": range(16,17)}, r)
            if fc != True:
                await websocket.send(util.error(fc, listener))
                continue
            r["username"] = r["username"].lower()
            r["invite_code"] = r["invite_code"].upper()
            if not re.fullmatch("[a-z0-9-_]{1,20}", r["username"]):
                await websocket.send(util.error("invalidUsername", listener))
                continue
            if r["invite_code"] not in invite_codes:
                await websocket.send(util.error("invalidInvite", listener))
                continue
            if db.acc.get(r["username"]) != "notExists":
                await websocket.send(util.error("usernameTaken", listener))
                continue
            data = {
                "_id": str(uuid.uuid4()),
                "username": r["username"],
                "display_name": r["username"],
                "avatar": None,
                "bot": False,
                "verified": False,
                "banned_until": 0,
                "profile": {
                    "bio": "",
                    "lastfm": "",
                    "banner": None,
                    "links": {}
                },
                "secure": {
                    "password": scrypt.hash(r["password"]),
                    "token": secrets.token_urlsafe(64),
                    "ban_reason": "",
                    "invite_code": r["invite_code"],
                    "support_code": secrets.token_hex(16)
                }
            }
            result = db.acc.add(data)
            if result != True:
                await websocket.send(util.error(result, listener))
                continue
            invite_codes.remove(r["invite_code"])
            await websocket.send(json.dumps({"error": False, "token": data["secure"]["token"], "listener": listener}))
        elif r["command"] == "login_pswd":
            fc = util.field_check({"username": range(1,21), "password": range(8,256)}, r)
            if fc != True:
                await websocket.send(util.error(fc, listener))
                continue
            r["username"] = r["username"].lower()

        else:
            await websocket.send(util.error("malformedJson", listener))
    if websocket in clients:
        clients.append(websocket)

async def main():
    async with serve(handler, addr, port) as server:
        await server.serve_forever()

asyncio.run(main())