import time
from backend.logger import setup_logger
from agents import AgentHooks, RunContextWrapper, Agent

logger = setup_logger(__name__)

class LegalAssistantHooks(AgentHooks):
    """Tracks execution time, tool calls and token usage per agent run."""
    
    def __init__(self):
        self.start_time: float = 0
        self.tool_calls: int = 0

    async def on_start(self, ctx: RunContextWrapper, agent: Agent) -> None:
        """Resets counters and logs agent start."""
        self.start_time = time.time()
        self.tool_calls = 0
        logger.info(f"[AGENT {agent.name} STARTED]")

    async def on_tool_start(self, ctx: RunContextWrapper, agent: Agent, tool) -> None:
        """Logs each tool invocation."""
        self.tool_calls += 1
        logger.info(f"[TOOL {tool.name} CALLED]")

    async def on_tool_end(self, ctx: RunContextWrapper, agent: Agent, tool, result) -> None:
        """Marker indicating the end of the tool invocation."""
        logger.info(f"[TOOL {tool.name} ENDED]")

    async def on_end(self, ctx: RunContextWrapper, agent: Agent, output) -> None:
        """Logs elapsed time and total tokens used."""
        elapsed = time.time() - self.start_time
        logger.info(f"[AGENT {agent.name} ENDED] | tools_called={self.tool_calls} | elapsed={elapsed:.2f}s")