"""Web Push, the protocol (RFC 8030 delivery, RFC 8291 encryption, RFC 8292 VAPID).

A browser that wants notifications gives us a push service URL (Google's,
Mozilla's, Apple's, Microsoft's) and two keys. A message is encrypted to those
keys, so the push service only carries it, and signed with our own key pair
(VAPID), so the push service knows it comes from the server the browser
subscribed to. Written here rather than taken from pywebpush, which needs a
newer ``cryptography`` than the rest of the backend pins; it is small, and
``tests/push`` checks it against the RFC 8291 worked example.
"""

from __future__ import annotations

import base64
import json
import os
import time
from dataclasses import dataclass
from urllib.parse import urlsplit

import httpx
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

RECORD_SIZE = 4096
# Push services only: the endpoint comes from the browser, i.e. from the user,
# and the server posts to it — anything else would let a user aim the server at
# an internal address (SSRF). Every browser in use delivers through one of these.
PUSH_HOSTS = (
    "fcm.googleapis.com",  # Chrome, Edge on Android, Opera, Samsung Internet
    "android.googleapis.com",
    ".google.com",  # Chromium builds subscribe at jmt17.google.com and the like
    ".push.services.mozilla.com",  # Firefox
    ".push.apple.com",  # Safari, and installed apps on iPhone/iPad
    ".notify.windows.com",  # Edge on Windows
)


def b64u_decode(value: str) -> bytes:
    return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4))


def b64u_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def allowed_endpoint(endpoint: str) -> bool:
    """Whether ``endpoint`` is an https URL of a known push service."""
    parts = urlsplit(endpoint)
    host = (parts.hostname or "").lower()
    if parts.scheme != "https" or not host or parts.port not in (None, 443):
        return False
    return any(host == h or (h.startswith(".") and host.endswith(h)) for h in PUSH_HOSTS)


def _public_bytes(key: ec.EllipticCurvePrivateKey) -> bytes:
    return key.public_key().public_bytes(
        serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint
    )


def _hkdf(salt: bytes, ikm: bytes, info: bytes, length: int) -> bytes:
    return HKDF(algorithm=hashes.SHA256(), length=length, salt=salt, info=info).derive(ikm)


def encrypt(
    payload: bytes,
    p256dh: str,
    auth: str,
    *,
    salt: bytes | None = None,
    sender: ec.EllipticCurvePrivateKey | None = None,
) -> bytes:
    """``payload`` encrypted for one browser (RFC 8291, ``aes128gcm``), as the
    request body: header (salt, record size, our one-off public key) + one
    record. ``salt`` and ``sender`` are for the RFC test vector only."""
    ua_public = b64u_decode(p256dh)
    auth_secret = b64u_decode(auth)
    sender = sender or ec.generate_private_key(ec.SECP256R1())
    as_public = _public_bytes(sender)
    ua_key = ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), ua_public)
    ecdh_secret = sender.exchange(ec.ECDH(), ua_key)
    ikm = _hkdf(auth_secret, ecdh_secret, b"WebPush: info\x00" + ua_public + as_public, 32)
    salt = salt or os.urandom(16)
    cek = _hkdf(salt, ikm, b"Content-Encoding: aes128gcm\x00", 16)
    nonce = _hkdf(salt, ikm, b"Content-Encoding: nonce\x00", 12)
    # One record, no padding: 0x02 marks the last (only) record.
    record = AESGCM(cek).encrypt(nonce, payload + b"\x02", None)
    header = salt + RECORD_SIZE.to_bytes(4, "big") + bytes([len(as_public)]) + as_public
    return header + record


# ---- VAPID (RFC 8292) ---------------------------------------------------------------


@dataclass(frozen=True)
class Vapid:
    """The server's signing key pair."""

    private: ec.EllipticCurvePrivateKey

    @classmethod
    def generate(cls) -> Vapid:
        return cls(ec.generate_private_key(ec.SECP256R1()))

    @classmethod
    def from_pem(cls, pem: bytes) -> Vapid:
        key = serialization.load_pem_private_key(pem, password=None)
        if not isinstance(key, ec.EllipticCurvePrivateKey):
            raise ValueError("VAPID key is not an EC key")
        return cls(key)

    def pem(self) -> bytes:
        return self.private.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.PKCS8,
            serialization.NoEncryption(),
        )

    @property
    def public_b64(self) -> str:
        """The public key as browsers take it (``applicationServerKey``)."""
        return b64u_encode(_public_bytes(self.private))

    def authorization(self, endpoint: str, subject: str, *, ttl: int = 12 * 3600) -> str:
        """The ``Authorization`` header for a push to ``endpoint``: a JWT (ES256)
        naming the push service's origin, with our public key."""
        parts = urlsplit(endpoint)
        claims = {
            "aud": f"{parts.scheme}://{parts.netloc}",
            "exp": int(time.time()) + ttl,
            "sub": subject,
        }
        header = {"typ": "JWT", "alg": "ES256"}
        signing_input = (
            b64u_encode(json.dumps(header, separators=(",", ":")).encode())
            + "."
            + b64u_encode(json.dumps(claims, separators=(",", ":")).encode())
        )
        der = self.private.sign(signing_input.encode(), ec.ECDSA(hashes.SHA256()))
        r, s = decode_dss_signature(der)
        signature = r.to_bytes(32, "big") + s.to_bytes(32, "big")
        return f"vapid t={signing_input}.{b64u_encode(signature)}, k={self.public_b64}"


# ---- delivery (RFC 8030) ------------------------------------------------------------


class Gone(Exception):
    """The subscription no longer exists (404/410): forget it."""


class Failed(Exception):
    """The push service refused or could not be reached; worth trying later."""


async def send(
    client: httpx.AsyncClient,
    *,
    endpoint: str,
    p256dh: str,
    auth: str,
    payload: dict[str, object],
    vapid: Vapid,
    subject: str,
    ttl: int = 24 * 3600,
    urgency: str = "normal",
    topic: str | None = None,
) -> None:
    """Deliver one encrypted message; raise :class:`Gone` or :class:`Failed`."""
    if not allowed_endpoint(endpoint):
        raise Gone("not a push service")
    body = encrypt(json.dumps(payload, separators=(",", ":")).encode(), p256dh, auth)
    headers = {
        "Authorization": vapid.authorization(endpoint, subject),
        "Content-Encoding": "aes128gcm",
        "Content-Type": "application/octet-stream",
        "TTL": str(ttl),
        "Urgency": urgency,
    }
    if topic:
        headers["Topic"] = topic  # replaces an undelivered message with the same topic
    try:
        resp = await client.post(endpoint, content=body, headers=headers)
    except httpx.HTTPError as err:
        raise Failed(str(err)) from err
    if resp.status_code in (404, 410):
        raise Gone(f"HTTP {resp.status_code}")
    if resp.status_code >= 400:
        raise Failed(f"HTTP {resp.status_code}: {resp.text[:200]}")