"""
测试多轮对话的 Prompt Caching
"""

import asyncio
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from dotenv import load_dotenv
load_dotenv()

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 = 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 确保 >1024 tokens
    system_prompt = """你是一个专业的 AI 助手，专注于帮助用户解决技术问题。

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

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

## 技术栈
- 编程语言：Python, JavaScript, TypeScript, Go, Rust, Java
- 前端框架：React, Vue, Angular, Svelte
- 后端框架：Node.js, Django, Flask, FastAPI, Spring Boot
- 数据库：PostgreSQL, MongoDB, Redis, MySQL, Elasticsearch
- 云平台：AWS, GCP, Azure
- DevOps：Docker, Kubernetes, CI/CD, Terraform
- 机器学习：TensorFlow, PyTorch, scikit-learn

## 响应格式
- 提供清晰的步骤说明
- 包含代码示例
- 解释关键概念
- 指出潜在问题
- 给出最佳实践建议

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

    messages = [
        {"role": "user", "content": "请用一句话介绍 Python"}
    ]

    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=5,  # 多轮对话
            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 messages: {item.total_messages}")
                print(f"  Total tokens: {item.total_tokens}")
                print(f"  Total cache creation: {item.total_cache_creation_tokens}")
                print(f"  Total cache read: {item.total_cache_read_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", "")
                    tool_calls = content.get("tool_calls")
                    if text and not tool_calls:
                        preview = text[:80] + "..." if len(text) > 80 else text
                        print(f"  Response: {preview}")
                    if tool_calls:
                        print(f"  Tool calls: {len(tool_calls)}")

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

    if trace_id:
        print("分析：")
        print("- 第 1 次调用：应该有 cache_creation_tokens > 0（创建缓存）")
        print("- 第 2+ 次调用：应该有 cache_read_tokens > 0（命中缓存）")
        print(f"\nTrace ID: {trace_id}")

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