"""Verificación de un webhook de FCM (JWS compacto) ANTES de parsear el cuerpo.

Requiere ``cryptography``. El JWKS se pide por HTTPS y se cachea; ante un ``kid`` desconocido, refrescá como máximo
una vez por minuto. Pasos: clave por ``kid`` con el ``alg`` del JWKS, firma, ``iat``/``exp`` (5 minutos con 30 s de
margen), ``body_sha256`` sobre los bytes crudos, ``aud`` igual a tu ``entity_id`` y ``jti`` no visto (anti-replay,
OBLIGATORIO). Los claims desconocidos se ignoran: FCM puede sumar claims sin aviso de versión.
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import threading
import time
from dataclasses import dataclass
from typing import Any, Iterable, Protocol

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec, ed25519
from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature

CLOCK_SKEW_SECONDS = 30
MAX_AGE_SECONDS = 300
REPLAY_TTL_SECONDS = 600


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


def _b64u(value: bytes) -> str:
    return base64.urlsafe_b64encode(value).rstrip(b"=").decode("ascii")


@dataclass(frozen=True)
class Jwk:
    kty: str
    crv: str
    x: str
    kid: str
    alg: str
    use: str
    y: str | None = None


@dataclass(frozen=True)
class WebhookClaims:
    iss: str
    aud: str
    iat: int
    exp: int
    jti: str
    evt: str
    body_sha256: str


class ReplayGuard(Protocol):
    """Anti-replay del ``jti``: devuelve False si ya lo vio en los últimos 10 minutos; si no, lo guarda.

    En producción usá un almacén compartido y persistente (Postgres, Redis) y guardá el ``jti`` en la misma
    transacción que el efecto de negocio.
    """

    def check_and_store(self, jti: str, now: int) -> bool: ...


class InMemoryReplayGuard:
    """Guarda de un solo proceso, segura con hilos. Sólo para empezar: se pierde al reiniciar."""

    def __init__(self) -> None:
        self._seen: dict[str, int] = {}
        self._lock = threading.Lock()

    def check_and_store(self, jti: str, now: int) -> bool:
        with self._lock:
            for seen_jti, seen_expires in list(self._seen.items()):
                if seen_expires <= now:
                    del self._seen[seen_jti]
            if jti in self._seen:
                return False
            self._seen[jti] = now + REPLAY_TTL_SECONDS
            return True


def _segment(value: str) -> dict[str, Any]:
    parsed = json.loads(_b64u_decode(value))
    if not isinstance(parsed, dict):
        raise ValueError("malformed JWS segment")
    return parsed


def _text(claims: dict[str, Any], name: str) -> str:
    value = claims.get(name)
    if not isinstance(value, str) or not value:
        raise ValueError(f"missing claim {name}")
    return value


def _integer(claims: dict[str, Any], name: str) -> int:
    value = claims.get(name)
    if not isinstance(value, int) or isinstance(value, bool):
        raise ValueError(f"invalid claim {name}")
    return value


def _known_claims(claims: dict[str, Any]) -> WebhookClaims:
    # Sólo los claims conocidos: uno nuevo (cambio aditivo) no rompe la verificación.
    return WebhookClaims(
        iss=_text(claims, "iss"),
        aud=_text(claims, "aud"),
        iat=_integer(claims, "iat"),
        exp=_integer(claims, "exp"),
        jti=_text(claims, "jti"),
        evt=_text(claims, "evt"),
        body_sha256=_text(claims, "body_sha256"),
    )


def _public_key(jwk: Jwk):
    x = _b64u_decode(jwk.x)
    if jwk.alg == "EdDSA" and (jwk.kty, jwk.crv) == ("OKP", "Ed25519"):
        return ed25519.Ed25519PublicKey.from_public_bytes(x)
    if jwk.alg == "ES256" and (jwk.kty, jwk.crv) == ("EC", "P-256") and jwk.y is not None:
        numbers = ec.EllipticCurvePublicNumbers(int.from_bytes(x, "big"), int.from_bytes(_b64u_decode(jwk.y), "big"), ec.SECP256R1())
        return numbers.public_key()
    raise ValueError("unsupported or malformed webhook JWK")


def verify_fcm_webhook(
    *,
    token: str,
    raw_body: bytes,
    jwks: Iterable[Jwk],
    entity_id: str,
    replay_guard: ReplayGuard,
    issuer: str = "https://api.fcmpayments.com",
    event_id: str | None = None,
    delivery_id: str | None = None,
    now: int | None = None,
) -> WebhookClaims:
    parts = token.split(".")
    if len(parts) != 3 or any(not part for part in parts):
        raise ValueError("malformed JWS")
    header = _segment(parts[0])
    alg = header.get("alg")
    if alg not in {"EdDSA", "ES256"}:
        raise ValueError("unsupported webhook algorithm")
    key = next((candidate for candidate in jwks if candidate.kid == header.get("kid")), None)
    if key is None or key.alg != alg or key.use != "sig":
        raise ValueError("unknown or mismatched webhook key")
    signature = _b64u_decode(parts[2])
    signing_input = f"{parts[0]}.{parts[1]}".encode("ascii")
    try:
        public = _public_key(key)
        if alg == "EdDSA":
            public.verify(signature, signing_input)
        else:
            if len(signature) != 64:
                raise ValueError("ES256 JWS signature must be 64 bytes")
            der = encode_dss_signature(int.from_bytes(signature[:32], "big"), int.from_bytes(signature[32:], "big"))
            public.verify(der, signing_input, ec.ECDSA(hashes.SHA256()))
    except InvalidSignature as exc:
        raise ValueError("invalid webhook signature") from exc

    claims = _known_claims(_segment(parts[1]))
    current = int(time.time()) if now is None else now
    if (
        claims.iat > current + CLOCK_SKEW_SECONDS
        or claims.iat < current - MAX_AGE_SECONDS - CLOCK_SKEW_SECONDS
        or claims.exp <= current - CLOCK_SKEW_SECONDS
        or claims.exp - claims.iat != MAX_AGE_SECONDS
    ):
        raise ValueError("expired or not-yet-valid webhook")
    if not hmac.compare_digest(_b64u(hashlib.sha256(raw_body).digest()), claims.body_sha256):
        raise ValueError("webhook body hash mismatch")
    if claims.iss != issuer or claims.aud != entity_id:
        raise ValueError("invalid webhook issuer or audience")
    if event_id is not None and claims.evt != event_id:
        raise ValueError("webhook event id mismatch")
    if delivery_id is not None and claims.jti != delivery_id:
        raise ValueError("webhook delivery id mismatch")
    # Último: un token inválido nunca consume el jti de una entrega legítima.
    if not replay_guard.check_and_store(claims.jti, current):
        raise ValueError("webhook replay")
    return claims