"""Edge attestation: proof that a machine request came through a P2HS Edge Worker.

Contract: docs/COMPANION_RUNTIME_CONTRACT.md section 2A.2 (EDGE-1, EDGE-2, C-6).

The Edge (for WordPress, the WP Gateway Worker) is a relay. It signs every request it
forwards to Core:

    X-P2HS-Edge-Attestation: v1;<kid>;<ts>;<nonce>;<mac>
    mac = HMAC-SHA256(secret[kid], method \\n path \\n sha256(body) \\n ts \\n nonce \\n request_id)

Core verifies the signature, the timestamp window and (in the FastAPI guard) nonce
non-reuse. The caller's own credential (installation/runtime bearer) is still validated
separately by the service; the attestation only proves the request passed the Edge.

This module is pure (no FastAPI, no I/O) so the signing rules are testable in isolation
and can be mirrored byte-for-byte by the Worker (`P2HS-WP-Gateway/src/attestation.ts`).
"""
from __future__ import annotations

import hashlib
import hmac
import json
import re
from collections.abc import Mapping
from dataclasses import dataclass

ATTESTATION_HEADER = "X-P2HS-Edge-Attestation"
REQUEST_ID_HEADER = "X-P2HS-Gateway-Request-Id"
VERSION = "v1"
DEFAULT_MAX_SKEW_SECONDS = 60
NONCE_TTL_SECONDS = 2 * DEFAULT_MAX_SKEW_SECONDS + 10

_NONCE_RE = re.compile(r"^[A-Za-z0-9_-]{16,64}$")
_KID_RE = re.compile(r"^[A-Za-z0-9_.-]{1,64}$")

# Edge classes. Each Edge has its own key, so one Edge cannot call another Edge's routes.
EDGE_WP_GATEWAY = "wp_gateway"
EDGE_SHOPIFY = "shopify_edge"
KNOWN_EDGES = frozenset({EDGE_WP_GATEWAY, EDGE_SHOPIFY})


class EdgeAttestationError(Exception):
    """Verification failed. ``code`` is for logs only and is never returned to callers."""

    def __init__(self, code: str):
        super().__init__(code)
        self.code = code


@dataclass(frozen=True)
class EdgeKey:
    kid: str
    secret: str
    edge: str


@dataclass(frozen=True)
class EdgeIdentity:
    kid: str
    edge: str
    nonce: str
    timestamp: int


def parse_keys(raw: str) -> dict[str, EdgeKey]:
    """Parse ``{"<kid>": {"secret": "...", "edge": "wp_gateway"}, ...}``.

    Malformed entries are dropped rather than raising, so a bad rotation entry cannot
    take the whole Core API down; an empty result makes enforce mode fail closed.
    """
    raw = (raw or "").strip()
    if not raw:
        return {}
    try:
        data = json.loads(raw)
    except ValueError:
        return {}
    if not isinstance(data, dict):
        return {}
    keys: dict[str, EdgeKey] = {}
    for kid, entry in data.items():
        if not isinstance(kid, str) or not _KID_RE.match(kid) or not isinstance(entry, dict):
            continue
        secret = entry.get("secret")
        edge = entry.get("edge")
        if not isinstance(secret, str) or len(secret) < 32 or edge not in KNOWN_EDGES:
            continue
        keys[kid] = EdgeKey(kid=kid, secret=secret, edge=edge)
    return keys


def body_sha256(body: bytes) -> str:
    return hashlib.sha256(body or b"").hexdigest()


def compute_mac(
    secret: str,
    *,
    method: str,
    path: str,
    body: bytes,
    timestamp: int,
    nonce: str,
    request_id: str,
) -> str:
    message = "\n".join(
        [method.upper(), path, body_sha256(body), str(int(timestamp)), nonce, request_id or ""]
    )
    return hmac.new(secret.encode(), message.encode(), hashlib.sha256).hexdigest()


def sign(
    key: EdgeKey,
    *,
    method: str,
    path: str,
    body: bytes,
    timestamp: int,
    nonce: str,
    request_id: str = "",
) -> str:
    """Build the header value. Used by tests and tooling; the Worker has its own copy."""
    mac = compute_mac(
        key.secret,
        method=method,
        path=path,
        body=body,
        timestamp=timestamp,
        nonce=nonce,
        request_id=request_id,
    )
    return f"{VERSION};{key.kid};{int(timestamp)};{nonce};{mac}"


def verify(
    header: str | None,
    *,
    method: str,
    path: str,
    body: bytes,
    request_id: str,
    keys: Mapping[str, EdgeKey],
    now: int,
    max_skew_seconds: int = DEFAULT_MAX_SKEW_SECONDS,
) -> EdgeIdentity:
    """Verify signature and freshness. Raises ``EdgeAttestationError``.

    Nonce non-reuse is the caller's job (it needs shared state); see ``edge_guard``.
    """
    if not keys:
        raise EdgeAttestationError("no_keys_configured")
    if not header:
        raise EdgeAttestationError("missing_header")

    parts = header.strip().split(";")
    if len(parts) != 5 or parts[0] != VERSION:
        raise EdgeAttestationError("malformed_header")
    _, kid, ts_raw, nonce, mac = parts

    if not _NONCE_RE.match(nonce):
        raise EdgeAttestationError("malformed_nonce")
    try:
        timestamp = int(ts_raw)
    except ValueError:
        raise EdgeAttestationError("malformed_timestamp") from None

    key = keys.get(kid)
    if key is None:
        raise EdgeAttestationError("unknown_kid")

    expected = compute_mac(
        key.secret,
        method=method,
        path=path,
        body=body,
        timestamp=timestamp,
        nonce=nonce,
        request_id=request_id,
    )
    if not hmac.compare_digest(expected, mac.lower()):
        raise EdgeAttestationError("bad_signature")

    # Freshness is checked after the signature so an unauthenticated caller learns nothing
    # about the clock window.
    if abs(now - timestamp) > max_skew_seconds:
        raise EdgeAttestationError("stale_timestamp")

    return EdgeIdentity(kid=kid, edge=key.edge, nonce=nonce, timestamp=timestamp)