# -*- coding: utf-8 -*-
"""对话分支 API 端到端测试"""
import sys, os, io, json, time, requests

sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")

BASE = "http://localhost:8038/api/v1"

def login():
    r = requests.post(f"{BASE}/auth/login", data={"username": "admin", "password": "123456"})
    r.raise_for_status()
    return r.json()["access_token"]

def test_branching():
    token = login()
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    print("=== 对话分支 API 测试 ===\n")

    # Step 1: Create a conversation
    print("1. 创建对话…")
    r1 = requests.post(f"{BASE}/agent-chat/bare", headers=headers,
                       json={"message": "你好，我叫小明，是一名Python程序员"}, timeout=120)
    r1.raise_for_status()
    d1 = r1.json()
    sid = d1["session_id"]
    print(f"   会话 ID: {sid}")
    print(f"   回复: {d1['content'][:100]}…\n")

    # Step 2: Create a branch
    print("2. 创建分支…")
    r2 = requests.post(f"{BASE}/agent-chat/branches", headers=headers,
                       json={"session_id": sid, "title": "测试分支"}, timeout=30)
    r2.raise_for_status()
    d2 = r2.json()
    bid = d2["id"]
    print(f"   分支 ID: {bid[:16]}…")
    print(f"   标题: {d2['title']}")
    print(f"   父会话: {d2['parent_session_id']}")
    print(f"   分支会话: {d2['branch_session_id']}")
    print(f"   消息数: {d2['message_count']}")
    print(f"   首条消息: {d2.get('first_user_message', 'N/A')[:80]}\n")

    # Step 3: List branches
    print("3. 列出分支…")
    r3 = requests.get(f"{BASE}/agent-chat/branches", headers=headers, timeout=10)
    r3.raise_for_status()
    d3 = r3.json()
    print(f"   共 {d3['total']} 个分支:")
    for b in d3["branches"]:
        print(f"   - {b['id'][:16]}… {b['title']} ({b['message_count']} 条消息)")
    print()

    # Step 4: Get branch details
    print("4. 获取分支详情…")
    r4 = requests.get(f"{BASE}/agent-chat/branches/{bid}", headers=headers, timeout=10)
    r4.raise_for_status()
    d4 = r4.json()
    msgs = d4.get("messages", [])
    print(f"   完整消息: {len(msgs)} 条")
    for m in msgs[:5]:
        role = m.get("role", "?")
        content = str(m.get("content", ""))[:80]
        print(f"   [{role}] {content}")

    # Step 5: Resume from branch
    print("\n5. 从分支恢复对话…")
    r5 = requests.post(f"{BASE}/agent-chat/branches/{bid}/resume", headers=headers,
                       json={"message": "我叫小红，不叫小明，请更正你的记忆"}, timeout=120)
    r5.raise_for_status()
    d5 = r5.json()
    print(f"   会话 ID: {d5['session_id']}")
    print(f"   回复: {d5['content'][:200]}…")
    print(f"   迭代: {d5['iterations_used']}, 工具调用: {d5['tool_calls_made']}\n")

    # Step 6: Delete branch
    print("6. 删除分支…")
    r6 = requests.delete(f"{BASE}/agent-chat/branches/{bid}", headers=headers, timeout=10)
    r6.raise_for_status()
    d6 = r6.json()
    print(f"   结果: {d6['message']}\n")

    # Verify deletion
    r7 = requests.get(f"{BASE}/agent-chat/branches", headers=headers, timeout=10)
    d7 = r7.json()
    remaining_ids = [b["id"] for b in d7["branches"]]
    if bid not in remaining_ids:
        print("✓ 分支删除验证通过")
    else:
        print("✗ 分支删除验证失败")

    print("\n=== 全部测试通过 ===")

if __name__ == "__main__":
    test_branching()