# retoor <retoor@molodetz.nl>

import base64
import fcntl
import json
import logging
import os
import random
import time
from typing import Any
from urllib.parse import urlparse

import httpx
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 devplacepy.config import (
    SECONDS_PER_DAY,
    VAPID_PRIVATE_KEY_FILE,
    VAPID_PRIVATE_KEY_PKCS8_FILE,
    VAPID_PUBLIC_KEY_FILE,
    VAPID_SUB,
)
from devplacepy.database import get_setting
from devplacepy.push.providers.base import (
    ACCEPTED,
    DEAD,
    REJECTED,
    Delivery,
    PushProvider,
)
from devplacepy.services.base import ConfigField
from devplacepy.utils import generate_uid

logger = logging.getLogger(__name__)

JWT_LIFETIME_SECONDS = 60 * 60
PUSH_TTL_SECONDS = str(SECONDS_PER_DAY)
DEAD_SUBSCRIPTION_STATUSES = (404, 410)
ACCEPTED_STATUSES = (200, 201)
SUBJECT_KEY = "push_webpush_subject"
PROVIDER_LABEL = "Web Push (VAPID)"


def generate_private_key() -> None:
    if not VAPID_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(),
        )
        VAPID_PRIVATE_KEY_FILE.write_bytes(pem)
        logger.info("Generated VAPID private key at %s", VAPID_PRIVATE_KEY_FILE)


def generate_pkcs8_private_key() -> None:
    if not VAPID_PRIVATE_KEY_PKCS8_FILE.exists():
        private_key = serialization.load_pem_private_key(
            VAPID_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(),
        )
        VAPID_PRIVATE_KEY_PKCS8_FILE.write_bytes(pem)
        logger.info(
            "Generated VAPID PKCS8 private key at %s", VAPID_PRIVATE_KEY_PKCS8_FILE
        )


def generate_public_key() -> None:
    if not VAPID_PUBLIC_KEY_FILE.exists():
        private_key = serialization.load_pem_private_key(
            VAPID_PRIVATE_KEY_FILE.read_bytes(),
            password=None,
            backend=default_backend(),
        )
        pem = private_key.public_key().public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        VAPID_PUBLIC_KEY_FILE.write_bytes(pem)
        logger.info("Generated VAPID public key at %s", VAPID_PUBLIC_KEY_FILE)


def ensure_certificates() -> None:
    if (
        VAPID_PRIVATE_KEY_FILE.exists()
        and VAPID_PRIVATE_KEY_PKCS8_FILE.exists()
        and VAPID_PUBLIC_KEY_FILE.exists()
    ):
        return
    lock_path = VAPID_PRIVATE_KEY_FILE.parent / ".vapid.lock"
    with open(lock_path, "w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        try:
            generate_private_key()
            generate_pkcs8_private_key()
            generate_public_key()
        finally:
            fcntl.flock(lock, fcntl.LOCK_UN)


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


def browser_base64(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode("utf-8").rstrip("=")


_keys: dict[str, Any] = {}


def _load_keys() -> dict[str, Any]:
    if _keys:
        return _keys
    ensure_certificates()
    private_key = serialization.load_pem_private_key(
        VAPID_PRIVATE_KEY_FILE.read_bytes(), password=None, backend=default_backend()
    )
    public_key = serialization.load_pem_public_key(
        VAPID_PUBLIC_KEY_FILE.read_bytes(), backend=default_backend()
    )
    uncompressed_point = public_key.public_bytes(
        encoding=serialization.Encoding.X962,
        format=serialization.PublicFormat.UncompressedPoint,
    )
    _keys["private_key_pem"] = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption(),
    )
    _keys["public_key_point"] = uncompressed_point
    _keys["public_key_base64"] = browser_base64(uncompressed_point)
    logger.debug("Loaded VAPID key material into cache")
    return _keys


def public_key_standard_b64() -> str:
    point = _load_keys()["public_key_point"]
    return base64.b64encode(point).decode("utf-8").rstrip("=")


def subject() -> str:
    return get_setting(SUBJECT_KEY, "").strip() or VAPID_SUB


def create_notification_authorization(push_url: str) -> str:
    target = urlparse(push_url)
    audience = f"{target.scheme}://{target.netloc}"
    issued_at = int(time.time())
    return jwt.encode(
        {
            "sub": subject(),
            "aud": audience,
            "exp": issued_at + JWT_LIFETIME_SECONDS,
            "nbf": issued_at,
            "iat": issued_at,
            "jti": generate_uid(),
        },
        _load_keys()["private_key_pem"],
        algorithm="ES256",
    )


def create_notification_info_with_payload(
    endpoint: str, auth: str, p256dh: str, payload: str
) -> dict[str, Any]:
    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 {create_notification_authorization(endpoint)}",
            "Crypto-Key": f"dh={browser_base64(message_public_key_bytes)}; p256ecdsa={_load_keys()['public_key_base64']}",
            "Encryption": f"salt={browser_base64(salt)}",
            "Content-Encoding": "aesgcm",
            "Content-Length": str(len(data)),
            "Content-Type": "application/octet-stream",
        },
        "data": data,
    }


class WebPushProvider(PushProvider):
    name = "webpush"
    label = PROVIDER_LABEL
    config_fields = [
        ConfigField(
            SUBJECT_KEY,
            "VAPID subject",
            type="str",
            default=VAPID_SUB,
            help="Contact sent as the JWT sub claim, a mailto: or https: URL. Blank uses the built-in default.",
            group=PROVIDER_LABEL,
        )
    ]

    def is_configured(self) -> bool:
        return True

    def client_config(self) -> dict[str, Any]:
        try:
            return {"publicKey": public_key_standard_b64()}
        except Exception as exc:
            logger.error("VAPID key material unavailable: %s", exc)
            return {}

    def parse_registration(self, body: dict[str, Any]) -> dict[str, Any] | None:
        keys = body.get("keys")
        if not isinstance(keys, dict):
            return None
        endpoint = body.get("endpoint")
        key_auth = keys.get("auth")
        key_p256dh = keys.get("p256dh")
        if not (
            isinstance(endpoint, str)
            and isinstance(key_auth, str)
            and isinstance(key_p256dh, str)
            and endpoint
            and key_auth
            and key_p256dh
        ):
            return None
        return {
            "endpoint": endpoint,
            "key_auth": key_auth,
            "key_p256dh": key_p256dh,
        }

    def prepare(self, payload: dict[str, Any]) -> str:
        return json.dumps(payload)

    async def deliver(
        self, client: httpx.AsyncClient, registration: dict[str, Any], prepared: str
    ) -> Delivery:
        endpoint = registration.get("endpoint") or ""
        if not endpoint:
            return Delivery(DEAD, "missing endpoint")
        try:
            notification_payload = create_notification_info_with_payload(
                endpoint,
                registration["key_auth"],
                registration["key_p256dh"],
                prepared,
            )
            headers = {**notification_payload["headers"], "TTL": PUSH_TTL_SECONDS}
            response = await client.post(
                endpoint, headers=headers, content=notification_payload["data"]
            )
        except (httpx.HTTPError, ValueError, KeyError, TypeError) as exc:
            return Delivery(REJECTED, str(exc))
        if response.status_code in ACCEPTED_STATUSES:
            return Delivery(ACCEPTED)
        if response.status_code in DEAD_SUBSCRIPTION_STATUSES:
            return Delivery(DEAD, str(response.status_code))
        return Delivery(REJECTED, str(response.status_code))
