"""
Z-Arch Service-to-Service Trust
-----------------------------------------------
Each service container has:
  SERVICE_ID           - its unique ID (string)
  S2S_PRIVATE_KEY      - PEM or base64 Ed25519 private key (build-derived)
  S2S_PUBLIC_KEYS_JSON - JSON mapping of trusted service_id → public_key
  S2S_ALLOWED_TARGETS  - comma-separated list of services this service may call

All values are generated by Z-Arch CLI at build time and never fetched at runtime.

Protocol:
  1. Caller signs a short-lived JWT: iss, aud, iat, exp, typ="zarch/s2s".
  2. Target verifies:
       • signature matches issuer's public key
       • aud == this SERVICE_ID
       • iss is allowed in policy
       • token fresh within bounds

Copyright © 2026 RAM Cloud Code LLC
Licensed under the Apache License, Version 2.0.
"""

import os, time, json, jwt
from jwt import ExpiredSignatureError, InvalidTokenError
from pathlib import Path
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
import base64

# ---------------------------------------------------------------------
# Environment
# ---------------------------------------------------------------------
ALG = os.getenv("S2S_ALG", "EdDSA")

# Trusted public keys for other services (JSON: {id: key})
TRUSTED_KEYS = json.loads(os.getenv("S2S_PUBLIC_KEYS_JSON", "{}"))
# Policy: which downstream targets this service can call
ALLOWED_TARGETS = set(
    t.strip() for t in os.getenv("S2S_ALLOWED_TARGETS", "").split(",") if t.strip()
)

S2S_TTL = int(os.getenv("S2S_TTL_SECS", "600"))      # 10 min default
MAX_SKEW = int(os.getenv("S2S_IAT_SKEW_SECS", "60")) # 1 min clock skew

def _must_env(name: str) -> str:
    v = os.getenv(name, "").strip()
    if not v:
        raise RuntimeError(f"Missing required environment variable: {name}")
    return v
# ---------------------------------------------------------------------
# Token minting
# ---------------------------------------------------------------------
def mint_s2s_token(target: str) -> str:
    """Sign and return a JWT authorizing a call to `target`."""

    SERVICE_ID = _must_env("SERVICE_ID")
    PRIVATE_KEY = Ed25519PrivateKey.from_private_bytes(
        base64.b64decode(Path("/zarch/secrets/s2s_private_key").read_text().strip())
    )

    if "*" not in ALLOWED_TARGETS and target not in ALLOWED_TARGETS:
        raise PermissionError(f"{SERVICE_ID} not authorized to call {target}")

    now = int(time.time())
    payload = {
        "iss": SERVICE_ID,
        "aud": target,
        "iat": now,
        "exp": now + S2S_TTL,
        "typ": "zarch/s2s",
        "v": 1,
    }
    
    return jwt.encode(payload, PRIVATE_KEY, algorithm=ALG)


# ---------------------------------------------------------------------
# Token verification
# ---------------------------------------------------------------------
S2S_HEADER = "x-zarch-s2s-token"


def verify_s2s_token(req) -> dict:
    """
    Verify an inbound S2S JWT from the x-zarch-s2s-token header.

    Selects the issuer's public key using the unverified `iss` claim,
    then performs full signature, audience, and freshness validation.

    Returns decoded claims if valid.
    Raises PermissionError on failure.
    """
    SERVICE_ID = _must_env("SERVICE_ID")

    token = req.headers.get(S2S_HEADER, "")
    if not token:
        raise PermissionError("missing S2S token")

    try:
        unverified = jwt.decode(token, options={"verify_signature": False})
    except Exception as e:
        raise PermissionError("invalid token format") from e

    iss = unverified.get("iss")
    if not iss:
        raise PermissionError("missing issuer")

    pubkey = TRUSTED_KEYS.get(iss)
    if not pubkey:
        raise PermissionError(f"unknown issuer {iss}")

    try:
        pubkey_bytes = base64.b64decode(pubkey)
        claims = jwt.decode(
            token,
            pubkey_bytes,
            algorithms=[ALG],
            audience=SERVICE_ID,
            leeway=MAX_SKEW,
        )
    except ExpiredSignatureError as e:
        raise PermissionError("token expired") from e
    except InvalidTokenError as e:
        raise PermissionError("invalid signature") from e

    aud = claims.get("aud")
    iat = int(claims.get("iat", 0))
    now = int(time.time())

    if aud != SERVICE_ID:
        raise PermissionError("audience mismatch")
    if abs(now - iat) > MAX_SKEW:
        raise PermissionError("iat skew too large")

    return claims
