"""Run a recipe in full on your own machines.

    datainventor run recipe.yaml --out ./curated --procs 16 --workers 32

Stages, each resumable (re-run the same command after an interruption):
  grade     grade a sample of every source with your model, fit the scorers   -> calibrator.json, grades.jsonl
  curate    stream every document: clean, decontaminate, dedup, score          -> curated/<source>/part-*.jsonl
  create    rephrase the best documents where a mix needs more tokens, invent  -> synthetic/<source>/part-*.jsonl
  compose   build each mix to its token budget                                 -> mixes/<mix>/part-*.jsonl
Then REPORT.md and report.json summarize what every stage did.

Token counts are estimates (four characters per token). Calibrate them against your tokenizer if the budget is tight.
"""
from __future__ import annotations

import itertools
import json
import os
import random
import shutil
import time
from concurrent.futures import ThreadPoolExecutor
from multiprocessing import get_context

from . import calibrate as cal
from .clean import Cleaner, Decontaminator, benchmark_texts
from .compose import default_weights, fmt, plan_mix
from .create import chunks, invent, rephrase, styles_for
from .preview import plan_topics
from .recipe import dump, endpoint_of, held_out_of, sources_of, targets_of
from .sources import Source, stream
from .text import MinHasher, NearDupIndex, doc_hash, est_tokens, h64

STAGES = ("grade", "curate", "create", "compose")


class Shards:
    """Append-only JSONL shards of a fixed number of documents."""

    def __init__(self, folder: str, per: int):
        os.makedirs(folder, exist_ok=True)
        self.folder, self.per, self.n, self.f, self.count = folder, per, 0, None, 0

    def write(self, doc: dict):
        if self.f is None or self.count >= self.per:
            if self.f:
                self.f.close()
            self.f = open(os.path.join(self.folder, f"part-{self.n:05d}.jsonl"), "w", encoding="utf-8")
            self.n += 1
            self.count = 0
        self.f.write(json.dumps(doc, ensure_ascii=False) + "\n")
        self.count += 1

    def close(self):
        if self.f:
            self.f.close()
            self.f = None


def read_dir(folder: str):
    if not os.path.isdir(folder):
        return
    for name in sorted(os.listdir(folder)):
        if name.endswith(".jsonl"):
            with open(os.path.join(folder, name), encoding="utf-8") as f:
                for line in f:
                    if line.strip():
                        yield json.loads(line)


class Run:
    def __init__(self, recipe: dict, out: str, procs: int = 0, workers: int | None = None, limit: int = 0, log=print):
        self.r, self.out, self.limit, self.log = recipe, os.path.abspath(out), limit, log
        self.procs = procs or max(1, (os.cpu_count() or 2) - 1)
        self.ep = endpoint_of(recipe, workers)
        os.makedirs(self.out, exist_ok=True)
        self.state_path = os.path.join(self.out, "state.json")
        self.state = json.load(open(self.state_path)) if os.path.exists(self.state_path) else {"done": [], "sources": {}}
        with open(os.path.join(self.out, "recipe.yaml"), "w", encoding="utf-8") as f:
            f.write(dump(recipe))

    def save(self):
        tmp = self.state_path + ".tmp"
        with open(tmp, "w") as f:
            json.dump(self.state, f, indent=1)
        os.replace(tmp, self.state_path)

    def path(self, *p) -> str:
        return os.path.join(self.out, *p)

    def go(self, stages=STAGES):
        if not self.ep.api_key and self.r["endpoint"]["base_url"].startswith("https://router.huggingface.co"):
            raise SystemExit("Set DI_API_KEY (or HF_TOKEN) for the model endpoint.")
        for st in STAGES:
            if st not in stages:
                continue
            if st in self.state["done"]:
                self.log(f"[{st}] already done, skipping (delete {self.state_path} to start over)")
                continue
            t = time.time()
            getattr(self, "stage_" + st)()
            self.state["done"].append(st)
            self.save()
            self.log(f"[{st}] done in {time.time() - t:.0f}s")
        self.write_report()

    # ------------------------------------------------------------ grade
    def stage_grade(self):
        r = self.r
        n_total = int(r["calibrate"]["grade_docs"])
        srcs = sources_of(r)
        per = max(4, n_total // len(srcs))
        cleaner = Cleaner(r["clean"]["rules"])
        pool: list[dict] = []
        for s in srcs:
            got = []
            for d in stream(s):
                d2, why = cleaner.apply(d)
                if d2:
                    got.append(d2)
                if len(got) >= per * 3:
                    break
            pool += random.Random(1).sample(got, min(per, len(got)))
            self.log(f"[grade] {s.name}: sampled {min(per, len(got))} documents")
        gpath = self.path("grades.jsonl")
        done = {}
        if os.path.exists(gpath):
            for line in open(gpath, encoding="utf-8"):
                g = json.loads(line)
                done[g["id"]] = g
        todo = [d for d in pool if d["id"] not in done]
        desc = r["targets"]["description"]
        self.log(f"[grade] grading {len(todo)} documents ({len(done)} already graded) with {self.ep.model}")
        with open(gpath, "a", encoding="utf-8") as f, ThreadPoolExecutor(self.ep.workers) as ex:
            for i, (d, g) in enumerate(zip(todo, ex.map(lambda d: cal.grade(self.ep, d["text"], desc), todo))):
                if g:
                    rec = {"id": d["id"], "source": d["source"], "text": d["text"][:12000], **g}
                    f.write(json.dumps(rec, ensure_ascii=False) + "\n")
                    done[d["id"]] = rec
                if (i + 1) % 200 == 0:
                    f.flush()
                    self.log(f"[grade] {i + 1}/{len(todo)}")
        graded = [(g["text"], g) for g in done.values()]
        target_texts = []
        for t in targets_of(r):
            for i, d in enumerate(stream(t)):
                if i >= int(r["targets"]["max_docs"]):
                    break
                target_texts.append(d["text"])
        calib = cal.Calibrator(float(r["calibrate"]["focus"]), desc, target_texts).fit(graded, [d["text"] for d in pool])
        with open(self.path("calibrator.json"), "w") as f:
            json.dump(calib.to_dict(), f)
        self.state["grade"] = {"graded": len(graded), "agreement": calib.agreement}
        self.log(f"[grade] scorer agreement with the model on held-out documents: {calib.agreement}")

    # ------------------------------------------------------------ curate
    def stage_curate(self):
        r = self.r
        dec_keys = r["clean"]["decontaminate"]
        index = NearDupIndex(MinHasher(), threshold=float(r["calibrate"]["dedup"]["threshold"])) if r["calibrate"]["dedup"]["near"] else None
        exact: set = set()
        ctx = get_context("fork") if hasattr(os, "fork") else get_context("spawn")
        held = {}
        for h in held_out_of(r):
            held[h.name] = [d["text"] for _, d in zip(range(200_000), stream(h))]
            self.log(f"[curate] held-out set {h.name}: {len(held[h.name]):,} items")
        init = (r["clean"]["rules"], dec_keys, held, int(r["clean"].get("ngram", 13)), self.path("calibrator.json"), bool(index))
        with ctx.Pool(self.procs, initializer=_worker_init, initargs=init) as pool:
            for s in sources_of(r):
                st = self.state["sources"].get(s.name, {})
                folder = self.path("curated", s.name)
                if st.get("done"):
                    self.log(f"[curate] {s.name}: done earlier, loading it into the dedup index")
                    for d in read_dir(folder):
                        exact.add(d["h"])
                    if index is not None:
                        for sig, d in zip(pool.imap(_sig_only, (d["text"] for d in read_dir(folder)), chunksize=64), read_dir(folder)):
                            index.check_add_sig(h64(d["h"]), sig)
                    continue
                if os.path.isdir(folder):
                    shutil.rmtree(folder)
                sh = Shards(folder, int(r["output"]["shard_docs"]))
                c = {"read": 0, "kept": 0, "dropped": {}, "decontaminated": {}, "tokens": 0}
                scores: list[float] = []
                t0 = time.time()
                llm_all = r["calibrate"].get("scorer") == "llm"
                calib = cal.Calibrator.from_dict(json.load(open(self.path("calibrator.json")))) if llm_all else None
                ex = ThreadPoolExecutor(self.ep.workers) if llm_all else None
                desc = r["targets"]["description"]
                for batch in pool.imap(_worker_batch, _batches(stream(s), 256, self.limit), chunksize=1):
                    if llm_all:  # grade every surviving document with the model instead of the learned scorer
                        live = [d for d, why, _ in batch if d]
                        for d, g in zip(live, ex.map(lambda d: cal.grade(self.ep, d["text"], desc), live)):
                            if g:
                                d.update(calib.score(d["text"], g))
                    for d, why, sig in batch:
                        c["read"] += 1
                        if why:
                            key = "decontaminated" if why.startswith("bench:") else "dropped"
                            name = why[6:] if key == "decontaminated" else why
                            c[key][name] = c[key].get(name, 0) + 1
                            continue
                        if d["h"] in exact:
                            c["dropped"]["exact duplicate"] = c["dropped"].get("exact duplicate", 0) + 1
                            continue
                        exact.add(d["h"])
                        if index is not None and index.check_add_sig(h64(d["h"]), sig) is not None:
                            c["dropped"]["near duplicate"] = c["dropped"].get("near duplicate", 0) + 1
                            continue
                        sh.write(d)
                        c["kept"] += 1
                        c["tokens"] += d["tokens"]
                        scores.append(d["score"])
                    if c["read"] % 50_000 < 256:
                        rate = c["read"] / max(1e-6, time.time() - t0)
                        self.log(f"[curate] {s.name}: {c['read']:,} read, {c['kept']:,} kept, {rate:,.0f} docs/s")
                sh.close()
                if ex:
                    ex.shutdown()
                scores.sort(reverse=True)
                keep_top, n = float(r["calibrate"]["keep_top"]), len(scores)
                c["cut"] = scores[max(0, int(n * keep_top) - 1)] if n else 1.0
                c["high_cut"] = scores[max(0, int(n * 0.2) - 1)] if n else 1.0
                c["selected_tokens"], c["high_tokens"] = 0, 0
                for d in read_dir(folder):
                    if d["score"] >= c["cut"]:
                        c["selected_tokens"] += d["tokens"]
                    if d["score"] >= c["high_cut"]:
                        c["high_tokens"] += d["tokens"]
                c["done"] = True
                self.state["sources"][s.name] = c
                self.save()
                self.log(f"[curate] {s.name}: kept {c['kept']:,} of {c['read']:,}; {fmt(c['selected_tokens'])} tokens above the score cut")

    # ------------------------------------------------------------ create
    def plans(self, synthetic: dict | None = None) -> list[dict]:
        rows = [{"name": s.name, "kind": s.kind, "available_tokens": self.state["sources"].get(s.name, {}).get("selected_tokens", 0),
                 "high_tokens": self.state["sources"].get(s.name, {}).get("high_tokens", 0)} for s in sources_of(self.r)]
        defaults = default_weights(rows)
        out = []
        for m in self.r["compose"]["mixes"]:
            w = {**defaults, **{k: float(v) for k, v in (m.get("weights") or {}).items()}}
            inv = float(m.get("invented_share", 0))
            p = plan_mix(int(int(m["budget_tokens"]) * (1 - inv)), rows, w, float(self.r["compose"]["max_epochs"]), float(self.r["compose"]["max_synthetic_share"]))
            p["name"], p["weights"], p["invented_share"], p["budget_tokens"] = m["name"], w, inv, int(m["budget_tokens"])
            out.append(p)
        return out

    def stage_create(self):
        r = self.r
        rc, ic = r["create"]["rephrase"], r["create"]["invent"]
        need: dict[str, int] = {}
        for p in self.plans():
            for row in p["rows"]:
                need[row["name"]] = max(need.get(row["name"], 0), row["synthetic_tokens"])
        made = self.state.setdefault("synthetic", {})
        rng = random.Random(3)
        if rc["enabled"]:
            for s in sources_of(r):
                want = need.get(s.name, 0)
                have = made.get(s.name, {}).get("tokens", 0)
                if want <= have:
                    continue
                self.log(f"[create] {s.name}: rephrasing to add {fmt(want - have)} tokens")
                self._rephrase_source(s, want, rc, rng)
        if ic["enabled"] and int(ic["docs"]) > 0 and ic["request"].strip():
            self._invent(ic, rng)

    def _rephrase_source(self, s: Source, want: int, rc: dict, rng: random.Random):
        folder = self.path("synthetic", s.name)
        st = self.state["synthetic"].setdefault(s.name, {"tokens": 0, "docs": 0, "flagged": 0, "seen": 0})
        cut = self.state["sources"][s.name]["high_cut"]
        styles = rc["styles"] or styles_for(s.kind)

        def gen():
            # each pass over the best documents uses a different style, so later passes add variety, not copies
            k = 0
            for pass_ in range(len(styles)):
                for d in read_dir(self.path("curated", s.name)):
                    if d["score"] >= cut:
                        for i, ch in enumerate(chunks(d["text"])):
                            yield d["id"], i, ch, styles[(pass_ + k) % len(styles)]
                        k += 1
        jobs = itertools.islice(gen(), st["seen"], None)
        sh = Shards(folder, int(self.r["output"]["shard_docs"]))
        sh.n = len([f for f in os.listdir(folder)]) if os.path.isdir(folder) else 0
        step = self.ep.workers * 4
        with ThreadPoolExecutor(self.ep.workers) as ex:
            while st["tokens"] < want:
                part = list(itertools.islice(jobs, step))
                if not part:
                    break
                for (oid, i, ch, style), o in zip(part, ex.map(lambda j: rephrase(self.ep, j[2], j[3], rc.get("language", ""), float(rc["max_new_share"])), part)):
                    st["seen"] += 1
                    if not o["text"]:
                        continue
                    if not o["ok"]:
                        st["flagged"] += 1
                        continue
                    t = est_tokens(o["text"])
                    sh.write({"id": f"{oid}:{i}:{style}", "text": o["text"], "source": s.name, "kind": "synthetic", "origin": oid,
                              "style": style, "tokens": t, "fidelity": o["fidelity"]})
                    st["tokens"] += t
                    st["docs"] += 1
                sh.close()
                self.save()
                self.log(f"[create] {s.name}: {fmt(st['tokens'])} of {fmt(want)} tokens, {st['flagged']} flagged")
        if st["tokens"] < want:
            self.log(f"[create] {s.name}: ran out of documents to rephrase at {fmt(st['tokens'])} tokens")

    def _invent(self, ic: dict, rng: random.Random):
        st = self.state["synthetic"].setdefault("invented", {"tokens": 0, "docs": 0, "flagged": 0, "seen": 0})
        n = int(ic["docs"]) - st["docs"]
        if n <= 0:
            return
        topics = plan_topics(self.ep, ic["request"], min(200, max(10, n // 5)))
        jobs = [(rng.choice(topics), rng.choice(ic["doc_types"]), rng.choice(ic["levels"])) for _ in range(n)]
        sh = Shards(self.path("synthetic", "invented"), int(self.r["output"]["shard_docs"]))
        sh.n = len(os.listdir(self.path("synthetic", "invented")))
        self.log(f"[create] inventing {n} documents over {len(topics)} topics")
        with ThreadPoolExecutor(self.ep.workers) as ex:
            for (t, dt, lv), o in zip(jobs, ex.map(lambda j: invent(self.ep, j[0], j[1], j[2], ic["request"], ic.get("language", "English")), jobs)):
                if o.get("text"):
                    tk = est_tokens(o["text"])
                    sh.write({"id": f"invented:{st['docs']}", "text": o["text"], "source": "invented", "kind": "synthetic",
                              "topic": t, "doc_type": dt, "level": lv, "tokens": tk})
                    st["docs"] += 1
                    st["tokens"] += tk
        sh.close()

    # ------------------------------------------------------------ compose
    def stage_compose(self):
        r = self.r
        synth = self.state.get("synthetic", {})
        results = []
        for p in self.plans():
            name = p["name"]
            folder = self.path("mixes", name)
            if os.path.isdir(folder):
                shutil.rmtree(folder)
            sh = Shards(folder, int(r["output"]["shard_docs"]))
            buckets = []
            for row in p["rows"]:
                s = row["name"]
                cut = self.state["sources"].get(s, {}).get("cut", 1.0)
                real_target = row["want_tokens"] - min(row["synthetic_tokens"], synth.get(s, {}).get("tokens", 0))
                buckets.append(_Bucket(s, self.path("curated", s), cut, real_target, float(r["compose"]["max_epochs"]),
                                       self.state["sources"].get(s, {}).get("selected_tokens", 0)))
                if row["synthetic_tokens"] and synth.get(s, {}).get("tokens"):
                    buckets.append(_Bucket(s + ":synthetic", self.path("synthetic", s), -1, min(row["synthetic_tokens"], synth[s]["tokens"]), 1.0,
                                           synth[s]["tokens"]))
            share = float(p.get("invented_share", 0))
            if share and synth.get("invented", {}).get("tokens"):
                buckets.append(_Bucket("invented", self.path("synthetic", "invented"), -1, int(p["budget_tokens"] * share), 1.0,
                                       synth["invented"]["tokens"]))
            rng = random.Random(11)
            buf: list[dict] = []
            live = [b for b in buckets if b.target > 0]
            while live:
                # draw the next bucket in proportion to the tokens it still owes, which interleaves sources evenly
                b = rng.choices(live, weights=[x.target - x.taken for x in live], k=1)[0]
                d = b.next()
                if d is None:
                    live.remove(b)
                    continue
                buf.append({"id": d["id"], "text": d["text"], "source": d["source"], "kind": d.get("kind", "")})
                if len(buf) >= 20_000:
                    rng.shuffle(buf)
                    for x in buf[:10_000]:
                        sh.write(x)
                    buf = buf[10_000:]
                if b.taken >= b.target:
                    live.remove(b)
            rng.shuffle(buf)
            for x in buf:
                sh.write(x)
            sh.close()
            manifest = {"name": name, "budget_tokens": p["budget_tokens"], "weights": p["weights"],
                        "buckets": [{"name": b.name, "target_tokens": b.target, "tokens": b.taken, "docs": b.docs, "epochs": b.epochs_used()} for b in buckets],
                        "tokens": sum(b.taken for b in buckets), "warnings": p["warnings"]}
            with open(os.path.join(folder, "manifest.json"), "w") as f:
                json.dump(manifest, f, indent=1)
            results.append(manifest)
            self.log(f"[compose] {name}: {fmt(manifest['tokens'])} tokens of {fmt(p['budget_tokens'])}")
        self.state["mixes"] = results

    # ------------------------------------------------------------ report
    def write_report(self):
        st = self.state
        lines = [f"# Curation report: {self.r['name']}", "", f"Model: `{self.ep.model}`. Stages done: {', '.join(st['done']) or 'none'}.", ""]
        if st.get("grade"):
            a = st["grade"]["agreement"]
            lines += [f"Graded {st['grade']['graded']:,} documents. Scorer rank agreement with the model on held-out documents: "
                      f"quality {a.get('quality')}, relevance {a.get('relevance')}.", ""]
        if st["sources"]:
            lines += ["## Sources", "", "| Source | Read | Kept | Decontaminated | Top drop reasons | Tokens above cut |", "|---|---|---|---|---|---|"]
            for n, c in st["sources"].items():
                top = ", ".join(f"{k} {v:,}" for k, v in sorted(c["dropped"].items(), key=lambda x: -x[1])[:3])
                dec = sum(c["decontaminated"].values())
                lines.append(f"| {n} | {c['read']:,} | {c['kept']:,} | {dec:,} | {top} | {fmt(c.get('selected_tokens', 0))} |")
            lines.append("")
        if st.get("synthetic"):
            lines += ["## Created", "", "| From | Documents | Tokens | Flagged |", "|---|---|---|---|"]
            for n, c in st["synthetic"].items():
                lines.append(f"| {n} | {c['docs']:,} | {fmt(c['tokens'])} | {c['flagged']:,} |")
            lines.append("")
        for m in st.get("mixes", []):
            lines += [f"## Mix `{m['name']}`: {fmt(m['tokens'])} of {fmt(m['budget_tokens'])} tokens", "",
                      "| Bucket | Target | Written | Epochs |", "|---|---|---|---|"]
            for b in m["buckets"]:
                lines.append(f"| {b['name']} | {fmt(b['target_tokens'])} | {fmt(b['tokens'])} | {b['epochs']} |")
            lines += [""] + [f"- {w}" for w in m["warnings"]] + [""]
        with open(self.path("REPORT.md"), "w") as f:
            f.write("\n".join(lines))
        with open(self.path("report.json"), "w") as f:
            json.dump(st, f, indent=1)


class _Bucket:
    """Documents of one source above its score cut, cycled up to max_epochs times."""

    def __init__(self, name: str, folder: str, cut: float, target: int, max_epochs: float, pass_tokens: int = 0):
        self.name, self.folder, self.cut, self.target, self.max_epochs = name, folder, cut, int(target), max_epochs
        self.pass_tokens = pass_tokens
        self.taken, self.docs, self.epoch, self.it, self.epoch_tokens = 0, 0, 0, None, 0

    def next(self):
        while True:
            if self.it is None:
                if self.epoch >= max(1, int(self.max_epochs + 0.999)):
                    return None
                self.it = iter(read_dir(self.folder))
                self.epoch += 1
            for d in self.it:
                if d.get("score", 1.0) >= self.cut:
                    self.taken += d.get("tokens") or est_tokens(d["text"])
                    self.docs += 1
                    return d
            self.it = None
            if self.taken == 0:
                return None

    def epochs_used(self) -> float:
        return round(self.taken / self.pass_tokens, 2) if self.pass_tokens else float(self.epoch)


# ---------------------------------------------------------------- worker processes

_W: dict = {}


def _worker_init(rules, dec_keys, held, ngram, calib_path, near):
    _W["clean"] = Cleaner(rules)
    dec = Decontaminator(ngram)
    for k in dec_keys:
        dec.add(k, benchmark_texts(k))
    for k, texts in held.items():
        dec.add(k, texts)
    _W["dec"] = dec
    _W["calib"] = cal.Calibrator.from_dict(json.load(open(calib_path)))
    _W["mh"] = MinHasher() if near else None


def _worker_batch(batch: list[dict]) -> list[tuple]:
    out = []
    for d in batch:
        d2, why = _W["clean"].apply(d)
        if not d2:
            out.append((None, why, None))
            continue
        b = _W["dec"].hit(d2["text"])
        if b:
            out.append((None, "bench:" + b, None))
            continue
        sc = _W["calib"].score(d2["text"])
        rec = {**d2, **sc, "tokens": est_tokens(d2["text"]), "h": doc_hash(d2["text"])}
        out.append((rec, "", _W["mh"].signature(d2["text"]) if _W["mh"] else None))
    return out


def _sig_only(text: str):
    return _W["mh"].signature(text) if _W.get("mh") else None


def _batches(it, n: int, limit: int = 0):
    buf, seen = [], 0
    for d in it:
        buf.append(d)
        seen += 1
        if len(buf) >= n:
            yield buf
            buf = []
        if limit and seen >= limit:
            break
    if buf:
        yield buf