# -*- coding: utf-8 -*-
"""
日志审计中间件

功能：
- 监听模型调用，在控制台打印审计日志
- 支持通过环境变量开关控制
- 打印内容：已消耗 token 数、本次调用传入的上下文

使用方法：
1. 在 .env 中设置 AUDIT_LOG_ENABLED=true 启用
2. 在 create_deep_agent 的 middleware 参数中传入
"""

from typing import Any, Dict, Optional
from langgraph.utils.runnable import RunnableCallable


def create_audit_middleware() -> RunnableCallable:
    """
    创建审计中间件（符合 DeepAgents 协议）

    Returns:
        RunnableCallable: 包装后的可调用对象
    """
    # 从环境变量读取配置
    from dotenv import dotenv_values
    from pathlib import Path

    _env_file = Path(__file__).parent.parent / ".env"
    _cfg = dotenv_values(str(_env_file))

    enabled = _cfg.get("AUDIT_LOG_ENABLED", "false").lower() == "true"
    total_tokens = [0]  # 使用列表以在闭包中修改

    async def audit_wrapper(input: Any, config: Optional[Dict] = None, **kwargs: Any) -> Any:
        """审计包装器"""
        if not enabled:
            # 如果未启用，直接调用下一个步骤
            from langgraph.utils.runnable import Runnable
            next_step = kwargs.get("_next") or kwargs.get("next_step")
            if next_step:
                return await next_step.ainvoke(input, config=config, **kwargs)
            return input

        # 提取输入消息
        messages = _extract_messages(input)

        # 打印上下文信息
        _log_context(messages)

        # 获取下一个步骤
        next_step = kwargs.pop("_next", None)
        if next_step is None:
            from langgraph.utils.runnable import Runnable
            for key in kwargs:
                if isinstance(kwargs[key], Runnable):
                    next_step = kwargs.pop(key)
                    break

        if next_step is None:
            return input

        # 调用下一个步骤
        result = await next_step.ainvoke(input, config=config, **kwargs)

        # 尝试提取并打印 token 使用情况
        _log_tokens(result, total_tokens)

        return result

    return RunnableCallable(ainvoke=audit_wrapper)


def _extract_messages(input: Any) -> list:
    """从输入中提取消息列表"""
    if isinstance(input, dict):
        messages = input.get("messages", [])
        if isinstance(messages, list):
            return messages
    elif isinstance(input, list):
        return input
    return []


def _log_context(messages: list):
    """打印上下文信息"""
    print("\n" + "=" * 60)
    print("【审计日志】模型调用上下文")
    print("=" * 60)

    for i, msg in enumerate(messages):
        role = getattr(msg, "role", "unknown")
        content = getattr(msg, "content", str(msg))
        # 限制内容长度
        if len(str(content)) > 500:
            content = str(content)[:500] + "...(截断)"
        print(f"\n[消息 {i+1}] 角色: {role}")
        print(f"内容: {content}")

    print("\n" + "-" * 60)


def _log_tokens(result: Any, total_tokens_ref: list):
    """打印 token 使用情况"""
    # 尝试从结果中提取 token 信息
    token_info = {}
    total_tokens = total_tokens_ref[0]

    # 尝试从 AIMessage 响应中提取
    if hasattr(result, "response_metadata"):
        metadata = result.response_metadata
        if hasattr(metadata, "get"):
            usage = metadata.get("usage", {})
            if isinstance(usage, dict):
                token_info["prompt_tokens"] = usage.get("prompt_tokens", 0)
                token_info["completion_tokens"] = usage.get("completion_tokens", 0)
                token_info["total_tokens"] = usage.get("total_tokens", 0)

    # 更新累计消耗
    if token_info.get("total_tokens"):
        total_tokens_ref[0] = total_tokens + token_info["total_tokens"]

    print(f"【审计日志】Token 使用情况:")
    print(f"  - 本次 Prompt Tokens: {token_info.get('prompt_tokens', 'N/A')}")
    print(f"  - 本次 Completion Tokens: {token_info.get('completion_tokens', 'N/A')}")
    print(f"  - 本次 Total Tokens: {token_info.get('total_tokens', 'N/A')}")
    print(f"  - 累计消耗 Tokens: {total_tokens_ref[0]}")
    print("=" * 60 + "\n")