#!/usr/bin/env python3
import json
import re
from pathlib import Path


documents = [
    {"id": "policy-2024", "text": "旧政策：退款期限为签收后7天。", "date": "2024-01-01"},
    {"id": "policy-2026", "text": "现行政策：退款期限为签收后14天。", "date": "2026-05-01"},
    {"id": "shipping", "text": "普通快递预计3至5天送达。", "date": "2026-01-01"},
] + [{"id": f"noise-{i}", "text": f"第{i}条无关产品介绍。", "date": "2025-01-01"} for i in range(17)]


def tokens(text):
    return set(re.findall(r"[\u4e00-\u9fff]{2,}|\d+", text))


def retrieve(query, top_k=3):
    q = tokens(query)
    ranked = sorted(
        documents,
        key=lambda d: (len(q & tokens(d["text"])), d["date"]),
        reverse=True,
    )
    return ranked[:top_k]


selected = retrieve("现在退款期限是多少天？")
current = max((d for d in selected if "退款期限" in d["text"]), key=lambda d: d["date"])
result = {
    "all_documents": len(documents),
    "selected_ids": [d["id"] for d in selected],
    "answer_evidence": current,
    "answer": "14天" if "14天" in current["text"] else "unknown",
}
Path(__file__).with_name("experiment-result.json").write_text(
    json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
)
print(json.dumps(result, ensure_ascii=False, indent=2))
