"""
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,
    Ed25519PublicKey,
)
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 ExpiredSignatureError as e:
        raise PermissionError("token expired") from e
    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)
        verify_key = Ed25519PublicKey.from_public_bytes(pubkey_bytes)
        claims = jwt.decode(
            token,
            verify_key,
            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
