"""从 中药图谱/*.json 重建 Alpaca 格式的微调数据集，不再调用 LLM。

01_extract_herb_graph.py 抽取时会顺手往 alpaca_herb_dataset.jsonl 追加样本，但那只
覆盖「那一次抽的」药材：断点续抽跳过的、以及加这个功能之前就抽好的，都不在里面。
这个脚本直接读已经落盘的图谱 json 重新生成全量数据集，所以是幂等的 —— 跑几次结果
都一样，也不花钱。

input 取源 txt 原文（图谱 json 里的 source_file 指回去），output 取 entities+relations
的 json 串，instruction 固定。

输出：
    _002__extract_information/alpaca_herb_dataset.jsonl   一行一条（默认）
    _002__extract_information/alpaca_herb_dataset.json    标准 json 数组（--json）

用法（在项目根目录执行）：
    python -X utf8 -m _002__extract_information.02_build_alpaca_dataset
    python -X utf8 -m _002__extract_information.02_build_alpaca_dataset --json
    python -X utf8 -m _002__extract_information.02_build_alpaca_dataset --skip-empty
"""

import argparse
import json
import os

from tqdm import tqdm

from common.path_utils import get_file_path
from _002__extract_information.__000__extract_graph_data_utils import (
    build_alpaca_record,
)

GRAPH_DIR = get_file_path("_002__extract_information/中药图谱")
HERB_DIR = get_file_path("_001__clawler/中药")
OUT_JSONL = get_file_path("_002__extract_information/alpaca_herb_dataset.jsonl")
OUT_JSON = get_file_path("_002__extract_information/alpaca_herb_dataset.json")


def load_source_text(graph, graph_path):
    """取这条图谱对应的中药详情原文，作为 Alpaca 的 input。

    优先用 json 里记的 source_file（相对仓库根），它不在或者文件没了就按
    「中药名.txt」回退到 中药/ 目录 —— 早期抽的 json 可能没这个字段。
    """
    rel = graph.get("source_file")
    if rel:
        # json 里存的是 Windows 风格的相对路径，normpath 一下跨平台都能用
        path = get_file_path(os.path.normpath(rel))
        if os.path.exists(path):
            return open(path, encoding="utf-8").read()

    name = graph.get("name") or os.path.splitext(os.path.basename(graph_path))[0]
    path = os.path.join(HERB_DIR, f"{name}.txt")
    if os.path.exists(path):
        return open(path, encoding="utf-8").read()
    return None


def build(skip_empty=False, as_json=False):
    if not os.path.isdir(GRAPH_DIR):
        print(f"图谱目录不存在: {GRAPH_DIR}")
        print("先跑 01_extract_herb_graph.py 抽取实体和关系。")
        return

    graph_files = sorted(
        os.path.join(GRAPH_DIR, f)
        for f in os.listdir(GRAPH_DIR)
        if f.endswith(".json")
    )
    if not graph_files:
        print(f"{GRAPH_DIR} 下没有图谱 json，先跑 01_extract_herb_graph.py。")
        return

    records = []
    empty = missing_text = broken = 0

    for path in tqdm(graph_files, desc="构建数据集", unit="条", ncols=90):
        try:
            with open(path, encoding="utf-8") as f:
                graph = json.load(f)
        except (OSError, json.JSONDecodeError):
            broken += 1
            continue

        entities = graph.get("entities") or []
        relations = graph.get("relations") or []
        # 空结果是合法抽取（prompt 明确要求单实体无关系时返回空），留着能教模型
        # 「这种文本不该硬凑关系」；真要练抽取能力可以 --skip-empty 去掉。
        if skip_empty and not entities and not relations:
            empty += 1
            continue
        if not entities and not relations:
            empty += 1

        text = load_source_text(graph, path)
        if text is None or not text.strip():
            missing_text += 1
            continue

        records.append(build_alpaca_record(text, entities, relations))

    if as_json:
        out_path = OUT_JSON
        payload = json.dumps(records, ensure_ascii=False, indent=2)
    else:
        out_path = OUT_JSONL
        payload = "\n".join(json.dumps(r, ensure_ascii=False) for r in records)
        payload += "\n" if records else ""

    # 整体重写，不是追加：这个脚本的语义是「按当前图谱重建」，幂等
    temp_path = f"{out_path}.part"
    with open(temp_path, "w", encoding="utf-8") as f:
        f.write(payload)
    os.replace(temp_path, out_path)

    print(f"\n图谱 json {len(graph_files)} 个 → 样本 {len(records)} 条")
    if empty:
        print(f"  其中空抽取结果 {empty} 条{'（已跳过）' if skip_empty else '（已保留）'}")
    if missing_text:
        print(f"  找不到源 txt 跳过 {missing_text} 条")
    if broken:
        print(f"  json 读取失败跳过 {broken} 条")
    print(f"输出: {out_path}")


def main():
    parser = argparse.ArgumentParser(
        description="从已抽取的图谱 json 重建 Alpaca 微调数据集"
    )
    parser.add_argument(
        "--json",
        action="store_true",
        dest="as_json",
        help="输出标准 json 数组，默认输出 jsonl（一行一条）",
    )
    parser.add_argument(
        "--skip-empty",
        action="store_true",
        help="丢掉 entities/relations 都为空的样本（默认保留）",
    )
    args = parser.parse_args()
    build(skip_empty=args.skip_empty, as_json=args.as_json)


if __name__ == "__main__":
    main()