from law_rag.retrieval import high_level_retrieval


def law_rag_run(questions, with_score: bool = False):
    retriever = high_level_retrieval.HighLevelRetriever()
    results = retriever.search(questions)

    print(f"\n总共找到 {results['total_results']} 条相关法律条文")
    print(f"耗时: {results['time_taken']:.2f}秒")

    law_retrieval_result_list = []
    for i, item in enumerate(results["results"][:10]):
        print(f"\n[{i + 1}] 相似度: {item.get('original_score', 0):.4f}")
        print(f"来源: {item.get('source', '未知')}")
        if "purpose" in item:
            print(f"目的: {item['purpose']}")

        law_id = item['metadata'].get('law_id', '未知')
        content = item['text']

        print(f"法律ID: {law_id}")
        print(f"内容: {content}")

        row = {"law_id": law_id, "content": content}
        if with_score:
            row["score"] = item.get("original_score", 0)
            row["source"] = item.get("source", "未知")
            if "purpose" in item:
                row["purpose"] = item["purpose"]
        law_retrieval_result_list.append(row)

    return law_retrieval_result_list

if __name__ == '__main__':
    result = law_rag_run("欠薪三月未还")
    print(str(result))