"""
AgentRunner: consume tasks from bus and execute via agent.
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any


class AgentRunner:
    def __init__(self, agent, bus, storage, tracer):
        self.agent = agent
        self.bus = bus
        self.storage = storage
        self.tracer = tracer
        self.log = logging.getLogger(agent.name)

    async def start(self):
        self.log.info(f"[{self.agent.name}] Listening for tasks...")
        while True:
            task = await self.bus.consume(self.agent.name)
            payload = task["payload"]
            task_id = task["task_id"]

            self.tracer.log(task_id, self.agent.name, "task_received", payload)
            try:
                result: Any = await self.agent.handle(payload)
                # Store large results in artifact storage
                body = str(result)
                if len(body) > 1_000_000:
                    ref = self.storage.put(result)
                    result = {"artifact_ref": ref}
                self.tracer.log(task_id, self.agent.name, "task_completed", result)
                self.log.info(f"[{self.agent.name}] Task {task_id} done.")
            except Exception as e:
                self.tracer.log(task_id, self.agent.name, "task_error", {"error": str(e)})
                self.log.exception(f"Task {task_id} failed: {e}")
