# -*- coding: utf-8 -*-
"""
对 queryData.jsonl 进行清洗：
1. 删除 questionScore < 7 的记录
2. 对 instruction 文本去重，保留分数最高的一条
3. 重新编号 id，写回原文件
"""
import json
import os


def cleanQueryData(targetFile, minScore=7):
    # ---------- 读取全部记录 ----------
    records = []
    with open(targetFile, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    print(f"原始记录数: {len(records)}")

    # ---------- 第一步：过滤低分记录 ----------
    filtered = []
    removedByScore = 0
    for rec in records:
        score = rec.get("questionScore", -1)
        if score >= minScore: # 打分大于等于阈值7分, 就保留
            filtered.append(rec)
        else:
            removedByScore += 1
    print(f"删除 questionScore < {minScore} 的记录: {removedByScore} 条")
    print(f"过滤后剩余: {len(filtered)} 条")

    # ---------- 第二步：按 instruction 去重 ----------
    # 用字典记录每个 instruction 对应的最佳记录（分数最高的）
    bestByInstruction = {}
    duplicateCount = 0
    for rec in filtered:
        instr = rec["instruction"]
        if instr not in bestByInstruction:
            bestByInstruction[instr] = rec
        else:
            # 保留分数更高的那条
            existing = bestByInstruction[instr]
            if rec.get("questionScore", 0) > existing.get("questionScore", 0):
                bestByInstruction[instr] = rec
            duplicateCount += 1

    deduped = list(bestByInstruction.values())
    print(f"去重删除重复记录: {duplicateCount} 条")
    print(f"去重后剩余: {len(deduped)} 条")

    # ---------- 第三步：重新编号 id ----------
    for i, rec in enumerate(deduped):
        rec["id"] = i + 1

    # ---------- 第四步：写回原文件 ----------
    with open(targetFile, 'w', encoding='utf-8') as f:
        for rec in deduped:
            f.write(json.dumps(rec, ensure_ascii=False) + '\n')

    print(f"\n清洗完成，共 {len(deduped)} 条记录已保存到 {targetFile}")

    # ---------- 打印各类别统计 ----------
    print(f"\n各类别剩余数量:")
    categoryCount = {}
    for rec in deduped:
        cate = rec.get("category", "")
        categoryCount[cate] = categoryCount.get(cate, 0) + 1
    for cate, count in categoryCount.items():
        print(f"  {cate}: {count} 条")

    # ---------- 打印被删除的低分记录 ----------
    if removedByScore > 0:
        print(f"\n被删除的低分记录(score<{minScore}):")
        for rec in records:
            if rec.get("questionScore", -1) < minScore:
                print(f"  id={rec['id']}, score={rec['questionScore']}, "
                      f"category={rec['category']}, instruction={rec['instruction'][:30]}...")

    # ---------- 打印被去重的记录 ----------
    if duplicateCount > 0:
        print(f"\n被去重的重复记录:")
        seen = set()
        for rec in filtered:
            instr = rec["instruction"]
            if instr in seen:
                print(f"  id={rec['id']}, score={rec['questionScore']}, "
                      f"category={rec['category']}, instruction={rec['instruction'][:30]}...")
            else:
                seen.add(instr)


if __name__ == "__main__":
    targetFile = "../data/testDataTemp/queryData.jsonl"
    cleanQueryData(targetFile, minScore=7)