"""
基础 Agent 模块

提供所有 Agent 的基类，内置：
- 防提示词注入检测
- 模型调用耗时日志
"""

import time
import logging

from google.adk.agents import LlmAgent
from google.adk.models.llm_response import LlmResponse
from google.genai import types

from common.logger import after_model_callback
from common.prompt_guard import check_prompt_injection

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


# ============================================================
# 防注入的 before_model_callback
# ============================================================
async def _before_model_callback(callback_context, llm_request):
    """在模型调用前检测提示词注入，并记录开始时间"""
    # 1. 防提示词注入检测：检查最新的用户消息
    contents = getattr(llm_request, "contents", None) or []
    # 找到最后一条 role="user" 的消息
    last_user_content = None
    for content in reversed(contents):
        if getattr(content, "role", None) == "user":
            last_user_content = content
            break

    if last_user_content:
        text = ""
        parts = getattr(last_user_content, "parts", None) or []
        for part in parts:
            part_text = getattr(part, "text", None)
            if part_text:
                text += part_text

        injection = check_prompt_injection(text)
        if injection:
            logger.warning(
                "检测到提示词注入 | agent=%s pattern='%s'",
                callback_context.agent_name,
                injection,
            )
            return LlmResponse(
                content=types.Content(
                    role="model",
                    parts=[types.Part(text="抱歉，我无法处理该请求。")],
                ),
                turn_complete=True,
            )

    # 2. 记录开始时间（供 after_model_callback 使用）
    callback_context.state["_log_model_call_start_time"] = time.perf_counter()
    return None


# ============================================================
# HuiYuBaseAgent 基类
# ============================================================
class HuiYuBaseAgent(LlmAgent):
    """
    所有 Agent 的基类，自动集成：
    - 防提示词注入检测
    - 模型调用耗时日志

    用法：
        agent = HuiYuBaseAgent(
            name="my_agent",
            model=model,
            instruction="...",
        )

    如果需要自定义回调，仍然可以传入 before_model_callback / after_model_callback，
    它们会在防注入检测和日志记录之间执行。
    """

    def __init__(self, **kwargs):
        # 获取用户传入的回调
        user_before = kwargs.pop("before_model_callback", None)
        user_after = kwargs.pop("after_model_callback", None)

        # 构建回调链：防注入 → 用户回调 → 日志
        before_chain = [_before_model_callback]
        if user_before:
            before_chain.append(user_before)

        after_chain = []
        if user_after:
            after_chain.append(user_after)
        after_chain.append(after_model_callback)

        kwargs["before_model_callback"] = before_chain
        kwargs["after_model_callback"] = after_chain

        # 强制模型用中文思考：在 instruction 前附加系统级要求
        instruction = kwargs.get("instruction", "")
        if instruction:
            kwargs["instruction"] = (
                "【系统要求：你的所有思考过程、分析推理、内心独白必须使用中文，"
                "禁止在思考中使用英文。你的最终回复也使用中文。】\n\n"
                + instruction
            )

        super().__init__(**kwargs)