"""Agent执行指标钩子"""
import time
from contextlib import asynccontextmanager
from typing import Optional

from app.middleware.prometheus_metrics import (
    AGENT_EXECUTIONS_TOTAL,
    AGENT_EXECUTION_DURATION_SECONDS,
    AGENT_RUNNING_TASKS,
)


@asynccontextmanager
async def agent_execution_context(agent_name: str):
    """Agent执行上下文管理器 - 自动记录指标"""
    # 增加运行任务计数
    AGENT_RUNNING_TASKS.labels(agent_name=agent_name).inc()

    start_time = time.perf_counter()
    status = "success"

    try:
        yield
    except Exception as e:
        status = "failure"
        raise
    finally:
        # 记录执行时间和状态
        duration = time.perf_counter() - start_time

        AGENT_EXECUTIONS_TOTAL.labels(
            agent_name=agent_name,
            status=status
        ).inc()

        AGENT_EXECUTION_DURATION_SECONDS.labels(
            agent_name=agent_name
        ).observe(duration)

        # 减少运行任务计数
        AGENT_RUNNING_TASKS.labels(agent_name=agent_name).dec()


def record_agent_execution(
    agent_name: str,
    status: str,
    duration: float
):
    """手动记录Agent执行指标"""
    AGENT_EXECUTIONS_TOTAL.labels(
        agent_name=agent_name,
        status=status
    ).inc()

    AGENT_EXECUTION_DURATION_SECONDS.labels(
        agent_name=agent_name
    ).observe(duration)