"""datainventor: curate a pretraining corpus on your own machines.

  datainventor init recipe.yaml                    write a starter recipe to edit
  datainventor preview recipe.yaml                 run every stage on a small sample and print what each did
  datainventor run recipe.yaml --out ./curated     run it in full (resumable)

The model endpoint is any OpenAI-compatible server: vLLM, SGLang, Ollama, a hosted API. Set its key in the
environment variable named by endpoint.api_key_env (DI_API_KEY by default).
"""
from __future__ import annotations

import argparse
import json
import os
import sys

from .compose import fmt
from .recipe import DEFAULT, dump, endpoint_of, load


def main(argv=None):
    ap = argparse.ArgumentParser(prog="datainventor", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    a = sub.add_parser("init", help="write a starter recipe")
    a.add_argument("path")
    a = sub.add_parser("preview", help="run every stage on a sample")
    a.add_argument("recipe")
    a.add_argument("--per-source", type=int, default=300)
    a.add_argument("--grade", type=int, default=64)
    a.add_argument("--rephrase", type=int, default=6)
    a.add_argument("--workers", type=int, default=None)
    a.add_argument("--json", help="also write the full report here")
    a = sub.add_parser("run", help="run the recipe in full")
    a.add_argument("recipe")
    a.add_argument("--out", default=None, help="output folder (default: output.dir in the recipe)")
    a.add_argument("--procs", type=int, default=0, help="CPU processes for cleaning and scoring (default: all but one)")
    a.add_argument("--workers", type=int, default=None, help="parallel model calls (default: endpoint.workers)")
    a.add_argument("--stages", default="grade,curate,create,compose")
    a.add_argument("--limit", type=int, default=0, help="read at most this many documents per source (for a trial run)")
    args = ap.parse_args(argv)

    if args.cmd == "init":
        if os.path.exists(args.path):
            sys.exit(f"{args.path} exists; not overwriting it.")
        r = json.loads(json.dumps(DEFAULT))
        r["sources"] = [
            {"name": "finepdfs", "dataset": "codelion/finepdfs-1B", "kind": "textbook"},
            {"name": "dclm", "dataset": "codelion/dclm-baseline-1B", "kind": "web"},
            {"name": "fineweb-edu", "dataset": "codelion/fineweb-edu-1B", "kind": "educational"},
        ]
        r["targets"]["description"] = "grade-school to university math and science reasoning"
        r["compose"]["mixes"][0]["budget_tokens"] = 2_000_000_000
        with open(args.path, "w") as f:
            f.write(dump(r))
        print(f"Wrote {args.path}. Edit the sources, targets and budget, then: datainventor preview {args.path}")
        return

    r = load(args.recipe)
    if args.cmd == "preview":
        from .preview import Preview, run_preview
        from .sources import inspect
        token = os.environ.get("HF_TOKEN", "")
        for s in r["sources"]:  # row counts make the token estimates possible
            if s.get("dataset") and not s.get("num_rows"):
                try:
                    s["num_rows"] = inspect(s["dataset"], token)["num_rows"] if (s.get("config", "default") == "default") else 0
                except Exception:
                    pass
        p = Preview(r)
        ep = endpoint_of(r, args.workers)
        run_preview(p, ep, token, per_source=args.per_source, grade_n=args.grade, rephrase_n=args.rephrase, invent_n=2)
        if p.state != "done":
            sys.exit(f"Preview {p.state}: {p.message}")
        rep = p.report
        for f in rep["sources"]:
            drops = ", ".join(f"{k} {v}" for k, v in sorted(f["dropped"].items(), key=lambda x: -x[1]))
            print(f"{f['name']}: read {f['read']}, selected {f['selected']}; dropped: {drops}; decontaminated: {f['decontaminated'] or 'none'}")
            if f.get("available_tokens"):
                print(f"  about {fmt(f['available_tokens'])} tokens available after curation")
        c = rep["stages"]["calibrate"]
        print(f"graded {c['graded']}; scorer agreement {c['agreement']}")
        for plan in rep["stages"]["compose"]["plans"]:
            print(f"mix {plan['name']} ({plan['budget']}): " + "; ".join(f"{x['name']} {int(x['share'] * 100)}% x{x['epochs']}" for x in plan["rows"]))
            for w in plan["warnings"]:
                print("  - " + w)
        if args.json:
            with open(args.json, "w") as fh:
                json.dump({"report": rep, "docs": p.docs}, fh, indent=1, ensure_ascii=False)
        return

    from .runner import Run
    out = args.out or r["output"]["dir"]
    Run(r, out, procs=args.procs, workers=args.workers, limit=args.limit).go(tuple(args.stages.split(",")))
    print(f"Done. See {os.path.join(out, 'REPORT.md')}")


if __name__ == "__main__":
    main()