import asyncio
import time
from datetime import datetime, timezone
from typing import Any, Dict, Optional

from .base import AgentContext, AgentResult, AgentStatus, BaseAgent
from src.config import settings
from src.telemetry import collector


class AgentExecutor:
    def __init__(self):
        self._semaphore = asyncio.Semaphore(settings.max_concurrent_agents)
        self._active_runs: Dict[str, asyncio.Task] = {}

    async def run(
        self,
        agent: BaseAgent,
        context: Optional[AgentContext] = None,
        timeout: Optional[int] = None,
        **kwargs,
    ) -> AgentResult:
        if context is None:
            context = AgentContext(agent_id=agent.agent_id)

        timeout = timeout or settings.agent_timeout
        started_at = datetime.now(timezone.utc)
        t0 = time.monotonic()

        result = AgentResult(
            run_id=context.run_id,
            agent_id=agent.agent_id,
            status=AgentStatus.PENDING,
            started_at=started_at,
        )

        async with self._semaphore:
            result.status = AgentStatus.RUNNING
            collector.agent_start(agent.agent_id, context.run_id, agent.agent_type)
            try:
                await agent.on_start(context)
                output = await asyncio.wait_for(
                    agent.execute(context, **kwargs),
                    timeout=timeout,
                )
                result.output = output
                result.status = AgentStatus.COMPLETED
            except asyncio.TimeoutError:
                result.status = AgentStatus.TIMEOUT
                result.error = f"Agent timed out after {timeout}s"
                collector.agent_failed(agent.agent_id, context.run_id, result.error)
                await agent.on_error(context, TimeoutError(result.error))
            except asyncio.CancelledError:
                result.status = AgentStatus.CANCELLED
                result.error = "Agent execution was cancelled"
                raise
            except Exception as exc:
                result.status = AgentStatus.FAILED
                result.error = str(exc)
                collector.agent_failed(agent.agent_id, context.run_id, result.error)
                await agent.on_error(context, exc)
            finally:
                result.completed_at = datetime.now(timezone.utc)
                result.duration_ms = (time.monotonic() - t0) * 1000
                if result.status == AgentStatus.COMPLETED:
                    collector.agent_complete(agent.agent_id, context.run_id, result.duration_ms)
                    await agent.on_complete(context, result)

        return result

    async def run_with_retry(
        self,
        agent: BaseAgent,
        context: Optional[AgentContext] = None,
        retries: Optional[int] = None,
        **kwargs,
    ) -> AgentResult:
        max_retries = retries if retries is not None else settings.retry_limit
        last_result: Optional[AgentResult] = None

        for attempt in range(max_retries + 1):
            ctx = context or AgentContext(agent_id=agent.agent_id)
            if attempt > 0:
                ctx.metadata["retry_attempt"] = attempt
            last_result = await self.run(agent, ctx, **kwargs)
            if last_result.status == AgentStatus.COMPLETED:
                return last_result
            if last_result.status == AgentStatus.CANCELLED:
                break

        return 
    

    async def cancel(self, run_id: str) -> bool:
        task = self._active_runs.get(run_id)
        if task and not task.done():
            task.cancel()
            return True
        return False


executor = AgentExecutor()