"""External checkpoint helpers for tamper-evidence beyond local storage."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any

from amcs.canonical_json import canonical_bytes
from amcs.crypto.ed25519 import b64decode, b64encode, sign, verify
from amcs.store.base import EventStore


@dataclass(frozen=True)
class RootCheckpoint:
    agent_id: str
    sequence: int
    memory_root: str
    checkpointed_at: str
    signer_kid: str
    signature_b64: str

    def as_dict(self) -> dict[str, Any]:
        return {
            "agent_id": self.agent_id,
            "sequence": self.sequence,
            "memory_root": self.memory_root,
            "checkpointed_at": self.checkpointed_at,
            "signer_kid": self.signer_kid,
            "signature_b64": self.signature_b64,
        }


def _utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _payload_for_signing(checkpoint: dict[str, Any]) -> dict[str, Any]:
    required = ("agent_id", "sequence", "memory_root", "checkpointed_at", "signer_kid")
    missing = [field for field in required if field not in checkpoint]
    if missing:
        raise ValueError(f"missing checkpoint fields: {', '.join(missing)}")

    payload = {
        "agent_id": checkpoint["agent_id"],
        "sequence": checkpoint["sequence"],
        "memory_root": checkpoint["memory_root"],
        "checkpointed_at": checkpoint["checkpointed_at"],
        "signer_kid": checkpoint["signer_kid"],
    }
    return payload


def build_root_checkpoint(
    *,
    store: EventStore,
    agent_id: str,
    signer_private_key_b64: str,
    signer_kid: str,
    sequence: int | None = None,
    checkpointed_at: str | None = None,
) -> RootCheckpoint:
    resolved_sequence = sequence if sequence is not None else store.get_latest_sequence(agent_id)
    if resolved_sequence is None:
        raise ValueError("cannot checkpoint empty agent log")

    memory_root = store.get_memory_root(agent_id, seq=resolved_sequence)
    if not memory_root:
        raise ValueError("memory root not found for requested sequence")

    payload = {
        "agent_id": agent_id,
        "sequence": int(resolved_sequence),
        "memory_root": str(memory_root),
        "checkpointed_at": checkpointed_at or _utc_now_iso(),
        "signer_kid": signer_kid,
    }
    message = canonical_bytes(_payload_for_signing(payload))
    signature_b64 = b64encode(sign(b64decode(signer_private_key_b64), message))
    return RootCheckpoint(signature_b64=signature_b64, **payload)


def verify_root_checkpoint(*, checkpoint: dict[str, Any], signer_public_key_b64: str) -> bool:
    signature_b64 = checkpoint.get("signature_b64")
    if not isinstance(signature_b64, str):
        return False
    try:
        message = canonical_bytes(_payload_for_signing(checkpoint))
        signature = b64decode(signature_b64)
        public_key = b64decode(signer_public_key_b64)
    except Exception:
        return False
    return verify(public_key, message, signature)
