"""
统一日志模块

为所有 Agent 提供一致的模型调用日志记录，包括：
- 模型调用耗时
- Token 使用量（prompt / completion）
- 模型版本信息
- Agent 名称

使用方式：
    from common.logger import get_model_callbacks

    agent = Agent(
        name="my_agent",
        model=model,
        instruction="...",
        **get_model_callbacks(),
    )
"""

import time
import logging

logger = logging.getLogger("adk.agent")


# ============================================================
# 模型调用耗时记录的 state key
# ============================================================
_STATE_KEY_START_TIME = "_log_model_call_start_time"


# ============================================================
# before_model_callback：记录模型调用开始时间
# ============================================================
async def before_model_callback(callback_context, llm_request):
    """在模型调用前记录开始时间到 session state"""
    callback_context.state[_STATE_KEY_START_TIME] = time.perf_counter()
    return None


# ============================================================
# after_model_callback：计算耗时并输出日志
# ============================================================
async def after_model_callback(callback_context, llm_response):
    """在模型调用完成后计算耗时，输出结构化日志"""
    # 流式模式下只在最终响应时记录（避免重复打印）
    if getattr(llm_response, "partial", None) and not getattr(
        llm_response, "turn_complete", None
    ):
        return None

    start_time = callback_context.state.get(_STATE_KEY_START_TIME)
    if start_time is None:
        return None
    callback_context.state[_STATE_KEY_START_TIME] = None

    elapsed = time.perf_counter() - start_time
    agent_name = callback_context.agent_name
    model_version = getattr(llm_response, "model_version", None) or "unknown"
    usage = getattr(llm_response, "usage_metadata", None)

    prompt_tokens = getattr(usage, "prompt_token_count", None) if usage else None
    candidates_tokens = (
        getattr(usage, "candidates_token_count", None) if usage else None
    )
    total_tokens = getattr(usage, "total_token_count", None) if usage else None

    # ── 打印模型输出的完整内容（用于排查问题） ──
    content = getattr(llm_response, "content", None)
    if content:
        parts = getattr(content, "parts", None) or []
        for i, part in enumerate(parts):
            text = getattr(part, "text", None)
            if text:
                logger.info(
                    "【模型输出】agent=%s part=%d text=%s",
                    agent_name,
                    i,
                    text[:2000] if len(text) > 2000 else text,
                )
            else:
                # 非文本 part（如工具调用）
                func_call = getattr(part, "function_call", None)
                if func_call:
                    logger.info(
                        "【工具调用】agent=%s function=%s args=%s",
                        agent_name,
                        getattr(func_call, "name", "?"),
                        getattr(func_call, "args", {}),
                    )

    logger.info(
        "模型调用完成 | agent=%s model=%s latency=%.3fs "
        "prompt_tokens=%s completion_tokens=%s total_tokens=%s",
        agent_name,
        model_version,
        elapsed,
        prompt_tokens,
        candidates_tokens,
        total_tokens,
    )

    return None


# ============================================================
# 便捷函数：获取回调字典，用于展开到 Agent 构造参数
# ============================================================
def get_model_callbacks():
    """
    返回 before_model_callback 和 after_model_callback 的字典。

    用法：
        agent = Agent(
            name="my_agent",
            model=model,
            instruction="...",
            **get_model_callbacks(),
        )
    """
    return {
        "before_model_callback": before_model_callback,
        "after_model_callback": after_model_callback,
    }