"""遍历 _001__clawler/方剂/ 下的方剂详情 txt，调用 LLM 抽取实体和关系，一方一个 json。

和 01_extract_herb_graph.py 是一对镜像脚本，逻辑一致，只换输入输出目录 —— 改一边
的行为时记得看另一边要不要同步。prompt 不用换：__000__ 里那套已经带了方剂的
FormulaAttributes 和 HAS_INGREDIENT / FROM_SOURCE，实测方剂文本能正确抽出组成药材
和出处。

支持断点续抽：已经存在对应 json 的方剂直接跳过，中断后重跑即可继续。

输出：
    _002__extract_information/方剂图谱/<方剂名>.json   成功（合法的空结果也算成功）
    _002__extract_information/方剂图谱/_failed.log     失败清单，每行「方剂名<TAB>原因」
    _002__extract_information/alpaca_formula_dataset.jsonl  顺便追加的 Alpaca 微调样本

数据集是抽取时顺手追加的，只覆盖「本次抽的」那些方剂。已经抽过的会被断点续抽跳过，
不会再写进数据集 —— 要补齐全量或去重，跑 04_build_formula_alpaca_dataset.py 从
方剂图谱/*.json 重建。

方剂文本比中药短得多（均值 377 字 vs 中药 1537），单条实测 1~2 分钟，
703 条 10 线程大约 2~3 小时。

用法（在项目根目录执行）：
    python -X utf8 -m _002__extract_information.03_extract_formula_graph --limit 5   # 先抽 5 条看效果
    python -X utf8 -m _002__extract_information.03_extract_formula_graph             # 全量，10 线程
    python -X utf8 -m _002__extract_information.03_extract_formula_graph --workers 20
    python -X utf8 -m _002__extract_information.03_extract_formula_graph --retry-failed
"""

import argparse
import json
import os
import threading
import time

from concurrent.futures import ThreadPoolExecutor, as_completed

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,
    extract_tcm_knowledge,
)

INPUT_DIR = get_file_path("_001__clawler/方剂")
OUTPUT_DIR = get_file_path("_002__extract_information/方剂图谱")
FAILED_LOG = os.path.join(OUTPUT_DIR, "_failed.log")
# 微调数据集用 jsonl 追加：一行一条，多线程各写一行不会互相截断，
# 中途崩了也只会丢最后一行，不会把前面写好的记录带坏。
ALPACA_JSONL = get_file_path("_002__extract_information/alpaca_formula_dataset.jsonl")

# 失败清单是多线程追加写的，加锁避免两行日志挤在一起
_failed_lock = threading.Lock()
# 数据集同理，一次只让一个线程追加
_alpaca_lock = threading.Lock()


def read_failed_names():
    """读失败清单里的方剂名，用于 --retry-failed。"""
    if not os.path.exists(FAILED_LOG):
        return set()
    names = set()
    with open(FAILED_LOG, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                names.add(line.split("\t")[0])
    return names


def record_failure(name, reason):
    """把失败原因追加到清单。写成一行，方便之后 --retry-failed 挑出来重跑。"""
    reason = " ".join(str(reason).split())[:200]
    with _failed_lock:
        with open(FAILED_LOG, "a", encoding="utf-8") as f:
            f.write(f"{name}\t{reason}\n")


def append_alpaca(text, entities, relations):
    """把这一条抽取结果追加成一行 Alpaca 样本。

    只追加、不去重 —— 断点续抽保证同一个方剂一轮里只会抽一次。真要补历史或
    去重，跑 04_build_formula_alpaca_dataset.py 从图谱 json 重建。
    """
    record = build_alpaca_record(text, entities, relations)
    line = json.dumps(record, ensure_ascii=False)
    with _alpaca_lock:
        with open(ALPACA_JSONL, "a", encoding="utf-8") as f:
            f.write(line + "\n")


def extract_one(name, retries=2):
    """抽一个方剂并落盘。返回 (name, ok, note, 实体数, 关系数)。由线程池并发调用。

    LLM 偶发超时/限流/返回非法 json，失败退避重试；仍然失败就记进清单，
    不写 json —— 这样下次跑还会被当成待抽项，断点续抽才是对的。
    """
    input_path = os.path.join(INPUT_DIR, f"{name}.txt")
    output_path = os.path.join(OUTPUT_DIR, f"{name}.json")

    try:
        with open(input_path, encoding="utf-8") as f:
            text = f.read()
    except OSError as exc:
        record_failure(name, f"读取失败: {exc}")
        return name, False, "读取失败", 0, 0

    if not text.strip():
        record_failure(name, "源文件为空")
        return name, False, "源文件为空", 0, 0

    last_exc = None
    for attempt in range(1, retries + 1):
        try:
            result = extract_tcm_knowledge(text)
            break
        except Exception as exc:  # 网络、限流、json 解析失败都在这
            last_exc = exc
            if attempt < retries:
                time.sleep(5 * attempt)
    else:
        record_failure(name, f"抽取失败: {last_exc}")
        return name, False, "抽取失败", 0, 0

    # parser 是 JsonOutputParser，正常返回 dict；模型偶尔会吐出 list 或别的形状
    if not isinstance(result, dict):
        record_failure(name, f"返回结构异常: {type(result).__name__}")
        return name, False, "结构异常", 0, 0

    entities = result.get("entities") or []
    relations = result.get("relations") or []
    # 空结果也落盘，否则每次重跑都会重复调用这些方剂
    payload = {
        "name": name,
        "source_file": os.path.relpath(input_path, get_file_path("")),
        "entities": entities,
        "relations": relations,
    }

    # 先写临时文件再改名，中断时不会留下半个 json 被误判成已完成
    temp_path = f"{output_path}.{threading.get_ident()}.part"
    with open(temp_path, "w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)
    os.replace(temp_path, output_path)

    # 图谱 json 落盘之后再追加数据集：反过来的话，万一追加完崩了、json 还没写成，
    # 续抽会把这个方剂重抽一遍，数据集里就多出一条重复样本。
    append_alpaca(text, entities, relations)

    note = f"{len(entities)}实体/{len(relations)}关系"
    return name, True, note, len(entities), len(relations)


def build_todo(limit=0, retry_failed=False):
    """列出本次要抽的方剂名。断点续抽的核心：已有 json 的直接排除。"""
    names = sorted(
        os.path.splitext(f)[0]
        for f in os.listdir(INPUT_DIR)
        if f.endswith(".txt")
    )
    done = {
        os.path.splitext(f)[0]
        for f in os.listdir(OUTPUT_DIR)
        if f.endswith(".json")
    }
    todo = [n for n in names if n not in done]

    if retry_failed:
        # 失败的本来就没写 json，已经在 todo 里；这里是缩小范围只跑它们
        failed = read_failed_names()
        todo = [n for n in todo if n in failed]

    if limit:
        todo = todo[:limit]
    return names, done, todo


def run(limit=0, workers=10, retry_failed=False):
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    names, done, todo = build_todo(limit=limit, retry_failed=retry_failed)

    print(f"源目录共 {len(names)} 个方剂，已抽取 {len(done)} 个，本次待抽 {len(todo)} 个")
    print(f"并发线程数: {workers}")
    if not todo:
        print("没有需要抽取的内容。")
        return

    ok_count = fail_count = 0
    entity_total = relation_total = 0

    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = {pool.submit(extract_one, name): name for name in todo}
        # tqdm 显示进度；每条结果用 write() 打，避免把进度条撕开
        with tqdm(total=len(todo), desc="抽取中", unit="方", ncols=90) as bar:
            for future in as_completed(futures):
                name = futures[future]
                try:
                    _, ok, note, n_entity, n_relation = future.result()
                except Exception as exc:  # 单条异常不该拖垮整轮
                    record_failure(name, f"未捕获异常: {exc}")
                    ok, note, n_entity, n_relation = False, "异常", 0, 0
                if ok:
                    ok_count += 1
                    entity_total += n_entity
                    relation_total += n_relation
                else:
                    fail_count += 1
                bar.write(f"  {'✓' if ok else '✗'} {name} {note}")
                bar.update(1)
                bar.set_postfix(成功=ok_count, 失败=fail_count)

    print(f"\n完成 {ok_count} 个，失败 {fail_count} 个")
    print(f"共抽出 {entity_total} 个实体、{relation_total} 条关系")
    print(f"输出目录: {OUTPUT_DIR}")
    print(f"微调数据集（追加 {ok_count} 条）: {ALPACA_JSONL}")
    if fail_count:
        print(f"失败清单: {FAILED_LOG}（可用 --retry-failed 只重跑这些）")


def main():
    parser = argparse.ArgumentParser(description="抽取方剂详情中的实体和关系")
    parser.add_argument(
        "--limit",
        type=int,
        default=0,
        help="本次最多抽多少个，0 表示不限制（默认 0）",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=40,
        help="并发线程数（默认 10，设 1 就是单线程）",
    )
    parser.add_argument(
        "--retry-failed",
        action="store_true",
        help="只重跑 _failed.log 里记过的方剂",
    )
    args = parser.parse_args()
    if args.workers < 1:
        parser.error("--workers 至少为 1")
    run(limit=args.limit, workers=args.workers, retry_failed=args.retry_failed)


if __name__ == "__main__":
    main()