"""Chain verification helpers for AMCS stores."""

from __future__ import annotations

from dataclasses import dataclass

from amcs.hashing import event_hash
from amcs.merkle import merkle_root_hex
from amcs.store.base import EventStore


@dataclass(frozen=True)
class VerificationResult:
    ok: bool
    checked_events: int
    failed_sequence: int | None = None
    error: str | None = None
    expected_root: str | None = None
    actual_root: str | None = None


def _event_for_hash(row: dict) -> dict:
    """Backward-compatible hash input for legacy rows missing sequence in JSON."""
    event = dict(row["event"])
    if "sequence" not in event:
        event["sequence"] = row["sequence"]
    return event


def verify_chain(
    store: EventStore,
    agent_id: str,
    from_seq: int = 1,
    to_seq: int | None = None,
) -> VerificationResult:
    """Verify event hashes and memory root consistency for an agent chain."""
    resolved_to_seq = to_seq
    if resolved_to_seq is None:
        resolved_to_seq = store.get_latest_sequence(agent_id)

    events = store.get_events(agent_id, from_seq=from_seq, to_seq=resolved_to_seq)

    for row in events:
        recomputed = event_hash(_event_for_hash(row))
        if recomputed != row["event_hash"]:
            return VerificationResult(
                ok=False,
                checked_events=row["sequence"] - from_seq,
                failed_sequence=row["sequence"],
                error="event hash mismatch",
            )

    all_events_for_root = store.get_events(agent_id, from_seq=1, to_seq=resolved_to_seq)
    if not all_events_for_root:
        actual_root = store.get_memory_root(agent_id, seq=resolved_to_seq)
        if actual_root is not None:
            return VerificationResult(
                ok=False,
                checked_events=0,
                error="unexpected stored root for empty chain",
                expected_root=None,
                actual_root=actual_root,
            )
        return VerificationResult(ok=True, checked_events=0)

    expected_root = merkle_root_hex(
        [event_hash(_event_for_hash(row)) for row in all_events_for_root]
    )
    actual_root = store.get_memory_root(agent_id, seq=resolved_to_seq)
    if expected_root != actual_root:
        return VerificationResult(
            ok=False,
            checked_events=len(events),
            error="memory root mismatch",
            expected_root=expected_root,
            actual_root=actual_root,
        )

    return VerificationResult(
        ok=True,
        checked_events=len(events),
        expected_root=expected_root,
        actual_root=actual_root,
    )
