"""
测试 Prompt Caching 功能
"""

import asyncio
import os
import sys
from pathlib import Path

# 添加项目根目录到 Python 路径
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from dotenv import load_dotenv
load_dotenv()

import logging
# 开启 DEBUG 日志查看缓存标记
logging.basicConfig(level=logging.DEBUG)

from agent.core.runner import AgentRunner, RunConfig
from agent.trace import FileSystemTraceStore, Trace, Message
from agent.llm import create_openrouter_llm_call

async def main():
    print("=" * 60)
    print("测试 Prompt Caching 功能")
    print("=" * 60)
    print()

    # 路径配置
    base_dir = Path(__file__).parent
    project_root = base_dir.parent.parent
    trace_dir = project_root / ".trace"

    # 创建 Runner
    runner = AgentRunner(
        trace_store=FileSystemTraceStore(base_path=str(trace_dir)),
        llm_call=create_openrouter_llm_call(model="anthropic/claude-sonnet-4.5"),
        debug=True
    )

    # 准备测试消息（足够长的 system prompt）
    system_prompt = """你是一个专业的 AI 助手。

## 核心能力
- 代码分析和生成
- 问题解决和调试
- 技术文档编写
- 架构设计建议

## 工作原则
1. 准确性优先：确保提供的信息和代码是正确的
2. 清晰表达：用简洁明了的语言解释复杂概念
3. 实用导向：提供可直接使用的解决方案
4. 持续学习：根据反馈不断改进

## 技术栈
- Python, JavaScript, TypeScript
- React, Vue, Node.js
- Docker, Kubernetes
- PostgreSQL, MongoDB, Redis
- AWS, GCP, Azure

这是一个足够长的 system prompt，用于测试 Anthropic Prompt Caching 功能。
缓存需要至少 1024 tokens 才能生效，所以我们需要让这个 prompt 足够长。
""" * 3  # 重复 3 次确保足够长

    messages = [
        {"role": "user", "content": "请简单介绍一下 Python 的特点，用 3 句话概括"}
    ]

    print("第一次调用（创建缓存）...")
    print("-" * 60)

    trace_id = None
    iteration = 0

    async for item in runner.run(
        messages=messages,
        config=RunConfig(
            system_prompt=system_prompt,
            model="anthropic/claude-sonnet-4.5",
            temperature=0.3,
            max_iterations=3,
            enable_prompt_caching=True,  # 启用缓存
            name="缓存测试"
        )
    ):
        if isinstance(item, Trace):
            trace_id = item.trace_id
            if item.status == "completed":
                print(f"\n✓ Trace 完成")
                print(f"  Total tokens: {item.total_tokens}")
                print(f"  Total cost: ${item.total_cost:.6f}")

        elif isinstance(item, Message):
            if item.role == "assistant":
                iteration += 1
                print(f"\n[Iteration {iteration}]")
                print(f"  Prompt tokens: {item.prompt_tokens}")
                print(f"  Completion tokens: {item.completion_tokens}")
                print(f"  Cache creation: {item.cache_creation_tokens}")
                print(f"  Cache read: {item.cache_read_tokens}")
                print(f"  Cost: ${item.cost:.6f}")

                content = item.content
                if isinstance(content, dict):
                    text = content.get("text", "")
                    if text:
                        preview = text[:100] + "..." if len(text) > 100 else text
                        print(f"  Response: {preview}")

    print()
    print("=" * 60)
    print("测试完成")
    print("=" * 60)
    print()

    if trace_id:
        print("验证要点：")
        print("1. 第一次调用应该有 cache_creation_tokens > 0")
        print("2. 后续调用应该有 cache_read_tokens > 0")
        print("3. cache_read_tokens 的成本应该是正常 input tokens 的 10%")
        print()
        print(f"Trace ID: {trace_id}")

if __name__ == "__main__":
    asyncio.run(main())