import base64
import requests
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

API_URL = 'https://prod.la-taniere.fun:30121/'

class GetServerTokenForDiscord:
    derived_key: bytes | None = None

    @staticmethod
    def authenticate(server = API_URL):

        if server is None:
            server = API_URL
        # ==========================
        # Génération clé ECDH client
        # ==========================
        client_private = ec.generate_private_key(ec.SECP256R1())
        client_public = client_private.public_key()

        client_pub_pem = client_public.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )

        # ==========================
        # AUTH
        # ==========================
        auth = requests.post(server + "/api_v2/auth", verify=False, json={
            "client_pub": base64.b64encode(client_pub_pem).decode()
        }).json()

        server_pub = serialization.load_pem_public_key(
            base64.b64decode(auth["server_pub"])
        )

        shared_key = client_private.exchange(ec.ECDH(), server_pub)

        GetServerTokenForDiscord.derived_key = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=None,
            info=b"fivem-private-server"
        ).derive(shared_key)

        return auth["session_id"]

    @staticmethod
    def get_token(session_id: bytes, server = API_URL):
        # ==========================
        # DISCORD TOKEN
        # ==========================
        if server is None:
            server = API_URL
        download = requests.post(server + "/api_v2/tkn_auth", verify=False, headers={
            "x-session-id": session_id
        }).json()

        nonce = base64.b64decode(download["nonce"])
        encrypted_data = base64.b64decode(download["data"])

        aesgcm = AESGCM(GetServerTokenForDiscord.derived_key) # type: ignore[arg-type]
        return aesgcm.decrypt(nonce, encrypted_data, None)

    # @staticmethod
    # def register_discord_user(user_id: str, server = API_URL) -> str:
    #     if server is None:
    #         server = API_URL
    #     registeredId = requests.post(server + "/api_v2/connection/register", verify=False, json={
    #         "x-session-id": user_id
    #     }).json()
    #
    #     return registeredId["discord_id"]
