"""Training mode of the task generator (stage 2 data).

Separate pool `tasks_gen/train` (ids G1000+). Same category and object type mix as the eval plan, plus
tasks for the common errors (docs/yol-haritasi.md step 3). A new bundle is checked against all eval tasks
and the earlier training tasks (`harness/overlap.py`); a too close bundle goes back to the model as a
repair message. K variants are not made here (they need the generic_v0 tool schema).

  python3 -m harness.trainset plan
  python3 -m harness.trainset run --part 0 --parts 3 [--target 200] [--stop-ledger 27]
"""
import argparse
import glob
import json
import os
import time

from .adt_client import load_env
from .evalset import SLOTS, RELEASES, accepted_goals
from .generator import ROOT, generate, make_k_variant
from .ledger import BudgetExceeded, spent
from . import overlap
from . import mix

POOL = os.path.join(ROOT, "tasks_gen", "train")
PLAN = os.path.join(POOL, "plan.json")
FIRST_ID = 1000
RUN_BASE = 32000          # run numbers stay below 36**3 * 1 (proxy prefix rule); 40 per slot
N_SLOTS = 224             # about 12 % more slots than the target: some tasks are not accepted
N_ERROR = 30
EVAL_HOLD = 3             # eval pool is loaded once per process

ERROR_HINTS = [
    ("named-type", "The contract has methods with short coded character parameters (zone code, currency code, "
     "status flag, region key) of fixed length. A correct solution declares named types for them: "
     "TYPE c LENGTH n is not allowed in a method signature."),
    ("reserved-word", "The business domain uses concepts that collide with reserved or SQL words (hours, mode, "
     "order, date, count, group, value). The contract uses safe field and parameter names (for example "
     "work_hours, calc_mode); the solution must also avoid reserved words as field names."),
    ("long-names", "The task has at least 10 separate business rules, so a good solution needs at least 10 "
     "own unit tests, each for one rule. Every ABAP name stays at 30 characters or less (the model must "
     "shorten long descriptive test method names)."),
]
ERROR_TYPES = ["CLAS", "CLAS", "FUNC", "DDLS", "PROG", "CLAS"]   # reserved-word fits DDLS and tables well


CATEGORY_SHARE = {"A": 10, "B": 15, "C": 15, "D": 10, "E": 15, "F": 10, "G": 10, "I": 5, "H": 10}  # eval 8.1
ERROR_CATEGORY = {"named-type": "A", "reserved-word": "D", "long-names": "C"}


def build_plan(n_slots=N_SLOTS, n_error=N_ERROR, hard_share=0.0, first_id=FIRST_ID, run_base=RUN_BASE):
    """Category shares as in the eval plan; object types inside a category as in evalset.SLOTS.
    n_error slots target the common errors (named-type, reserved-word, long-names); hard_share of the other
    slots get difficulty 3 (every 10th slot pattern, deterministic)."""
    share_sum = sum(CATEGORY_SHARE.values())
    items = []
    for cat, share in CATEGORY_SHARE.items():
        k_cat = round(n_slots * share / share_sum)
        per_kind = n_error // 3
        k_err = sum(1 for k in ERROR_CATEGORY if ERROR_CATEGORY[k] == cat) * per_kind
        slots = [x for x in SLOTS if x[0] == cat]
        weight = sum(n for _, _, n, _ in slots)
        k_rest = k_cat - k_err
        quota = [(t, hints, round(k_rest * n / weight)) for _, t, n, hints in slots]
        quota[0] = (quota[0][0], quota[0][1], quota[0][2] + k_rest - sum(q[2] for q in quota))
        pos = 0
        for t, hints, k in quota:
            for i in range(k):
                topic = hints[i % len(hints)] if hints else None
                if cat == "G":
                    topic = (topic + "; " if topic else "") + f"release target {RELEASES[pos % 2]}"
                pos += 1
                items.append({"category": cat, "object_type": t, "topic": topic, "frac": pos / (k_cat + 1)})
        for kind, text in ERROR_HINTS:
            if ERROR_CATEGORY[kind] != cat:
                continue
            for i in range(per_kind):
                t = ERROR_TYPES[i % len(ERROR_TYPES)]
                pos += 1
                items.append({"category": cat, "object_type": t, "topic": text, "error_kind": kind,
                              "frac": pos / (k_cat + 1)})
    items.sort(key=lambda x: (x["frac"], x["category"], x["object_type"]))   # interleave: a cut keeps the mix
    hard_every = round(1 / hard_share) if hard_share else 0
    for n, it in enumerate(items):
        it.pop("frac")
        it["id"] = f"G{first_id + n:04d}"
        it["run_base"] = run_base + 40 * n
        if hard_every and it["category"] != "H" and n % hard_every == hard_every // 2:
            it["difficulty"] = 3
    return items


def accepted_count():
    n = 0
    for f in glob.glob(os.path.join(POOL, "G*", "generation.json")):
        try:
            n += bool(json.load(open(f)).get("accepted"))
        except (OSError, ValueError):
            pass
    return n


STOP_FLAG = os.path.join(ROOT, "runs", "pipeline", "STOP")   # the pipeline writes it: workers end after the current slot
BACKLOG_LIMIT = 40        # accepted tasks without a first trajectory run: generation waits above this
PLAN2 = {"n_slots": 500, "n_error": 100, "hard_share": 0.3, "first_id": 1400, "run_base": 100000}
# plan 2 (2026-10-05, Opus 5.5 rules): error tasks 20 % (plan 1: 13 %, +50 %), hard (difficulty 3) 30 % of the
# other slots (plan 1: 20 % assumed from the pilot hard batch), same category shares, ids G1400+, no fixed target


def plan_path(name):
    return os.path.join(POOL, name + ".json")


def ensure_plan(name, stop_ledger=27.0, target=200):
    os.makedirs(POOL, exist_ok=True)
    if not os.path.exists(plan_path(name)):
        if name == "plan":
            slots, tgt, stop = build_plan(), target, stop_ledger
        else:
            slots, tgt, stop = build_plan(**PLAN2), None, 1e6   # only the budget guard and the deadline stop it
        json.dump({"slots": slots, "ledger_at_start": spent(), "stop_ledger": stop, "target": tgt,
                   "created": time.strftime("%F %T")}, open(plan_path(name), "w"), indent=1)
    return json.load(open(plan_path(name)))


def backlog():
    """Accepted training tasks that have no first trajectory run yet."""
    done = set()
    sp = os.path.join(ROOT, "runs", "traj", "summary.jsonl")
    if os.path.exists(sp):
        done = {json.loads(l)["task"] for l in open(sp) if json.loads(l)["attempt"] == 0}
    acc = {os.path.basename(os.path.dirname(f)) for f in glob.glob(os.path.join(POOL, "G*", "generation.json"))
           if json.load(open(f)).get("accepted")}
    return len(acc - done)


def backlog_by_kind():
    """{kind: accepted tasks without a first trajectory run}."""
    done = set()
    sp = os.path.join(ROOT, "runs", "traj", "summary.jsonl")
    if os.path.exists(sp):
        done = {json.loads(l)["task"] for l in open(sp) if json.loads(l)["attempt"] == 0}
    out = {}
    for f in glob.glob(os.path.join(POOL, "G*", "generation.json")):
        tid = os.path.basename(os.path.dirname(f))
        try:
            ok = json.load(open(f)).get("accepted")
        except (OSError, ValueError):
            ok = False
        if ok and tid not in done:
            k = mix.kind_of_task_dir(tid)
            out[k] = out.get(k, 0) + 1
    return out


BAL_KIND_BACKLOG = 8      # a kind with more waiting tasks than this is not generated (its trajectories come first)


def run(part, parts, target, stop_ledger, plan_name="plan", deadline=None):
    plan = ensure_plan(plan_name, stop_ledger, target)
    stop_at = plan["ledger_at_start"] + plan["stop_ledger"]
    base_url = os.environ.get("LLM_BASE_URL", "http://127.0.0.1:11434/v1")
    evals = overlap.load_pool("eval")
    for idx, s in enumerate(plan["slots"]):
        if idx % parts != part:
            continue
        if os.path.exists(os.path.join(POOL, "_logs", s["id"] + ".json")):
            continue  # done or failed: not tried again (delete the log to retry)
        while backlog() > BACKLOG_LIMIT and not os.path.exists(STOP_FLAG) \
                and not (deadline and time.time() > deadline):
            time.sleep(120)   # trajectories are the slower side: do not generate tasks that wait for days
        if os.path.exists(STOP_FLAG):
            print("STOP flag", flush=True)
            break
        if deadline and time.time() > deadline:
            print("DEADLINE", flush=True)
            break
        if plan["target"] and accepted_count() >= plan["target"]:
            print("TARGET reached", flush=True)
            break
        if spent() >= stop_at:
            print("PHASE LIMIT", spent(), stop_at, flush=True)
            break
        goals = accepted_goals()
        avoid = [g for g in goals if g]   # eval goals and all training goals
        topic = ((s["topic"] + ". ") if s["topic"] else "Choose a new, realistic business topic. ") + \
            "Do not repeat these existing topics: " + "; ".join(avoid[-170:])
        pool_now = evals + overlap.load_pool("train")

        def extra(b, _pool=pool_now):
            hits = overlap.check(overlap.load_bundle(b), _pool)
            return [f"Too close to task {i} (similarity spec {sc['spec']:.2f}, rules {sc['core']:.2f}, "
                    f"names {sc['name']:.2f}). Choose a different business topic and different object names."
                    for i, sc in hits[:3]]
        try:
            log = generate(s["id"], "train", s["object_type"], s["category"], s.get("difficulty", 2),
                           "deepseek-v4.1-flash:cloud",
                           base_url, s["run_base"], topic, extra_check=extra)
        except BudgetExceeded as e:
            print("BUDGET", e, flush=True)
            break
        except Exception as e:  # noqa: BLE001
            log = {"id": s["id"], "error": str(e)[:500]}
        log.update(error_kind=s.get("error_kind"), difficulty=s.get("difficulty", 2), spent_total=spent())
        os.makedirs(os.path.join(POOL, "_logs"), exist_ok=True)
        json.dump(log, open(os.path.join(POOL, "_logs", s["id"] + ".json"), "w"), indent=1)
        stray = os.path.join(POOL, "generation.json")  # generate() writes here when no task dir exists
        if os.path.exists(stray):
            os.remove(stray)
        print(json.dumps(log), flush=True)


BAL_FIRST_ID = 1910
BAL_RUN_BASE = 370000     # 40 per slot; below 466560 (a digit must lead the 4-char base36 run)
BAL_ERROR_KINDS = {"CLAS": ["named-type", "long-names"], "FUNC": ["named-type"], "DDLS": ["reserved-word"],
                   "TABL": ["reserved-word"], "STRU": ["reserved-word"]}
BAL_HINTS = dict(((c, t), h) for c, t, _, h in SLOTS if h)


def _claims_dir():
    d = os.path.join(POOL, "_claims")
    os.makedirs(d, exist_ok=True)
    return d


def _claim_slot():
    """Next free balanced slot number, claimed with O_EXCL (several workers). Returns (n, claim path)."""
    for n in range(BAL_FIRST_ID, BAL_FIRST_ID + 600):
        sid = "G%04d" % n
        if os.path.exists(os.path.join(POOL, "_logs", sid + ".json")):
            continue
        path = os.path.join(_claims_dir(), sid + ".json")
        try:
            fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
        except FileExistsError:
            continue
        os.close(fd)
        return n, path
    return None, None


def _kind_stats():
    """({kind: accepted}, {kind: attempted}) from the generation logs; claims of running slots count as attempted."""
    acc = mix.accepted_task_counts()
    att = {}
    for f in glob.glob(os.path.join(POOL, "_logs", "G*.json")):
        try:
            l = json.load(open(f))
        except (OSError, ValueError):
            continue
        k = l.get("kind") or l.get("object_type")
        if k and l.get("category") != "K":
            att[k] = att.get(k, 0) + 1
    return acc, att


def run_balanced(part, parts, deadline):
    """Generation without a fixed plan: each slot takes the kind with the biggest deficit against mix.TYPE_SHARE.
    A kind with 6 or more tries and an acceptance below 20 % is skipped (a harness or prompt problem: do not burn budget)."""
    base_url = os.environ.get("LLM_BASE_URL", "http://127.0.0.1:11434/v1")
    evals = overlap.load_pool("eval")
    while True:
        if os.path.exists(STOP_FLAG):
            print("STOP flag", flush=True)
            return
        if deadline and time.time() > deadline:
            print("DEADLINE", flush=True)
            return
        acc, att = _kind_stats()
        running = {}
        for f in glob.glob(os.path.join(_claims_dir(), "G*.json")):
            try:
                k = json.load(open(f)).get("kind")
            except (OSError, ValueError):
                k = None
            if k:
                running[k] = running.get(k, 0) + 1
        counts = {k: acc.get(k, 0) + running.get(k, 0) for k in set(acc) | set(running) | set(mix.TYPE_SHARE)}
        blocked = {k for k in mix.TYPE_SHARE if att.get(k, 0) >= 6 and acc.get(k, 0) < 0.2 * att.get(k, 0)}
        if blocked:
            print("kinds skipped (low acceptance):", sorted(blocked), flush=True)
        waiting = {k for k, n in backlog_by_kind().items() if n > BAL_KIND_BACKLOG}
        allowed = set(mix.TYPE_SHARE) - blocked - waiting
        if not allowed:
            time.sleep(120)     # every kind has a backlog: the trajectories are the slower side
            continue
        kind = mix.deficit_pick(counts, allowed=allowed)
        n, claim = _claim_slot()
        if n is None:
            print("no free slot", flush=True)
            return
        sid = "G%04d" % n
        json.dump({"kind": kind}, open(claim, "w"))
        otype = "CLAS" if kind == "EXC" else kind
        cats = mix.KIND_CATEGORIES[kind]
        logs = [json.load(open(f)) for f in glob.glob(os.path.join(POOL, "_logs", "G*.json"))]
        ccount = {c: sum(1 for l in logs if l.get("accepted") and l.get("category") == c) for c in cats}
        tot = sum(ccount.values()) + 1
        cat = max(cats, key=lambda c: tot * mix.CATEGORY_SHARE[c] / sum(mix.CATEGORY_SHARE[x] for x in cats) - ccount[c])
        idx = n - BAL_FIRST_ID
        error_kind = None
        if idx % 5 == 2 and kind in BAL_ERROR_KINDS:   # 20 % error-targeted slots
            error_kind = BAL_ERROR_KINDS[kind][(idx // 5) % len(BAL_ERROR_KINDS[kind])]
            cat = ERROR_CATEGORY[error_kind] if kind in ("CLAS", "FUNC") else cat
        hard = cat != "H" and idx % 10 in (3, 6, 9)       # 30 % hard
        topic = None
        if error_kind:
            topic = dict(ERROR_HINTS)[error_kind]
        elif kind == "EXC":
            topic = "exception class (CX_...): " + ["a domain exception with context attributes and message texts",
                                                    "an exception hierarchy with a common super class",
                                                    "an exception that wraps a previous exception"][idx % 3]
        elif (cat, otype) in BAL_HINTS:
            h = BAL_HINTS[(cat, otype)]
            topic = h[idx % len(h)]
        if cat == "G":
            topic = (topic + "; " if topic else "") + f"release target {RELEASES[idx % 2]}"
        avoid = [g for g in accepted_goals() if g][-170:]
        full = ((topic + ". ") if topic else "Choose a new, realistic business topic. ") + \
            "Do not repeat these existing topics: " + "; ".join(avoid)
        pool_now = evals + overlap.load_pool("train")

        def extra(b, _pool=pool_now):
            hits = overlap.check(overlap.load_bundle(b), _pool)
            return [f"Too close to task {i} (similarity spec {sc['spec']:.2f}, rules {sc['core']:.2f}, "
                    f"names {sc['name']:.2f}). Choose a different business topic and different object names."
                    for i, sc in hits[:3]]
        try:
            log = generate(sid, "train", otype, cat, 3 if hard else 2, "deepseek-v4.1-flash:cloud", base_url,
                           BAL_RUN_BASE + 40 * idx, full, extra_check=extra)
        except BudgetExceeded as e:
            print("BUDGET", e, flush=True)
            os.remove(claim)
            return
        except Exception as e:  # noqa: BLE001
            log = {"id": sid, "error": str(e)[:500]}
        log.update(kind=kind, error_kind=error_kind, difficulty=3 if hard else 2, spent_total=spent())
        os.makedirs(os.path.join(POOL, "_logs"), exist_ok=True)
        json.dump(log, open(os.path.join(POOL, "_logs", sid + ".json"), "w"), indent=1)
        stray = os.path.join(POOL, "generation.json")
        if os.path.exists(stray):
            os.remove(stray)
        os.remove(claim)
        print(json.dumps(log), flush=True)


K_FIRST_ID = 1300
K_RUN_BASE = 41800        # 20 per variant; above the trajectory run numbers (41000-41700)
K_COUNT = 18              # K share of the eval plan: 10 of 110 (9 %); counted inside the 200 accepted tasks
K_STYLES_CYCLE = ["free_text", "incomplete"]


def run_k(count):
    """K tasks for training: free-text or incomplete spec of an accepted training task, EPOD tool names
    (no generic_v0). Same reference and hidden tests as the source task."""
    base_url = os.environ.get("LLM_BASE_URL", "http://127.0.0.1:11434/v1")
    for n in range(count):
        new_id = f"G{K_FIRST_ID + n}"
        if os.path.exists(os.path.join(POOL, "_logs", new_id + ".json")):
            continue
        used = {json.load(open(f)).get("base_task") for f in glob.glob(os.path.join(POOL, "_logs", "G13*.json"))}
        cands = []   # accepted, not K, not H (a stop task has no free-text form), not used yet
        for f in sorted(glob.glob(os.path.join(POOL, "_logs", "G1*.json"))):
            lg = json.load(open(f))
            if lg.get("accepted") and lg.get("category") not in ("H", "K") and lg["id"] not in used:
                cands.append(lg)
        if not cands:
            print("no source task left", flush=True)
            return
        kinds = {}
        for lg in cands:      # spread over object types: take the type with the fewest K variants so far
            kinds.setdefault(lg["object_type"], []).append(lg)
        done_types = [json.load(open(f)).get("object_type") for f in glob.glob(os.path.join(POOL, "_logs", "G13*.json"))]
        otype = min(kinds, key=lambda t: done_types.count(t))
        src = kinds[otype][0]
        style = K_STYLES_CYCLE[n % 2]
        if os.path.exists(STOP_FLAG):
            return
        try:
            log = make_k_variant(src["id"], new_id, style, "deepseek-v4.1-flash:cloud", base_url,
                                 K_RUN_BASE + 20 * n, tool_schema=None, pool="train")
        except BudgetExceeded as e:
            print("BUDGET", e, flush=True)
            return
        log.update(object_type=otype, spent_total=spent())
        stray = os.path.join(POOL, "generation.json")
        if os.path.exists(stray):
            os.remove(stray)
        os.makedirs(os.path.join(POOL, "_logs"), exist_ok=True)
        json.dump(log, open(os.path.join(POOL, "_logs", new_id + ".json"), "w"), indent=1)
        print(json.dumps(log), flush=True)


def main():
    load_env(os.path.join(ROOT, ".env"))
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["plan", "run", "k"])
    ap.add_argument("--part", type=int, default=0)
    ap.add_argument("--parts", type=int, default=1)
    ap.add_argument("--target", type=int, default=200)
    ap.add_argument("--plan", default="plan", help="plan (first 223 slots), plan2 (hard and error share raised) or balanced (kind with the biggest deficit)")
    ap.add_argument("--deadline", help="YYYY-MM-DDTHH:MM local time: no new slot after it")
    ap.add_argument("--k-count", type=int, default=K_COUNT)
    ap.add_argument("--stop-ledger", type=float, default=27.0, help="ledger USD for this phase (10 USD usage = 27)")
    a = ap.parse_args()
    if a.cmd == "plan":
        import collections
        p = build_plan()
        print(len(p), "slots", collections.Counter(x["category"] for x in p))
        print(collections.Counter(x["object_type"] for x in p), collections.Counter(x.get("error_kind") for x in p))
        return
    if a.cmd == "k":
        run_k(a.k_count)
        return
    dl = time.mktime(time.strptime(a.deadline, "%Y-%m-%dT%H:%M")) if a.deadline else None
    if a.plan == "balanced":
        run_balanced(a.part, a.parts, dl)
        return
    run(a.part, a.parts, a.target, a.stop_ledger, a.plan, dl)


if __name__ == "__main__":
    main()