from __future__ import annotations

import time
from typing import Any

from agents.lifecycle import RunHooksBase


class ArmorerRunHooks(RunHooksBase[Any, Any]):
    def __init__(self, *, verification_ledger: list[dict[str, Any]] | None = None):
        self._verification_ledger = verification_ledger if verification_ledger is not None else []

    async def on_tool_start(self, context, agent, tool) -> None:
        self._verification_ledger.append(
            {
                "event": "tool_start",
                "agent": getattr(agent, "name", ""),
                "tool": getattr(tool, "name", ""),
                "timestamp": time.time(),
            }
        )

    async def on_tool_end(self, context, agent, tool, result: str) -> None:
        self._verification_ledger.append(
            {
                "event": "tool_end",
                "agent": getattr(agent, "name", ""),
                "tool": getattr(tool, "name", ""),
                "timestamp": time.time(),
                "result_preview": str(result or "")[:400],
            }
        )

    async def on_agent_end(self, context, agent, output: Any) -> None:
        self._verification_ledger.append(
            {
                "event": "agent_end",
                "agent": getattr(agent, "name", ""),
                "timestamp": time.time(),
                "output_preview": str(output or "")[:400],
            }
        )
