import base64
import json
import os.path
import random
import time
import uuid
from pathlib import Path
from urllib.parse import urlparse

import aiohttp
import jwt
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.hashes import SHA256
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

from snek.system.service import BaseService

# The only reason to persist the keys is to be able to use them in the web push

PRIVATE_KEY_FILE = Path("./notification-private.pem")
PRIVATE_KEY_PKCS8_FILE = Path("./notification-private.pkcs8.pem")
PUBLIC_KEY_FILE = Path("./notification-public.pem")


def generate_private_key():
    if not PRIVATE_KEY_FILE.exists():
        private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())

        pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )

        PRIVATE_KEY_FILE.write_bytes(pem)


def generate_pcks8_private_key():
    if not PRIVATE_KEY_PKCS8_FILE.exists():
        private_key = serialization.load_pem_private_key(
            PRIVATE_KEY_FILE.read_bytes(), password=None, backend=default_backend()
        )

        pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )

        PRIVATE_KEY_PKCS8_FILE.write_bytes(pem)


def generate_public_key():
    if not PUBLIC_KEY_FILE.exists():
        private_key = serialization.load_pem_private_key(
            PRIVATE_KEY_FILE.read_bytes(), password=None, backend=default_backend()
        )

        public_key = private_key.public_key()

        pem = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )

        PUBLIC_KEY_FILE.write_bytes(pem)


def ensure_certificates():
    generate_private_key()
    generate_pcks8_private_key()
    generate_public_key()


def hkdf(input_key, salt, info, length):
    return HKDF(
        algorithm=SHA256(),
        length=length,
        salt=salt,
        info=info,
        backend=default_backend(),
    ).derive(input_key)


def _browser_base64(data):
    return base64.urlsafe_b64encode(data).decode("utf-8").rstrip("=")


class PushService(BaseService):
    mapper_name = "push"

    private_key_pem = None
    public_key = None
    public_key_base64 = None

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        ensure_certificates()

        private_key = serialization.load_pem_private_key(
            PRIVATE_KEY_FILE.read_bytes(), password=None, backend=default_backend()
        )
        self.private_key_pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )

        self.public_key = serialization.load_pem_public_key(
            PUBLIC_KEY_FILE.read_bytes(), backend=default_backend()
        )

        self.public_key_base64 = _browser_base64(
            self.public_key.public_bytes(
                encoding=serialization.Encoding.X962,
                format=serialization.PublicFormat.UncompressedPoint,
            )
        )

    def create_notification_authorization(self, push_url):
        target = urlparse(push_url)
        aud = f"{target.scheme}://{target.netloc}"
        sub = "mailto:admin@molodetz.nl"

        identifier = str(uuid.uuid4())

        print(
            f"Creating notification authorization for {aud} with identifier {identifier}"
        )

        return jwt.encode(
            {
                "sub": sub,
                "aud": aud,
                "exp": int(time.time()) + 60 * 60,
                "nbf": int(time.time()),
                "iat": int(time.time()),
                "jti": identifier,
            },
            self.private_key_pem,
            algorithm="ES256",
        )

    def create_notification_info_with_payload(
        self, endpoint: str, auth: str, p256dh: str, payload: str
    ):
        message_private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())

        message_public_key_bytes = message_private_key.public_key().public_bytes(
            encoding=serialization.Encoding.X962,
            format=serialization.PublicFormat.UncompressedPoint,
        )

        salt = os.urandom(16)

        user_key_bytes = base64.urlsafe_b64decode(p256dh + "==")
        shared_secret = message_private_key.exchange(
            ec.ECDH(),
            ec.EllipticCurvePublicKey.from_encoded_point(
                ec.SECP256R1(), user_key_bytes
            ),
        )

        encryption_key = hkdf(
            shared_secret,
            base64.urlsafe_b64decode(auth + "=="),
            b"Content-Encoding: auth\x00",
            32,
        )

        context = (
            b"P-256\x00"
            + len(user_key_bytes).to_bytes(2, "big")
            + user_key_bytes
            + len(message_public_key_bytes).to_bytes(2, "big")
            + message_public_key_bytes
        )

        nonce = hkdf(encryption_key, salt, b"Content-Encoding: nonce\x00" + context, 12)
        content_encryption_key = hkdf(
            encryption_key, salt, b"Content-Encoding: aesgcm\x00" + context, 16
        )

        padding_length = random.randint(0, 16)
        padding = padding_length.to_bytes(2, "big") + b"\x00" * padding_length

        data = AESGCM(content_encryption_key).encrypt(
            nonce, padding + payload.encode("utf-8"), None
        )

        return {
            "headers": {
                "Authorization": f"WebPush {self.create_notification_authorization(endpoint)}",
                "Crypto-Key": f"dh={_browser_base64(message_public_key_bytes)}; p256ecdsa={self.public_key_base64}",
                "Encryption": f"salt={_browser_base64(salt)}",
                "Content-Encoding": "aesgcm",
                "Content-Length": str(len(data)),
                "Content-Type": "application/octet-stream",
            },
            "data": data,
        }

    async def notify_user(self, user_uid: str, payload: dict):
        async with aiohttp.ClientSession() as session:
            async for subscription in self.find(user_uid=user_uid):
                endpoint = subscription["endpoint"]
                key_auth = subscription["key_auth"]
                key_p256dh = subscription["key_p256dh"]

                notification_info = self.create_notification_info_with_payload(
                    endpoint, key_auth, key_p256dh, json.dumps(payload)
                )

                headers = {
                    **notification_info["headers"],
                    "TTL": "60",
                }
                data = notification_info["data"]

                async with session.post(
                    endpoint,
                    headers=headers,
                    data=data,
                ) as response:
                    if response.status == 201 or response.status == 200:
                        print(
                            f"Notification sent to user {user_uid} via endpoint {endpoint}"
                        )
                    else:
                        print(
                            f"Failed to send notification to user {user_uid} via endpoint {endpoint}: {response.status}"
                        )

    async def register(
        self, user_uid: str, endpoint: str, key_auth: str, key_p256dh: str
    ):
        if await self.exists(
            user_uid=user_uid,
            endpoint=endpoint,
            key_auth=key_auth,
            key_p256dh=key_p256dh,
        ):
            return

        model = await self.new()
        model["user_uid"] = user_uid
        model["endpoint"] = endpoint
        model["key_auth"] = key_auth
        model["key_p256dh"] = key_p256dh

        print(
            f"Registering push subscription for user {user_uid} with endpoint {endpoint}"
        )

        if await self.save(model=model) and model:
            print(
                f"Push subscription registered for user {user_uid} with endpoint {endpoint}"
            )

            return model

        raise Exception(
            f"Failed to register push subscription for user {user_uid} with endpoint {endpoint}"
        )
