"""Signed grant tokens for cross-agent workspace handoff.

A grant is a small, self-contained, signed claim issued by one agent that
the platform (or the receiving agent) can verify without a registry round-trip.

Wire format::

    "<base64url(json(payload))>.<base64url(signature)>"

The payload describes *what* the callee is allowed to do, *whose* workspace
they can see, and *for how long*. The runtime on the receiving side
materializes a :class:`WorkspaceClient` scoped to that grant.

Production deployments should set ``A2A_GRANT_SIGNING_KEY`` only on trusted
platform components and ``A2A_GRANT_VERIFYING_KEY`` on agents that need to
verify grants. ``A2A_PLATFORM_SECRET`` remains a legacy HMAC fallback for
local/dev compatibility.
"""
from __future__ import annotations

import base64
import binascii
import hashlib
import hmac
import json
import os
import secrets
import time
from typing import Any

from pydantic import BaseModel, ConfigDict, Field, NonNegativeFloat, NonNegativeInt, PositiveInt

from .workspace import WorkspaceMode, _normalize_write_prefixes

DEFAULT_TTL_SECONDS = 5 * 60


class GrantInvalid(PermissionError):
    """Raised by :func:`verify_grant` when a grant is bad/expired/forged."""


class Grant(BaseModel):
    """The payload of a signed grant token.

    A grant binds *who* (issuer) gave *whom* (audience) access to *which*
    workspace files (bucket + allow/deny patterns) under *what* mode and
    *how long*. The runtime enforces every line of this payload.
    """

    model_config = ConfigDict(extra="forbid", frozen=True)

    grant_id: str
    issuer: str  # caller agent name or URL
    audience: str  # callee agent name or URL
    bucket: str  # workspace bucket the grant covers
    mode: WorkspaceMode = WorkspaceMode.READ_ONLY
    allow_patterns: tuple[str, ...] = ("**",)
    deny_patterns: tuple[str, ...] = ()
    outputs_prefix: str | None = None  # if set, callee writes only here
    write_prefixes: tuple[str, ...] = ()
    llm_models: tuple[str, ...] = ()
    llm_max_budget_usd: NonNegativeFloat | None = None
    llm_rpm_limit: PositiveInt | None = None
    llm_tpm_limit: PositiveInt | None = None
    expires_at: NonNegativeInt = 0
    issued_at: NonNegativeInt = 0
    nonce: str = Field(default_factory=lambda: secrets.token_hex(8))


def _b64encode(b: bytes) -> str:
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode("ascii")


def _b64decode(s: str) -> bytes:
    clean = s.strip()
    pad = "=" * (-len(clean) % 4)
    decoded = base64.urlsafe_b64decode(clean + pad)
    if _b64encode(decoded) != clean.rstrip("="):
        raise ValueError("non-canonical base64url encoding")
    return decoded


def _platform_secret() -> bytes:
    secret = os.environ.get("A2A_PLATFORM_SECRET", "")
    if not secret:
        raise RuntimeError("A2A_PLATFORM_SECRET is required")
    return secret.encode("utf-8")


def _grant_signing_key_env() -> str:
    return os.environ.get("A2A_GRANT_SIGNING_KEY", "").strip()


def _grant_verifying_key_env() -> str:
    return os.environ.get("A2A_GRANT_VERIFYING_KEY", "").strip()


def _decode_key_material(value: str) -> bytes:
    clean = value.strip()
    if clean.startswith("base64:"):
        clean = clean[len("base64:") :].strip()
    try:
        return base64.b64decode(clean, validate=True)
    except binascii.Error:
        pad = "=" * (-len(clean) % 4)
        return base64.urlsafe_b64decode(clean + pad)


def _ed25519_private_key(value: str) -> Any:
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

    if "BEGIN" in value:
        return serialization.load_pem_private_key(value.encode("utf-8"), password=None)
    return Ed25519PrivateKey.from_private_bytes(_decode_key_material(value))


def _ed25519_public_key(value: str) -> Any:
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

    if "BEGIN" in value:
        return serialization.load_pem_public_key(value.encode("utf-8"))
    return Ed25519PublicKey.from_public_bytes(_decode_key_material(value))


def _ed25519_verifying_key() -> Any | None:
    verifying_key = _grant_verifying_key_env()
    if verifying_key:
        return _ed25519_public_key(verifying_key)
    signing_key = _grant_signing_key_env()
    if signing_key:
        return _ed25519_private_key(signing_key).public_key()
    return None


def mint_grant(
    *,
    issuer: str,
    audience: str,
    bucket: str,
    mode: WorkspaceMode = WorkspaceMode.READ_ONLY,
    allow_patterns: tuple[str, ...] = ("**",),
    deny_patterns: tuple[str, ...] = (),
    outputs_prefix: str | None = None,
    write_prefixes: tuple[str, ...] = (),
    llm_models: tuple[str, ...] = (),
    llm_max_budget_usd: float | None = None,
    llm_rpm_limit: int | None = None,
    llm_tpm_limit: int | None = None,
    ttl_seconds: int = DEFAULT_TTL_SECONDS,
    secret: bytes | None = None,
) -> tuple[Grant, str]:
    """Build a :class:`Grant` and return it together with its signed token."""
    now = int(time.time())
    grant = Grant(
        grant_id=secrets.token_hex(8),
        issuer=issuer,
        audience=audience,
        bucket=bucket,
        mode=mode,
        allow_patterns=tuple(allow_patterns),
        deny_patterns=tuple(deny_patterns),
        outputs_prefix=outputs_prefix,
        write_prefixes=_normalize_write_prefixes(outputs_prefix, write_prefixes),
        llm_models=tuple(llm_models),
        llm_max_budget_usd=llm_max_budget_usd,
        llm_rpm_limit=llm_rpm_limit,
        llm_tpm_limit=llm_tpm_limit,
        expires_at=now + ttl_seconds,
        issued_at=now,
    )
    return grant, sign_grant(grant, secret=secret)


def sign_grant(grant: Grant, *, secret: bytes | None = None) -> str:
    normalized_write_prefixes = _normalize_write_prefixes(
        grant.outputs_prefix,
        grant.write_prefixes,
    )
    if grant.write_prefixes != normalized_write_prefixes:
        grant = grant.model_copy(update={"write_prefixes": normalized_write_prefixes})
    payload = grant.model_dump_json(exclude_none=False).encode("utf-8")
    signing_key = _grant_signing_key_env() if secret is None else ""
    if signing_key:
        sig = _ed25519_private_key(signing_key).sign(payload)
    else:
        sig = hmac.new(secret or _platform_secret(), payload, hashlib.sha256).digest()
    return f"{_b64encode(payload)}.{_b64encode(sig)}"


def verify_grant(token: str, *, secret: bytes | None = None) -> Grant:
    """Parse + verify ``token``. Raises :class:`GrantInvalid` on any failure.

    Checks signature, expiry, and minimal structural shape. Caller-specific
    audience checks are layered on top by the server adapter.
    """
    if not token or "." not in token:
        raise GrantInvalid("malformed grant token")
    payload_b64, sig_b64 = token.rsplit(".", 1)
    try:
        payload = _b64decode(payload_b64)
        sig = _b64decode(sig_b64)
    except (ValueError, base64.binascii.Error) as exc:  # type: ignore[attr-defined]
        raise GrantInvalid(f"grant decode failed: {exc}") from exc

    verifying_key = _ed25519_verifying_key() if secret is None else None
    if verifying_key is not None:
        try:
            verifying_key.verify(sig, payload)
        except Exception as exc:  # noqa: BLE001
            raise GrantInvalid("grant signature mismatch") from exc
    else:
        expected = hmac.new(secret or _platform_secret(), payload, hashlib.sha256).digest()
        if not hmac.compare_digest(expected, sig):
            raise GrantInvalid("grant signature mismatch")

    try:
        data = json.loads(payload)
        grant = Grant.model_validate(data)
    except Exception as exc:  # noqa: BLE001
        raise GrantInvalid(f"grant payload invalid: {exc}") from exc

    if not grant.write_prefixes and grant.outputs_prefix:
        grant = grant.model_copy(
            update={
                "write_prefixes": _normalize_write_prefixes(grant.outputs_prefix)
            }
        )

    if grant.expires_at and grant.expires_at < int(time.time()):
        raise GrantInvalid(f"grant expired at {grant.expires_at}")
    return grant


__all__ = [
    "Grant",
    "GrantInvalid",
    "mint_grant",
    "sign_grant",
    "verify_grant",
    "DEFAULT_TTL_SECONDS",
]
