"""Export the mixed finance dataset (from prepare_datasets) to LLaMA-Factory's
ShareGPT format, and register it in dataset_info.json.

    python -m src.data.export_llamafactory --data data/finance_sft --out data/llamafactory

Then point LLaMA-Factory at it:
    llamafactory-cli train configs/llamafactory/sft_general_qwen25_7b.yaml
(the config sets dataset_dir: data/llamafactory, dataset: finance_sft)
"""

import argparse
import json
import pathlib

from datasets import load_from_disk

ROLE_MAP = {"user": "human", "assistant": "gpt"}


def to_sharegpt(example):
    system = ""
    conversations = []
    for m in example["messages"]:
        if m["role"] == "system":
            system = m["content"]
        else:
            conversations.append({"from": ROLE_MAP[m["role"]], "value": m["content"]})
    return {"conversations": conversations, "system": system}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="data/finance_sft")
    ap.add_argument("--out", default="data/llamafactory")
    ap.add_argument("--name", default="finance_sft")
    args = ap.parse_args()

    outdir = pathlib.Path(args.out)
    outdir.mkdir(parents=True, exist_ok=True)
    ds = load_from_disk(args.data)

    for split, fname in [("train", f"{args.name}.json"), ("test", f"{args.name}_eval.json")]:
        if split not in ds:
            continue
        rows = [to_sharegpt(ex) for ex in ds[split]]
        (outdir / fname).write_text(json.dumps(rows, ensure_ascii=False, indent=1))
        print(f"[ok] {split}: {len(rows)} -> {outdir / fname}")

    info_path = outdir / "dataset_info.json"
    info = json.loads(info_path.read_text()) if info_path.exists() else {}
    entry_tags = {
        "role_tag": "from",
        "content_tag": "value",
        "user_tag": "human",
        "assistant_tag": "gpt",
        "system_tag": "system",
    }
    info[args.name] = {
        "file_name": f"{args.name}.json",
        "formatting": "sharegpt",
        "columns": {"messages": "conversations", "system": "system"},
        "tags": entry_tags,
    }
    if (outdir / f"{args.name}_eval.json").exists():
        info[f"{args.name}_eval"] = {
            "file_name": f"{args.name}_eval.json",
            "formatting": "sharegpt",
            "columns": {"messages": "conversations", "system": "system"},
            "tags": entry_tags,
        }
    info_path.write_text(json.dumps(info, indent=2))
    print(f"[ok] registered '{args.name}' in {info_path}")


if __name__ == "__main__":
    main()