"""Nova IdP token-vend Lambda — PAT/session → KMS-signed OIDC token
(REQ-336, C-6.1/C-7.1 ABAC FAIL-CLOSED, D-229 revocation).

Accepts a PAT (or session token) and returns a KMS-signed OIDC token
with claims ``sub, aud, iss, exp, iat, jti, roles`` (REQ-336).

## ABAC fail-closed (C-6.1/C-7.1 — INV-17 runtime enforcement)

The grill's #1 finding: the token-vend Lambda MUST fail closed on ABAC
evaluation failure. Concretely, a token is vended **only** when:

1. The PAT is active (``nova-pats.GetItem(jti, ConsistentRead=True)``
   returns an item with ``status == "active"`` — D-229; strong read on
   the main table, GSIs don't support strong reads).
2. ``KyvernoJsonEngine.is_configured()`` returns ``True`` **AND**
   ``evaluate_token_vend_policy()`` returns ``allowed=True`` without
   raising.

If (2) fails for **any** reason — ``kj`` absent, ``kj`` error, policy
parse error, engine raise — the Lambda returns **403** + audit
``token.vend.denied`` (reason ``abac_eval_failed``). **Never fail
open.** This is verified by ``tests/test_abac_fail_closed.py`` — the
most important test of the milestone.

## Dual-use (REQ-329 pattern)

Mirrors ``nova_idp_auth.py``: lazy boto3, env-var table names,
``NOVA_LAMBDA_LOCAL_BYPASS``, ``__main__`` CLI block, audit emission.
"""

from __future__ import annotations

import datetime
import json
import os
import sys
import time

import boto3

# ---------------------------------------------------------------------------
# Config (env-var table names, mirroring nova_idp_auth.py)
# ---------------------------------------------------------------------------

PATS_TABLE = os.environ.get("NOVA_PATS_TABLE", "nova-pats")
SESSIONS_TABLE = os.environ.get("NOVA_SESSIONS_TABLE", "nova-sessions")
OIDC_KMS_KEY_ID = os.environ.get("NOVA_OIDC_KMS_KEY_ID", "alias/nova-oidc-signing")
OIDC_ISSUER = os.environ.get("NOVA_OIDC_ISSUER", "nova-idp")
OIDC_AUDIENCE = os.environ.get("NOVA_OIDC_AUDIENCE", "nova-cli")
# OIDC token lifetime (seconds). Default 15 min.
OIDC_TTL_SECONDS = int(os.environ.get("NOVA_OIDC_TTL_SECONDS", "900"))

_dynamodb = None
_kms_client = None


def _get_dynamodb():
    """Lazy boto3 DynamoDB resource singleton (mirrors contract_ingestor)."""
    global _dynamodb
    if _dynamodb is None:
        _dynamodb = boto3.resource("dynamodb")
    return _dynamodb


def _iso8601_now() -> str:
    return datetime.datetime.now(datetime.timezone.utc).strftime(
        "%Y-%m-%dT%H:%M:%SZ"
    )


def _epoch_now() -> int:
    return int(datetime.datetime.now(datetime.timezone.utc).timestamp())


def _emit_audit(event_type: str, **fields) -> None:
    """Emit an audit event to stderr as JSON (never the raw PAT/token)."""
    payload = {"event": event_type, "ts": _iso8601_now(), **fields}
    # Defense-in-depth: scrub raw token fields (INV-16/INV-17 spirit).
    for _k in ("pat", "session_token", "token", "raw_pat"):
        payload.pop(_k, None)
    sys.stderr.write(json.dumps(payload, sort_keys=True) + "\n")
    sys.stderr.flush()


# ---------------------------------------------------------------------------
# PAT / session decoding (decode WITHOUT verifying — signature verified
# by KMS public key separately at the JWKS verifier; the revocation
# check is the trust anchor here, not the JWT signature).
# ---------------------------------------------------------------------------


def _decode_jwt_unverified(token: str) -> dict:
    """Decode a JWT's payload without verifying the signature."""
    try:
        import jwt as pyjwt
        return pyjwt.decode(token, options={"verify_signature": False})
    except Exception:
        # Fallback: manual base64url decode of the payload segment.
        parts = token.split(".")
        if len(parts) < 2:
            raise ValueError("malformed JWT (expected 3 segments)")
        import base64
        pad = parts[1] + "=" * (-len(parts[1]) % 4)
        return json.loads(base64.urlsafe_b64decode(pad))


def _extract_pat_claims(token: str) -> dict:
    """Decode a PAT/session JWT → extract jti, sub, typ, roles, owner, exp."""
    claims = _decode_jwt_unverified(token)
    required = ("jti", "sub", "exp")
    for f in required:
        if f not in claims:
            raise ValueError(f"token missing claim: {f}")
    return claims


# ---------------------------------------------------------------------------
# Revocation check (D-229 — strong read on the main table)
# ---------------------------------------------------------------------------


def _check_pat_active(jti: str) -> tuple[bool, str]:
    """Return ``(active, reason)``. Strong read on nova-pats main table.

    D-229: GSIs don't support strongly-consistent reads, so the
    revocation check uses ``GetItem(PK=jti, ConsistentRead=True)`` on
    the main table. This satisfies the 60s SLO synchronously (the
    strong read reflects the latest write — revocation is instant).
    """
    table = _get_dynamodb().Table(PATS_TABLE)
    resp = table.get_item(
        TableName=PATS_TABLE,
        Key={"jti": jti},
        ConsistentRead=True,
    )
    item = resp.get("Item")
    if item is None:
        return False, "pat_unknown"
    status = item.get("status", "active")
    if status != "active":
        return False, f"pat_{status}"  # pat_revoked, pat_expired, etc.
    # Expired? (defense-in-depth; TTL may not have reaped it yet)
    expires_at = item.get("expires_at")
    if expires_at is not None:
        try:
            if int(expires_at) < _epoch_now():
                return False, "pat_expired"
        except (ValueError, TypeError):
            pass
    return True, "active"


# ---------------------------------------------------------------------------
# ABAC fail-closed (C-6.1/C-7.1)
# ---------------------------------------------------------------------------


def _build_abac_payload(claims: dict, requested_claims: list[str],
                        target_resource: dict, environment: str,
                        policy_version: str) -> dict:
    """Build the ABAC authorization payload (REQ-339, C-5.1)."""
    return {
        "subject": {
            "id": claims.get("sub", ""),
            "role": (claims.get("roles") or ["unknown"])[0],
            "owner": claims.get("owner", ""),
        },
        "requested_claims": requested_claims,
        "target_resource": target_resource,
        "environment": environment,
        "pat_jti": claims.get("jti", ""),
        "policy_version": policy_version,
    }


def _evaluate_abac_fail_closed(payload: dict) -> tuple[bool, list, str, str]:
    """Evaluate ABAC with fail-closed semantics (C-6.1).

    Returns ``(allowed, pcrs, policy_sha, reason)``. On ANY failure
    (engine not configured, evaluate raises, policy parse error) returns
    ``(False, [], "", "abac_eval_failed")``. **Never fails open.**
    """
    # Lazy imports so the module imports without the engine adapter.
    from core.policy_engine import get_engine

    # C-6.1: is_configured() check. If kj is absent → fail closed.
    try:
        engine = get_engine()
        if not engine.is_configured():
            _emit_audit(
                "token.vend.abac_engine_not_configured",
                pat_jti=payload.get("pat_jti", ""),
            )
            return False, [], "", "abac_eval_failed"
    except Exception:  # noqa: BLE001 - fail closed on any engine check error
        return False, [], "", "abac_eval_failed"

    # C-6.1: evaluate() raising → fail closed.
    try:
        from core.abac_evaluator import evaluate_token_vend_policy
        allowed, pcrs, policy_sha = evaluate_token_vend_policy(payload)
        reason = "abac_denied" if not allowed else "ok"
        return allowed, pcrs, policy_sha, reason
    except Exception:  # noqa: BLE001 - fail closed on any eval error
        return False, [], "", "abac_eval_failed"


# ---------------------------------------------------------------------------
# Token vend (REQ-336)
# ---------------------------------------------------------------------------


def _build_oidc_claims(pat_claims: dict) -> dict:
    """Build the OIDC token claims (REQ-336)."""
    now = _epoch_now()
    return {
        "sub": pat_claims["sub"],
        "aud": OIDC_AUDIENCE,
        "iss": OIDC_ISSUER,
        "exp": now + OIDC_TTL_SECONDS,
        "iat": now,
        "jti": pat_claims.get("jti", ""),  # carry the PAT jti for tracing
        "roles": pat_claims.get("roles", []),
        "typ": "nova_oidc_token",  # INV-14: distinguish from developer_pat
    }


def vend_token(
    token: str,
    requested_claims: list[str] | None = None,
    target_resource: dict | None = None,
    environment: str | None = None,
    policy_version: str = "",
) -> dict:
    """Vend a KMS-signed OIDC token for a PAT/session (REQ-336).

    Returns ``{"token": ..., "expires_at": ...}`` on success. Raises
    ``_DeniedError`` (→ 403) on revocation / ABAC denial.
    """
    requested_claims = requested_claims or ["sub", "roles"]
    environment = environment or "dev"

    # 1. Decode the PAT/session (without verifying — D-229).
    pat_claims = _extract_pat_claims(token)
    jti = pat_claims["jti"]

    # Default target_resource: owner inherits from the PAT subject so
    # the owner-matches ABAC rule passes for same-tenant vends. Callers
    # can override with an explicit target_resource.
    if target_resource is None:
        target_resource = {
            "type": "contract",
            "id": "*",
            "owner": pat_claims.get("owner", "*"),
            "environment": environment,
        }

    # 2. Revocation check (D-229, strong read).
    active, reason = _check_pat_active(jti)
    if not active:
        _emit_audit("token.vend.denied", pat_jti=jti, reason=reason)
        raise _DeniedError(reason)

    # 3. ABAC eval (C-6.1 FAIL-CLOSED).
    abac_payload = _build_abac_payload(
        pat_claims, requested_claims, target_resource, environment, policy_version
    )
    allowed, _pcrs, policy_sha, abac_reason = _evaluate_abac_fail_closed(abac_payload)
    if not allowed:
        _emit_audit(
            "token.vend.denied",
            pat_jti=jti,
            reason=abac_reason,
            policy_sha=policy_sha,
        )
        raise _DeniedError(abac_reason)

    # 4. KMS sign (REQ-337).
    from core.kms_signing import sign_jwt
    oidc_claims = _build_oidc_claims(pat_claims)
    oidc_token = sign_jwt(oidc_claims, key_id=OIDC_KMS_KEY_ID)
    _emit_audit(
        "token.vend.allowed",
        pat_jti=jti,
        sub=oidc_claims["sub"],
        policy_sha=policy_sha,
        expires_at=oidc_claims["exp"],
    )
    return {"token": oidc_token, "expires_at": oidc_claims["exp"]}


class _DeniedError(Exception):
    """Raised on revocation / ABAC denial → 403."""

    def __init__(self, reason: str):
        self.reason = reason
        super().__init__(f"token vend denied: {reason}")


# ---------------------------------------------------------------------------
# Lambda handler + HTTP mapping
# ---------------------------------------------------------------------------


def _to_http_response(result_or_error):
    if isinstance(result_or_error, Exception):
        if isinstance(result_or_error, _DeniedError):
            return {
                "statusCode": 403,
                "body": json.dumps({"error": "token_vend_denied", "reason": result_or_error.reason}),
            }
        if isinstance(result_or_error, ValueError):
            return {
                "statusCode": 400,
                "body": json.dumps({"error": str(result_or_error)}),
            }
        return {
            "statusCode": 500,
            "body": json.dumps({"error": str(result_or_error)}),
        }
    return {"statusCode": 200, "body": json.dumps(result_or_error)}


def lambda_handler(event, context):
    """AWS Lambda handler entry point (thin wrapper, REQ-329 dual-use)."""
    try:
        body = event.get("body", "{}")
        payload = json.loads(body) if isinstance(body, str) else body
        token = payload.get("token") or payload.get("pat") or payload.get("session_token")
        if not token:
            raise ValueError("missing field: token (or pat / session_token)")
        result = vend_token(
            token=token,
            requested_claims=payload.get("requested_claims"),
            target_resource=payload.get("target_resource"),
            environment=payload.get("environment"),
            policy_version=payload.get("policy_version", ""),
        )
        return _to_http_response(result)
    except Exception as e:
        return _to_http_response(e)


# ---------------------------------------------------------------------------
# CLI (dual-use, REQ-329 pattern)
# ---------------------------------------------------------------------------


def cli_main(argv=None):
    """CLI entry point for the token-vend Lambda (REQ-329 dual-use)."""
    raw = argv if argv is not None else sys.argv[1:]
    local_bypass = os.environ.get("NOVA_LAMBDA_LOCAL_BYPASS")
    if not local_bypass:
        os.environ["NOVA_LAMBDA_LOCAL_BYPASS"] = "1"
    try:
        if "--vend-stdin" in raw:
            payload = json.loads(sys.stdin.read())
        elif "--vend" in raw:
            idx = raw.index("--vend")
            path = raw[idx + 1] if idx + 1 < len(raw) else None
            if not path:
                print("Usage: --vend <payload.json>", file=sys.stderr)
                return 2
            with open(path) as fh:
                payload = json.loads(fh.read())
        else:
            print(
                "Usage: python3 -m core.lambda.nova_idp_token_vend "
                "--vend <payload.json> | --vend-stdin < <payload.json>",
                file=sys.stderr,
            )
            return 2
        token = payload.get("token") or payload.get("pat") or payload.get("session_token")
        if not token:
            print("error: missing token in payload", file=sys.stderr)
            return 1
        result = vend_token(
            token=token,
            requested_claims=payload.get("requested_claims"),
            target_resource=payload.get("target_resource"),
            environment=payload.get("environment"),
            policy_version=payload.get("policy_version", ""),
        )
        sys.stdout.write(json.dumps(result, indent=2) + "\n")
        return 0
    except _DeniedError as e:
        sys.stderr.write(f"error: token vend denied ({e.reason})\n")
        return 3  # 403-class
    except ValueError as e:
        sys.stderr.write(f"error: {e}\n")
        return 1
    except Exception as e:  # pragma: no cover - defensive top-level guard
        sys.stderr.write(f"internal error: {e}\n")
        return 2
    finally:
        if not local_bypass:
            os.environ.pop("NOVA_LAMBDA_LOCAL_BYPASS", None)


if __name__ == "__main__":  # pragma: no cover - CLI entry
    sys.exit(cli_main())