# -*- coding: utf-8 -*-
"""
Phase 1 + Phase 2 边界风险修复验证测试

测试覆盖：
1. RRF 融合跨库同名文件不吞结果
2. search_multiple 跨库去重正确
3. DocStore 路径含 collection 前缀
4. citation 构建兼容 _collection 和 collection
5. _collection 回退逻辑正确
"""

import sys
import os
import json
import tempfile
import shutil

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

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


# ======================================================================
# 测试 1：RRF 融合 — 跨库同名文件不被吞
# ======================================================================
print("\n=== 测试 1：RRF 融合复合键去重 ===")

# 模拟 engine 的 reciprocal_rank_fusion 逻辑（提取纯函数测试）
def reciprocal_rank_fusion_test(results_list, weights=None, k=60):
    """复刻 engine.py 的 RRF 逻辑（含修复）"""
    if not results_list:
        return {'ids': [[]], 'documents': [[]], 'metadatas': [[]], 'distances': [[]]}
    if weights is None:
        weights = [1.0] * len(results_list)

    doc_scores = {}
    for results, weight in zip(results_list, weights):
        if not results['documents'] or not results['documents'][0]:
            continue
        for rank, (doc_id, doc, meta) in enumerate(zip(
            results['ids'][0], results['documents'][0], results['metadatas'][0]
        )):
            rrf_score = weight / (k + rank + 1)
            coll = meta.get('_collection') or meta.get('collection') or ''
            composite_key = f"{coll}\x00{doc_id}" if coll else doc_id
            if composite_key not in doc_scores:
                doc_scores[composite_key] = {'score': 0.0, 'doc': doc, 'meta': meta, 'coll': coll, 'raw_id': doc_id}
            doc_scores[composite_key]['score'] += rrf_score

    sorted_items = sorted(doc_scores.items(), key=lambda x: x[1]['score'], reverse=True)

    out_ids = []
    for item in sorted_items:
        coll = item[1]['coll']
        raw_id = item[1]['raw_id']
        if coll and not raw_id.startswith(f"{coll}/"):
            out_ids.append(f"{coll}/{raw_id}")
        else:
            out_ids.append(raw_id)

    return {
        'ids': [out_ids],
        'documents': [[item[1]['doc'] for item in sorted_items]],
        'metadatas': [[item[1]['meta'] for item in sorted_items]],
        'distances': [[item[1]['score'] for item in sorted_items]],
    }


# 场景：public_kb 和 dept_1_kb 都有 "规章制度.pdf_0"
result_public = {
    'ids': [['规章制度.pdf_0', '规章制度.pdf_1']],
    'documents': [['public版内容_0', 'public版内容_1']],
    'metadatas': [[
        {'source': '规章制度.pdf', 'collection': 'public_kb', '_collection': 'public_kb', 'chunk_index': 0},
        {'source': '规章制度.pdf', 'collection': 'public_kb', '_collection': 'public_kb', 'chunk_index': 1},
    ]],
    'distances': [[0.9, 0.8]],
}

result_dept = {
    'ids': [['规章制度.pdf_0', '规章制度.pdf_1']],
    'documents': [['dept版内容_0', 'dept版内容_1']],
    'metadatas': [[
        {'source': '规章制度.pdf', 'collection': 'dept_1_kb', '_collection': 'dept_1_kb', 'chunk_index': 0},
        {'source': '规章制度.pdf', 'collection': 'dept_1_kb', '_collection': 'dept_1_kb', 'chunk_index': 1},
    ]],
    'distances': [[0.85, 0.75]],
}

rrf_result = reciprocal_rank_fusion_test([result_public, result_dept])
rrf_ids = rrf_result['ids'][0]

check("跨库同名文件：结果数应为 4（不是 2）", len(rrf_ids) == 4, f"实际: {len(rrf_ids)}")
check("ID 包含 collection 前缀", all('/' in i for i in rrf_ids), f"IDs: {rrf_ids}")
check("public_kb 的结果存在", any('public_kb/' in i for i in rrf_ids))
check("dept_1_kb 的结果存在", any('dept_1_kb/' in i for i in rrf_ids))

# 验证文档内容保留完整（不被覆盖）
rrf_docs = rrf_result['documents'][0]
check("public版内容保留", any('public版' in d for d in rrf_docs))
check("dept版内容保留", any('dept版' in d for d in rrf_docs))

# 场景：同库向量+BM25同名chunk应合并分数
result_vec = {
    'ids': [['规章制度.pdf_0']],
    'documents': [['内容A']],
    'metadatas': [[{'source': '规章制度.pdf', '_collection': 'public_kb', 'collection': 'public_kb'}]],
    'distances': [[0.9]],
}
result_bm25 = {
    'ids': [['规章制度.pdf_0']],
    'documents': [['内容A']],
    'metadatas': [[{'source': '规章制度.pdf', '_collection': 'public_kb', 'collection': 'public_kb'}]],
    'distances': [[0.7]],
}

rrf_same = reciprocal_rank_fusion_test([result_vec, result_bm25])
check("同库向量+BM25合并：结果数应为 1", len(rrf_same['ids'][0]) == 1, f"实际: {len(rrf_same['ids'][0])}")

# chunk_index 解析测试（带前缀的 ID）
print("\n--- chunk_index 解析兼容性 ---")
for test_id, expected in [
    ("public_kb/规章制度.pdf_0", 0),
    ("dept_1_kb/规章制度.pdf_5", 5),
    ("规章制度.pdf_3", 3),
    ("test_.pdf_0", 0),
]:
    chunk_id_raw = test_id
    try:
        chunk_index = int(str(chunk_id_raw).rsplit('_', 1)[-1])
    except (ValueError, IndexError):
        chunk_index = None
    check(f"解析 '{test_id}' -> chunk_index={expected}", chunk_index == expected, f"实际: {chunk_index}")


# ======================================================================
# 测试 2：search_multiple 去重
# ======================================================================
print("\n=== 测试 2：search_multiple 复合键去重 ===")

# 模拟 search.py 的 _merge_multiple_results 逻辑
class FakeSearchResult:
    def __init__(self, ids, documents, metadatas, distances, collection_name=""):
        self.ids = ids
        self.documents = documents
        self.metadatas = metadatas
        self.distances = distances
        self.collection_name = collection_name


def merge_multiple_results_test(results, top_k=10):
    """复刻 search.py 的 _merge_multiple_results 逻辑（含修复）"""
    if not results:
        return []
    all_items = []
    for result in results:
        for i, doc_id in enumerate(result.ids):
            all_items.append({
                'id': doc_id,
                'doc': result.documents[i],
                'meta': result.metadatas[i],
                'score': result.distances[i],
                'collection': result.collection_name
            })
    all_items.sort(key=lambda x: x['score'], reverse=True)

    seen = set()
    unique_items = []
    for item in all_items:
        composite_key = (item['collection'], item['id'])
        if composite_key not in seen:
            seen.add(composite_key)
            unique_items.append(item)
    return unique_items[:top_k]


result_a = FakeSearchResult(
    ids=['规章制度.pdf_0', '规章制度.pdf_1'],
    documents=['A内容0', 'A内容1'],
    metadatas=[{'source': '规章制度.pdf'}, {'source': '规章制度.pdf'}],
    distances=[0.9, 0.8],
    collection_name='public_kb'
)
result_b = FakeSearchResult(
    ids=['规章制度.pdf_0', '规章制度.pdf_1'],
    documents=['B内容0', 'B内容1'],
    metadatas=[{'source': '规章制度.pdf'}, {'source': '规章制度.pdf'}],
    distances=[0.85, 0.75],
    collection_name='dept_1_kb'
)

merged = merge_multiple_results_test([result_a, result_b])
check("跨库去重：结果数应为 4", len(merged) == 4, f"实际: {len(merged)}")
check("包含 public_kb 的内容", any('A内容' in m['doc'] for m in merged))
check("包含 dept_1_kb 的内容", any('B内容' in m['doc'] for m in merged))

# 同库应去重
result_a2 = FakeSearchResult(
    ids=['规章制度.pdf_0'],
    documents=['A内容0-vec'],
    metadatas=[{'source': '规章制度.pdf'}],
    distances=[0.9],
    collection_name='public_kb'
)
result_a3 = FakeSearchResult(
    ids=['规章制度.pdf_0'],
    documents=['A内容0-bm25'],
    metadatas=[{'source': '规章制度.pdf'}],
    distances=[0.7],
    collection_name='public_kb'
)
merged_same = merge_multiple_results_test([result_a2, result_a3])
check("同库去重：结果数应为 1", len(merged_same) == 1, f"实际: {len(merged_same)}")


# ======================================================================
# 测试 3：DocStore 路径含 collection 前缀
# ======================================================================
print("\n=== 测试 3：DocStore 路径生成 ===")

def docstore_path_test(doc_id, metadata):
    """复刻 processing.py 的路径逻辑"""
    coll = metadata.get('collection', '')
    safe_id = f"{coll}_{doc_id}" if coll else doc_id
    return f"{safe_id}.json"


path1 = docstore_path_test("规章制度.pdf_2", {"collection": "public_kb"})
path2 = docstore_path_test("规章制度.pdf_2", {"collection": "dept_1_kb"})
check("跨库路径不同", path1 != path2, f"path1={path1}, path2={path2}")
check("public_kb 路径含前缀", path1 == "public_kb_规章制度.pdf_2.json", f"实际: {path1}")
check("dept_1_kb 路径含前缀", path2 == "dept_1_kb_规章制度.pdf_2.json", f"实际: {path2}")

# 无 collection 的兼容
path3 = docstore_path_test("test.pdf_0", {})
check("无 collection 时保持原始路径", path3 == "test.pdf_0.json", f"实际: {path3}")


# ======================================================================
# 测试 4：citation 构建 — collection 字段兼容
# ======================================================================
print("\n=== 测试 4：citation collection 字段兼容 ===")

def citation_collection_test(meta):
    """复刻 chat_routes.py 的 collection 取值逻辑"""
    return meta.get('_collection') or meta.get('collection', '')


# 场景 A：只有 _collection（多库检索路径）
check("_collection 优先", citation_collection_test({'_collection': 'dept_1_kb', 'collection': 'public_kb'}) == 'dept_1_kb')
# 场景 B：只有 collection（入库时的 metadata）
check("回退到 collection", citation_collection_test({'collection': 'public_kb'}) == 'public_kb')
# 场景 C：两者都没有
check("都没有时返回空", citation_collection_test({}) == '')
# 场景 D：_collection 为空字符串
check("_collection 空字符串时回退", citation_collection_test({'_collection': '', 'collection': 'public_kb'}) == 'public_kb')


# ======================================================================
# 测试 5：_collection 回退逻辑
# ======================================================================
print("\n=== 测试 5：_collection 回退逻辑 ===")

def fallback_collection_test(meta, collections):
    """复刻 chat_routes.py 修复后的回退逻辑"""
    if not meta.get('_collection'):
        meta['_collection'] = meta.get('collection') or (collections[0] if collections else 'public_kb')
    return meta['_collection']


# 场景：meta 有 collection 字段但无 _collection
m1 = {'collection': 'dept_1_kb'}
check("使用 meta 中的 collection", fallback_collection_test(m1, ['public_kb', 'dept_1_kb']) == 'dept_1_kb')

# 场景：meta 两者都没有
m2 = {}
check("回退到 collections[0]", fallback_collection_test(m2, ['public_kb', 'dept_1_kb']) == 'public_kb')

# 场景：已有 _collection 不覆盖
m3 = {'_collection': 'dept_1_kb', 'collection': 'public_kb'}
check("已有 _collection 不覆盖", fallback_collection_test(m3, ['public_kb']) == 'dept_1_kb')


# ======================================================================
# 测试 6：_expand_contiguous_chunks ID 匹配兼容
# ======================================================================
print("\n=== 测试 6：上下文扩展 ID 匹配兼容 ===")

# 模拟 existing_ids 和 n_id 的匹配逻辑
existing_ids = {'public_kb/规章制度.pdf_0', 'public_kb/规章制度.pdf_1', 'dept_1_kb/制度.pdf_3'}
_raw_id_set = set()
for _eid in existing_ids:
    if '/' in _eid:
        _raw_id_set.add(_eid.split('/', 1)[1])
    else:
        _raw_id_set.add(_eid)

# 邻居查询返回的是原始 ID（无前缀）
n_id_raw = "规章制度.pdf_0"
n_id_new = "规章制度.pdf_5"

is_dup_raw = n_id_raw in existing_ids or n_id_raw in _raw_id_set
is_dup_new = n_id_new in existing_ids or n_id_new in _raw_id_set

check("原始 ID 能匹配带前缀的 existing_ids", is_dup_raw == True)
check("新 ID 不误判为已存在", is_dup_new == False)


# ======================================================================
# 汇总
# ======================================================================
print(f"\n{'='*50}")
print(f"测试结果：{passed} 通过 / {failed} 失败 / {passed+failed} 总计")
if failed > 0:
    print("存在失败项，请检查！")
    sys.exit(1)
else:
    print("全部通过！")