# -*- coding: utf-8 -*-
"""
端到端风险测试脚本 - 直接调用运行中的服务 API

测试场景：
1. 上传同名文件覆盖 → 版本记录是否正确（P0 修复验证）
2. 废止文档 → 切片状态是否正确标记（Part 4 修复验证）
3. 废止后检索 → 废止文档不应出现在结果中（过滤逻辑验证）
4. 恢复文档 → 切片状态应恢复为 active
5. SQLite/ChromaDB 状态一致性
6. 边缘操作：重复废止、恢复非废止文档等
"""

import requests
import json
import time
import os
import sys
import tempfile

API = "http://localhost:5001"
KB = "test1"  # 使用隔离的测试知识库

passed = 0
failed = 0


def check(name, condition, detail=""):
    global passed, failed
    if condition:
        print(f"  [PASS] {name}")
        passed += 1
    else:
        print(f"  [FAIL] {name}  {detail}")
        failed += 1


def upload_file(collection, filename, content):
    """上传文件到指定知识库"""
    # 创建临时文件
    tmpdir = tempfile.mkdtemp()
    filepath = os.path.join(tmpdir, filename)
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(content)

    with open(filepath, 'rb') as f:
        resp = requests.post(
            f"{API}/documents/upload",
            files={"file": (filename, f, "text/plain")},
            data={"collection": collection}
        )
    # 清理临时文件
    os.remove(filepath)
    os.rmdir(tmpdir)
    return resp.json()


def get_doc_chunks(collection, filename):
    """获取文档的切片信息"""
    resp = requests.get(
        f"{API}/documents/{collection}/{filename}/chunks"
    )
    if resp.status_code == 200:
        return resp.json()
    return {"error": resp.text, "chunks": []}


def get_doc_status(collection, filename):
    """获取文档状态"""
    resp = requests.get(
        f"{API}/documents/{collection}/{filename}/status"
    )
    if resp.status_code == 200:
        return resp.json()
    return {"error": resp.text}


def get_version_history(collection, filename):
    """获取版本历史"""
    resp = requests.get(
        f"{API}/collections/{collection}/documents/{filename}/versions"
    )
    if resp.status_code == 200:
        return resp.json()
    return {"error": resp.text, "versions": []}


def deprecate_doc(collection, filename, reason="测试废止"):
    """废止文档"""
    resp = requests.post(
        f"{API}/collections/{collection}/documents/{filename}/deprecate",
        json={"reason": reason}
    )
    return resp.json()


def restore_doc(collection, filename):
    """恢复文档"""
    resp = requests.post(
        f"{API}/collections/{collection}/documents/{filename}/restore"
    )
    return resp.json()


def delete_doc(collection, filename):
    """删除文档"""
    resp = requests.delete(
        f"{API}/documents/{collection}/{filename}"
    )
    return resp.json()


def rag_query(collection, query):
    """发送 RAG 查询"""
    resp = requests.post(
        f"{API}/rag",
        json={
            "question": query,
            "collection": collection,
            "stream": False
        },
        timeout=30
    )
    if resp.status_code == 200:
        return resp.json()
    return {"error": resp.text}


# ======================================================================
# 清理：确保 test1 知识库是干净的
# ======================================================================
print("\n=== 准备：清理 test1 知识库 ===")
list_resp = requests.get(f"{API}/documents/list?collection={KB}").json()
for doc in list_resp.get("documents", []):
    src = doc.get("source", "")
    if src:
        delete_doc(KB, src)
        print(f"  清理旧文档: {src}")
time.sleep(1)
print("  知识库已清理")


# ======================================================================
# 风险场景 1：上传同名文件覆盖 → 版本记录
# ======================================================================
print("\n=== 风险场景 1：上传覆盖 → SQLite 版本记录 ===")

# 1a. 上传 v1
v1_content = "这是版本测试文档的第一版内容。\n包含一些独特的v1信息用于后续检索验证。\nAlpha Bravo Charlie Delta."
r1 = upload_file(KB, "version_test.txt", v1_content)
check("v1 上传成功", r1.get("success") is True or "chunks" in str(r1),
      f"resp: {json.dumps(r1, ensure_ascii=False)[:200]}")
time.sleep(1)

# 1b. 检查 v1 的 chunks
chunks_v1 = get_doc_chunks(KB, "version_test.txt")
v1_chunks_list = chunks_v1.get("chunks", [])
check("v1 有切片",
      len(v1_chunks_list) > 0,
      f"chunks count: {len(v1_chunks_list)}")
if v1_chunks_list:
    first_meta = v1_chunks_list[0].get("metadata", {})
    check("v1 切片 status=active",
          first_meta.get("status") == "active",
          f"status: {first_meta.get('status')}")

# 1c. 上传 v2（同名覆盖）
v2_content = "这是版本测试文档的第二版内容。\nv2版本包含了全新的信息。\nEcho Foxtrot Golf Hotel."
r2 = upload_file(KB, "version_test.txt", v2_content)
check("v2 覆盖上传成功",
      r2.get("success") is True or "chunks" in str(r2),
      f"resp: {json.dumps(r2, ensure_ascii=False)[:200]}")
check("v2 标记为 replaced",
      r2.get("data", {}).get("file", {}).get("replaced") is True,
      f"replaced: {r2.get('data', {}).get('file', {}).get('replaced')}")
time.sleep(1)

# 1d. 检查 v2 的 chunks（应该是新版本内容）
chunks_v2 = get_doc_chunks(KB, "version_test.txt")
v2_chunks_list = chunks_v2.get("chunks", [])
check("v2 有切片",
      len(v2_chunks_list) > 0,
      f"chunks count: {len(v2_chunks_list)}")
if v2_chunks_list:
    first_meta = v2_chunks_list[0].get("metadata", {})
    check("v2 切片 status=active",
          first_meta.get("status") == "active",
          f"status: {first_meta.get('status')}")

# 1e. 检查版本历史（SQLite）
versions = get_version_history(KB, "version_test.txt")
ver_list = versions.get("versions", [])
check("版本历史有记录",
      len(ver_list) > 0,
      f"versions: {json.dumps(ver_list, ensure_ascii=False)[:300]}")

if len(ver_list) >= 2:
    # 应该有 v1(superseded) 和 v2(active)
    statuses = [v.get("status") for v in ver_list]
    check("版本历史包含 superseded 和 active 状态",
          "superseded" in str(statuses) and "active" in str(statuses),
          f"statuses: {statuses}")
elif len(ver_list) == 1:
    check("至少有一条 active 版本记录",
          ver_list[0].get("status") in ("active", "superseded"),
          f"version: {ver_list[0]}")


# ======================================================================
# 风险场景 2：废止文档 → 切片状态标记
# ======================================================================
print("\n=== 风险场景 2：废止文档 → ChromaDB + SQLite 状态同步 ===")

# 2a. 上传一个专门用于废止测试的文档
deprecate_content = "这份文档将被废止。\n包含独特的废止测试关键词 XYZ123ABC。\nIndigo Juliet Kilo Lima."
r3 = upload_file(KB, "deprecate_test.txt", deprecate_content)
check("废止测试文档上传成功",
      r3.get("success") is True or "chunks" in str(r3),
      f"resp: {json.dumps(r3, ensure_ascii=False)[:200]}")
time.sleep(1)

# 2b. 确认上传后状态为 active
chunks_before = get_doc_chunks(KB, "deprecate_test.txt")
if chunks_before.get("chunks"):
    check("上传后切片状态为 active",
          all(c.get("metadata", {}).get("status") == "active"
              for c in chunks_before["chunks"]),
          f"statuses: {[c.get('metadata', {}).get('status') for c in chunks_before['chunks']]}")

# 2c. 执行废止
dep_result = deprecate_doc(KB, "deprecate_test.txt", reason="测试废止操作")
check("废止操作返回 success",
      dep_result.get("success") is True,
      f"resp: {json.dumps(dep_result, ensure_ascii=False)[:200]}")
check("废止标记了切片",
      dep_result.get("deprecated_chunks", 0) > 0,
      f"deprecated_chunks: {dep_result.get('deprecated_chunks')}")

# 2d. 验证 ChromaDB 中的切片状态
time.sleep(0.5)
chunks_after_dep = get_doc_chunks(KB, "deprecate_test.txt")
if chunks_after_dep.get("chunks"):
    dep_statuses = [c.get("metadata", {}).get("status", "") for c in chunks_after_dep["chunks"]]
    check("ChromaDB 切片状态已改为 deprecated",
          all(s == "deprecated" for s in dep_statuses),
          f"statuses: {dep_statuses}")

# 2e. 验证 SQLite 版本记录也同步了
dep_versions = get_version_history(KB, "deprecate_test.txt")
dep_ver_list = dep_versions.get("versions", [])
if dep_ver_list:
    has_deprecated = any(
        v.get("status") in ("deprecated",) or
        str(v.get("status", "")).lower() == "deprecated"
        for v in dep_ver_list
    )
    check("SQLite 版本记录中有 deprecated 状态",
          has_deprecated,
          f"versions: {json.dumps(dep_ver_list, ensure_ascii=False)[:300]}")
else:
    check("SQLite 版本记录存在", False, "版本历史为空")


# ======================================================================
# 风险场景 3：废止后检索 → 不应出现废止文档内容
# ======================================================================
print("\n=== 风险场景 3：废止后检索过滤 ===")

# 用废止文档中的独特关键词检索
search_result = rag_query(KB, "XYZ123ABC 废止测试关键词")
if "error" not in search_result:
    answer = search_result.get("answer", "")
    sources = search_result.get("sources", [])
    source_files = [s.get("source", s.get("file", "")) for s in sources] if sources else []
    check("检索结果不包含废止文档",
          "deprecate_test.txt" not in source_files,
          f"sources: {source_files}")
    # 也检查 citations
    citations = search_result.get("citations", [])
    cite_sources = [c.get("source", "") for c in citations] if citations else []
    check("citations 不包含废止文档",
          "deprecate_test.txt" not in cite_sources,
          f"cite_sources: {cite_sources}")
else:
    print(f"  [SKIP] RAG 查询失败（可能 LLM 不可用）: {search_result.get('error', '')[:100]}")
    # 备选方案：直接检查 chunks 的 status 字段
    dep_check = get_doc_chunks(KB, "deprecate_test.txt")
    if dep_check.get("chunks"):
        all_dep = all(
            c.get("status") == "deprecated" or c.get("metadata", {}).get("status") == "deprecated"
            for c in dep_check["chunks"]
        )
        check("（备选）废止文档所有切片 status=deprecated",
              all_dep,
              f"statuses: {[c.get('status', c.get('metadata', {}).get('status')) for c in dep_check['chunks']]}")


# ======================================================================
# 风险场景 4：恢复文档 → 切片状态恢复
# ======================================================================
print("\n=== 风险场景 4：恢复已废止文档 ===")

restore_result = restore_doc(KB, "deprecate_test.txt")
check("恢复操作返回 success",
      restore_result.get("success") is True,
      f"resp: {json.dumps(restore_result, ensure_ascii=False)[:200]}")
check("恢复了切片",
      restore_result.get("restored_chunks", 0) > 0,
      f"restored_chunks: {restore_result.get('restored_chunks')}")

# 验证 ChromaDB 切片恢复
time.sleep(0.5)
chunks_after_restore = get_doc_chunks(KB, "deprecate_test.txt")
if chunks_after_restore.get("chunks"):
    restored_statuses = [c.get("metadata", {}).get("status", "") for c in chunks_after_restore["chunks"]]
    check("ChromaDB 切片状态恢复为 active",
          all(s == "active" for s in restored_statuses),
          f"statuses: {restored_statuses}")

# 验证 SQLite 版本记录
rest_versions = get_version_history(KB, "deprecate_test.txt")
rest_ver_list = rest_versions.get("versions", [])
if rest_ver_list:
    has_active = any(
        str(v.get("status", "")).lower() == "active"
        for v in rest_ver_list
    )
    check("SQLite 版本记录恢复为 active",
          has_active,
          f"versions: {json.dumps(rest_ver_list, ensure_ascii=False)[:300]}")


# ======================================================================
# 风险场景 5：SQLite 与 ChromaDB 状态一致性
# ======================================================================
print("\n=== 风险场景 5：SQLite/ChromaDB 状态一致性 ===")

# 对 version_test.txt 做一致性检查
ver_chunks = get_doc_chunks(KB, "version_test.txt")
ver_versions = get_version_history(KB, "version_test.txt")

if ver_chunks.get("chunks") and ver_versions.get("versions"):
    # ChromaDB 中所有切片应该是 active（只有当前版本在 ChromaDB 中）
    chroma_statuses = set(
        c.get("metadata", {}).get("status", "active")
        for c in ver_chunks["chunks"]
    )
    check("ChromaDB 中 version_test 切片全为 active",
          chroma_statuses == {"active"} or chroma_statuses == set(),
          f"chroma_statuses: {chroma_statuses}")

    # SQLite 中应该至少有一条 active 记录
    sqlite_statuses = [v.get("status") for v in ver_versions["versions"]]
    check("SQLite 中有 active 版本记录",
          "active" in sqlite_statuses,
          f"sqlite_statuses: {sqlite_statuses}")
else:
    check("能获取到切片和版本信息", False,
          f"chunks: {bool(ver_chunks.get('chunks'))}, versions: {bool(ver_versions.get('versions'))}")


# ======================================================================
# 风险场景 6：边缘操作（不应崩溃）
# ======================================================================
print("\n=== 风险场景 6：边缘操作容错 ===")

# 6a. 废止不存在的文档
dep_nonexist = deprecate_doc(KB, "nonexistent_file_xyz.txt")
check("废止不存在文档不崩溃",
      "success" in dep_nonexist or "error" in dep_nonexist,
      f"resp: {json.dumps(dep_nonexist, ensure_ascii=False)[:200]}")

# 6b. 恢复非废止状态的文档
# deprecate_test.txt 已经被恢复了，再恢复一次应该报错但不崩溃
restore_again = restore_doc(KB, "deprecate_test.txt")
check("重复恢复不崩溃",
      "success" in restore_again or "error" in restore_again,
      f"resp: {json.dumps(restore_again, ensure_ascii=False)[:200]}")

# 6c. 连续两次废止同一文档
dep1 = deprecate_doc(KB, "deprecate_test.txt", reason="第一次废止")
check("第一次废止成功",
      dep1.get("success") is True,
      f"resp: {json.dumps(dep1, ensure_ascii=False)[:200]}")
time.sleep(0.5)
dep2 = deprecate_doc(KB, "deprecate_test.txt", reason="第二次废止")
check("第二次废止不崩溃（已废止状态）",
      "success" in dep2 or "error" in dep2,
      f"resp: {json.dumps(dep2, ensure_ascii=False)[:200]}")

# 恢复（为后续清理准备）
restore_doc(KB, "deprecate_test.txt")
time.sleep(0.5)

# 6d. 上传空文件
empty_result = upload_file(KB, "empty_file.txt", "")
check("上传空文件不崩溃",
      "success" in empty_result or "error" in empty_result,
      f"resp: {json.dumps(empty_result, ensure_ascii=False)[:200]}")


# ======================================================================
# 清理：删除测试文档
# ======================================================================
print("\n=== 清理测试数据 ===")
for fname in ["version_test.txt", "deprecate_test.txt", "empty_file.txt"]:
    r = delete_doc(KB, fname)
    print(f"  删除 {fname}: {r.get('success', r.get('deleted', r.get('error', '?')))}")


# ======================================================================
# 汇总
# ======================================================================
print(f"\n{'='*60}")
print(f"风险测试完成: {passed} 通过, {failed} 失败, 共 {passed + failed} 条")
print(f"{'='*60}")

if failed > 0:
    sys.exit(1)