"""Deterministic replay sessions.

A :class:`ReplaySession` is the ordered event log of a single skill run —
enough material to re-render the timeline (so a UI can scrub through it) and,
combined with double-injected side effects, to re-execute deterministically.

Determinism contract
--------------------

The runtime injects a per-run ``random_seed`` into the :class:`RunContext`;
skill code that wants its random draws to survive a replay must use
:meth:`RunContext.random` (which returns a :class:`random.Random` seeded from
that value). Direct ``random.random()`` / ``secrets.token_hex()`` calls bypass
the seed and will diverge under replay.

Wire format mirrors :mod:`a2a_pack.receipts` and :mod:`a2a_pack.grants`::

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

Signing uses Ed25519 when ``A2A_REPLAY_SIGNING_KEY`` is configured (verifying
side reads ``A2A_REPLAY_VERIFYING_KEY``), with HMAC-SHA256 over
``A2A_PLATFORM_SECRET`` as a local/dev fallback.

What this module ships:
    * :class:`ReplayEvent` schema (per-step record).
    * :class:`ReplaySession` schema (the full ordered log + run metadata).
    * :class:`EventRecorder` — append-only in-memory recorder the runtime
      drives during a live run.
    * ``seal_replay_session`` / ``sign_replay_session`` / ``verify_replay_session``
      — the wire primitives.
    * ``iter_events`` — synchronous iterator over a verified session, used to
      drive UI scrub timelines.

What this module deliberately does *not* ship:
    * Re-execution of the agent. That requires the runtime to swap real LLM /
      tool calls for replay-doubles; the session log is the *input* to that
      machine, but the machine lives in :mod:`a2a_pack.runtime` / the server
      adapter.
    * Storage. Sessions are signed inline; the control plane picks where they
      land (object store, Postgres, filesystem).
"""
from __future__ import annotations

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

from pydantic import BaseModel, ConfigDict, Field, NonNegativeInt


class ReplayInvalid(ValueError):
    """Raised by :func:`verify_replay_session` on any failure."""


# Canonical event kinds. The set is a closed enum on the wire (validated by
# pydantic), so adding a new kind is an explicit, reviewable change.
EVENT_KINDS: tuple[str, ...] = (
    "skill_start",
    "llm_call",
    "llm_response",
    "tool_call",
    "tool_response",
    "workspace_read",
    "workspace_write",
    "scope_request",
    "scope_approve",
    "handoff_start",
    "handoff_end",
    "artifact_write",
    "eval",
    "error",
    "skill_end",
)


class ReplayEvent(BaseModel):
    """One event in a replay session.

    Ordering is by ``idx`` (monotonic per session). ``ts_ms`` is the unix
    millisecond stamp of when the event was recorded; useful for the scrub UI
    but not authoritative for ordering (a re-run with cached LLM responses
    will produce earlier stamps).

    ``payload`` stays free-form so the schema can evolve without breaking
    older sessions; UIs render by ``kind`` and inspect ``payload`` defensively.
    """

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

    idx: NonNegativeInt
    kind: str
    ts_ms: NonNegativeInt
    payload: dict[str, Any] = Field(default_factory=dict)


class ReplaySession(BaseModel):
    """A signed, ordered event log for a single skill run."""

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

    session_id: str
    schema_version: int = 1

    # Identity (matches ExecutionReceipt fields so a session and its receipt
    # can be cross-linked by the control plane without a join table).
    agent_name: str
    agent_version: str = ""
    caller: str = ""
    task_id: str = ""
    skill_name: str

    # Determinism inputs. ``random_seed`` is what the runtime injected at
    # start; recording it lets a re-run reproduce the same RNG stream.
    random_seed: str = ""
    input_hash: str = ""

    # Timing
    started_at: NonNegativeInt = 0
    ended_at: NonNegativeInt = 0

    # The log itself
    events: tuple[ReplayEvent, ...] = ()

    # Cross-link to the run's signed receipt, if one was sealed.
    receipt_id: str = ""

    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 _replay_signing_key_env() -> str:
    return os.environ.get("A2A_REPLAY_SIGNING_KEY", "").strip()


def _replay_verifying_key_env() -> str:
    return os.environ.get("A2A_REPLAY_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 = _replay_verifying_key_env()
    if verifying_key:
        return _ed25519_public_key(verifying_key)
    signing_key = _replay_signing_key_env()
    if signing_key:
        return _ed25519_private_key(signing_key).public_key()
    return None


class EventRecorder:
    """Append-only event recorder. The runtime drives this during a run.

    Thread-safety: not guaranteed. A skill invocation is single-task by
    design; concurrent handoffs spawn their own recorder.
    """

    def __init__(
        self,
        *,
        agent_name: str,
        skill_name: str,
        agent_version: str = "",
        caller: str = "",
        task_id: str = "",
        random_seed: str = "",
        input_hash: str = "",
        receipt_id: str = "",
        clock: Any = None,
    ) -> None:
        self._agent_name = agent_name
        self._agent_version = agent_version
        self._skill_name = skill_name
        self._caller = caller
        self._task_id = task_id
        self._random_seed = random_seed
        self._input_hash = input_hash
        self._receipt_id = receipt_id
        self._clock = clock or (lambda: int(time.time() * 1000))
        self._started_at = self._clock() // 1000
        self._events: list[ReplayEvent] = []
        self._idx = 0

    def record(self, kind: str, payload: dict[str, Any] | None = None) -> ReplayEvent:
        if kind not in EVENT_KINDS:
            raise ValueError(f"unknown event kind: {kind!r}")
        event = ReplayEvent(
            idx=self._idx,
            kind=kind,
            ts_ms=self._clock(),
            payload=dict(payload or {}),
        )
        self._events.append(event)
        self._idx += 1
        return event

    @property
    def events(self) -> tuple[ReplayEvent, ...]:
        return tuple(self._events)

    def build_session(
        self,
        *,
        session_id: str | None = None,
        ended_at: int | None = None,
        receipt_id: str | None = None,
    ) -> ReplaySession:
        end = ended_at if ended_at is not None else self._clock() // 1000
        return ReplaySession(
            session_id=session_id or secrets.token_hex(8),
            agent_name=self._agent_name,
            agent_version=self._agent_version,
            caller=self._caller,
            task_id=self._task_id,
            skill_name=self._skill_name,
            random_seed=self._random_seed,
            input_hash=self._input_hash,
            started_at=self._started_at,
            ended_at=end,
            events=tuple(self._events),
            receipt_id=receipt_id if receipt_id is not None else self._receipt_id,
        )


def seal_replay_session(
    session: ReplaySession,
    *,
    secret: bytes | None = None,
) -> tuple[ReplaySession, str]:
    """Return ``(session, signed_token)``."""
    return session, sign_replay_session(session, secret=secret)


def sign_replay_session(
    session: ReplaySession,
    *,
    secret: bytes | None = None,
) -> str:
    payload = session.model_dump_json(exclude_none=False).encode("utf-8")
    signing_key = _replay_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_replay_session(
    token: str,
    *,
    secret: bytes | None = None,
) -> ReplaySession:
    """Parse + verify ``token``. Raises :class:`ReplayInvalid` on any failure.

    Checks signature and structural shape. Also validates that ``events`` is
    monotonically indexed and uses only known event kinds — replay sessions
    are forever; an unindexed log is a forged log.
    """
    if not token or "." not in token:
        raise ReplayInvalid("malformed replay 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 ReplayInvalid(f"replay 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 ReplayInvalid("replay signature mismatch") from exc
    else:
        expected = hmac.new(secret or _platform_secret(), payload, hashlib.sha256).digest()
        if not hmac.compare_digest(expected, sig):
            raise ReplayInvalid("replay signature mismatch")

    try:
        data = json.loads(payload)
        session = ReplaySession.model_validate(data)
    except Exception as exc:  # noqa: BLE001
        raise ReplayInvalid(f"replay payload invalid: {exc}") from exc

    for i, event in enumerate(session.events):
        if event.idx != i:
            raise ReplayInvalid(
                f"replay event {i} has non-monotonic idx {event.idx}"
            )
        if event.kind not in EVENT_KINDS:
            raise ReplayInvalid(f"replay event {i} has unknown kind {event.kind!r}")

    return session


def iter_events(session: ReplaySession) -> Iterator[ReplayEvent]:
    """Yield events in their recorded order. Drives the UI scrub timeline."""
    yield from session.events


def filter_events(
    session: ReplaySession, *, kinds: Iterable[str] | None = None
) -> Iterator[ReplayEvent]:
    """Yield events in order, optionally filtered to ``kinds``."""
    allowed = set(kinds) if kinds else None
    for event in session.events:
        if allowed is None or event.kind in allowed:
            yield event


__all__ = [
    "EVENT_KINDS",
    "EventRecorder",
    "ReplayEvent",
    "ReplayInvalid",
    "ReplaySession",
    "filter_events",
    "iter_events",
    "seal_replay_session",
    "sign_replay_session",
    "verify_replay_session",
]
