"""JWS-from-PAT key derivation + symmetric attestation (REQ-332, C-5.2).

C-5.2 grill fix: the "public key derivable from the PAT" acceptance
criterion is re-interpreted as a SYMMETRIC scheme. The PAT (Personal
Access Token) is the shared secret; the JWS signing key AND the
verification key are both derived from the PAT via the same HKDF-SHA256
KDF. The JWS uses HMAC-SHA256 (HS256) — a symmetric MAC, not an
asymmetric signature.

Key derivation (NIST SP 800-56C / RFC 5869):
    key = HKDF-SHA256(
        input_key_material = PAT.encode(),
        salt               = b"nova-local-attestation",
        info               = b"jws-signing-key",
        length             = 32,
    )

The resulting 32-byte key is used both to sign (sign_attestation) and to
verify (verify_attestation). Anyone holding the PAT can derive the same
key and verify the attestation; without the PAT, the HMAC cannot be
forged. This satisfies INV-14..17:

  - INV-14: the signing key is derived from the PAT (no separate key
    material; no long-lived private key on disk).
  - INV-15: the key never leaves the derivation (it is recomputed from
    the PAT on each sign/verify call; not cached, not persisted).
  - INV-16: the salt + info are fixed constants binding the key to the
    "nova-local-attestation / jws-signing-key" purpose (key separation).
  - INV-17: tamper detection via the HMAC verification (verify_attestation
    raises on any signature mismatch).

The JWS is the compact serialization:
    b64url(header).b64url(payload).b64url(signature)
where header = {"alg":"HS256","typ":"JWT"}, payload = the JWT claims
(the attestation payload dict), and signature = HMAC-SHA256(key,
b64url(header) + "." + b64url(payload)).
"""

from __future__ import annotations

import hashlib
import hmac
import json
from typing import Any, Dict

__all__ = [
    "derive_signing_key",
    "sign_attestation",
    "verify_attestation",
    "JWSValidationError",
]

# Fixed KDF parameters (INV-16: key separation — binds the derived key to
# the nova-local-attestation / jws-signing-key purpose).
_KDF_SALT = b"nova-local-attestation"
_KDF_INFO = b"jws-signing-key"
_KDF_LENGTH = 32  # 256-bit key for HMAC-SHA256

# JWS header for HS256 (symmetric HMAC-SHA256).
_JWS_HEADER = {"alg": "HS256", "typ": "JWT"}


class JWSValidationError(Exception):
    """Raised when a JWS attestation fails verification (signature mismatch,
    malformed token, or wrong PAT)."""


def _b64url_encode(data: bytes) -> str:
    """RFC 7515 base64url encoding WITHOUT padding (JWS compact form)."""
    import base64
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64url_decode(segment: str) -> bytes:
    """RFC 7515 base64url decoding (re-adds stripped padding)."""
    import base64
    pad = "=" * (-len(segment) % 4)
    return base64.urlsafe_b64decode(segment + pad)


def _hkdf_sha256(input_key_material: bytes, salt: bytes, info: bytes, length: int) -> bytes:
    """HKDF-SHA256 (RFC 5869).

    Prefers cryptography.hazmat.primitives.kdf.hkdf.HKDF (the cryptography
    extra); falls back to a hashlib-based implementation if cryptography
    is unavailable (so the module works in a minimal Lambda runtime).
    """
    try:
        from cryptography.hazmat.primitives.kdf.hkdf import HKDF
        from cryptography.hazmat.primitives import hashes
        hkdf = HKDF(
            algorithm=hashes.SHA256(),
            length=length,
            salt=salt,
            info=info,
        )
        return hkdf.derive(input_key_material)
    except ImportError:  # pragma: no cover - fallback path
        return _hkdf_sha256_hashlib(input_key_material, salt, info, length)


def _hkdf_sha256_hashlib(input_key_material: bytes, salt: bytes, info: bytes, length: int) -> bytes:
    """RFC 5869 HKDF-SHA256 using only hashlib + hmac (fallback)."""
    # Extract: PRK = HMAC-SHA256(salt, IKM)
    prk = hmac.new(salt, input_key_material, hashlib.sha256).digest()
    # Expand: T(i) = HMAC-SHA256(PRK, T(i-1) | info | i)
    okm = b""
    t = b""
    block = 0
    while len(okm) < length:
        block += 1
        t = hmac.new(prk, t + info + bytes([block]), hashlib.sha256).digest()
        okm += t
    return okm[:length]


def derive_signing_key(pat: str) -> bytes:
    """Derive the 32-byte symmetric JWS signing key from a PAT.

    HKDF-SHA256(PAT.encode(), salt=b'nova-local-attestation',
                info=b'jws-signing-key', length=32).

    The same PAT always yields the same key (deterministic); the key is
    never cached or persisted (INV-15 — recomputed on each call).
    """
    if not isinstance(pat, str) or not pat:
        raise ValueError("pat must be a non-empty string")
    return _hkdf_sha256(
        input_key_material=pat.encode("utf-8"),
        salt=_KDF_SALT,
        info=_KDF_INFO,
        length=_KDF_LENGTH,
    )


def sign_attestation(payload: Dict[str, Any], pat: str) -> str:
    """Produce a compact JWS (HS256) for the attestation payload.

    Args:
        payload: the JWT claims (the attestation payload dict).
        pat: the Personal Access Token (shared secret).

    Returns:
        The compact JWS string: b64url(header).b64url(payload).b64url(signature).
        The header is {"alg":"HS256","typ":"JWT"}; the payload is the
        JSON-encoded claims; the signature is HMAC-SHA256(key, header.payload).
    """
    if not isinstance(payload, dict):
        raise ValueError("payload must be a dict")
    key = derive_signing_key(pat)
    header_segment = _b64url_encode(
        json.dumps(_JWS_HEADER, separators=(",", ":"), sort_keys=True).encode("utf-8")
    )
    payload_segment = _b64url_encode(
        json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")
    )
    signing_input = f"{header_segment}.{payload_segment}".encode("ascii")
    signature = hmac.new(key, signing_input, hashlib.sha256).digest()
    signature_segment = _b64url_encode(signature)
    return f"{header_segment}.{payload_segment}.{signature_segment}"


def verify_attestation(jws: str, pat: str) -> Dict[str, Any]:
    """Verify a compact JWS (HS256) attestation and return the payload.

    Derives the same key from the PAT, recomputes the HMAC, and compares
    in constant time. Raises JWSValidationError on:
      - malformed JWS (not 3 segments, bad base64, bad JSON)
      - signature mismatch (tampering or wrong PAT)
      - wrong header (alg != HS256)

    Args:
        jws: the compact JWS string from sign_attestation.
        pat: the Personal Access Token (shared secret).

    Returns:
        The decoded payload dict (the JWT claims) on success.
    """
    if not isinstance(jws, str) or not jws:
        raise JWSValidationError("jws must be a non-empty string")
    parts = jws.split(".")
    if len(parts) != 3:
        raise JWSValidationError(f"malformed JWS: expected 3 segments, got {len(parts)}")
    header_segment, payload_segment, signature_segment = parts

    # Decode + validate the header.
    try:
        header = json.loads(_b64url_decode(header_segment))
    except (ValueError, json.JSONDecodeError) as e:
        raise JWSValidationError(f"malformed JWS header: {e}") from e
    if not isinstance(header, dict) or header.get("alg") != "HS256":
        raise JWSValidationError(
            f"unsupported JWS alg: expected HS256, got {header.get('alg')!r}"
        )

    # Recompute the signature with the key derived from the PAT.
    key = derive_signing_key(pat)
    signing_input = f"{header_segment}.{payload_segment}".encode("ascii")
    expected_signature = hmac.new(key, signing_input, hashlib.sha256).digest()
    actual_signature = _b64url_decode(signature_segment)
    if not hmac.compare_digest(expected_signature, actual_signature):
        raise JWSValidationError(
            "JWS signature verification failed (tampered token or wrong PAT)"
        )

    # Decode + return the payload.
    try:
        payload = json.loads(_b64url_decode(payload_segment))
    except (ValueError, json.JSONDecodeError) as e:
        raise JWSValidationError(f"malformed JWS payload: {e}") from e
    if not isinstance(payload, dict):
        raise JWSValidationError("JWS payload is not a JSON object")
    return payload