# /mnt/petrelfs/zhengzhijie/mutiAgent4Fraud/utils/dataset_make/export_bad_pairs_jsonl.py
# -*- coding: utf-8 -*-

import os
import sys
import json

# 写死参数（按需修改）
DB_PATH = "/mnt/petrelfs/zhengzhijie/mutiAgent4Fraud/data/simu_db/expierement/different_bad_model/deepseek-r1.db"
BAD_ID_START = 100
BAD_ID_END = 109
INCLUDE_EVENTS = True
MIN_DEPTH = 3  # 新增：深度阈值，小于该值不导出

# 输出文件命名
OUTPUT_DIR = "/mnt/petrelfs/zhengzhijie/mutiAgent4Fraud/data/datasets"
MODEL_BASENAME = os.path.basename(DB_PATH).replace(".db", "")
OUTPUT_JSONL = os.path.join(OUTPUT_DIR, f"{MODEL_BASENAME}_bad_{BAD_ID_START}_{BAD_ID_END}_min{MIN_DEPTH}.jsonl")

# 便捷导入 build_pair_trajectory（与本文件同目录）
THIS_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, THIS_DIR)
from mutiAgent4Fraud.utils.dataset_make.cooperation.get_history import build_pair_trajectory  # noqa: E402


def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    line_id = 1
    exported = 0
    scanned_pairs = 0

    with open(OUTPUT_JSONL, "w", encoding="utf-8") as fout:
        for uid1 in range(BAD_ID_START, BAD_ID_END + 1):
            for uid2 in range(uid1 + 1, BAD_ID_END + 1):
                scanned_pairs += 1
                result = build_pair_trajectory(
                    db_path=DB_PATH,
                    uid1=uid1,
                    uid2=uid2,
                    include_events=INCLUDE_EVENTS,
                )
                if result is None:
                    continue  # 无对话，跳过

                depth = int(result.get("对话的深度", 0))
                if depth < MIN_DEPTH:
                    continue  # 过滤深度小于 MIN_DEPTH 的对话

                obj = {
                    "id": line_id,
                    "模型种类": result["模型种类"],
                    "对话的深度": depth,
                    "对话的内容": result["对话的内容"],
                }
                fout.write(json.dumps(obj, ensure_ascii=False) + "\n")
                line_id += 1
                exported += 1

    print(f"扫描对数: {scanned_pairs}，导出条数(深度≥{MIN_DEPTH}): {exported}，保存至: {OUTPUT_JSONL}")


if __name__ == "__main__":
    main()