#!/usr/bin/env python3
"""GUTS wall/matcap texture pack — MODELBEAST flux_local, on-device, $0 (Lane D).

Ported from PROCITY's gen_skins.py (same operator, same discipline: dry-run first, idempotent,
resumable, review-then-harvest). Differences from PROCITY: square 1024s, grayscale-authored
(the biome tint lives in Lane A's shader — ART_BIBLE), and every result must TILE, so harvest
runs through derive_maps.py (seam-blend + Sobel normals + WebP).

  <mflux-py> pipeline/gen_textures.py --dry-run    # list the pack, no GPU, no spend
  <mflux-py> pipeline/gen_textures.py --local      # generate -> pipeline/.genraw/*.png
  <mflux-py> pipeline/gen_textures.py --local --only esophagus
  <mflux-py> pipeline/gen_textures.py --harvest    # .genraw -> web/assets/gen/*.webp (+ _n)

<mflux-py> = ~/Documents/MODELBEAST/venvs/mflux/bin/python   (numpy + PIL live there)

Review .genraw/ by eye and DELETE REJECTS before --harvest — committed junk is forever.
Provenance (prompt/seed/model/time per asset) is written to pipeline/batch_textures.json and
committed: regeneration must always be possible.
"""
import json, os, subprocess, sys, time
import shutil as _sh

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RAW = os.path.join(ROOT, "pipeline", ".genraw")          # git-ignored scratch
BATCH = os.path.join(ROOT, "pipeline", "batch_textures.json")
os.makedirs(RAW, exist_ok=True)

MB = os.environ.get("MB_HOME", os.path.expanduser("~/Documents/MODELBEAST"))
FLUX_PY = os.path.join(MB, "venvs/mflux/bin/python")
FLUX_RUN = os.path.join(MB, "server/operators/flux_local/run.py")

# ── the style stem — verbatim from ART_BIBLE §FLUX prompt kit ────────────────────────────
STEM = ("scanning electron microscope micrograph, monochrome, high detail organic tissue, "
        "dark background, dramatic rim lighting, seamless tiling texture")

# ── the pack. tile = [repeats_around_theta, units_of_s_per_repeat] — see NOTES §tiling ────
# Two per biome (ART_BIBLE: "generate ~2 per biome, not ten"). One texture can serve two
# biomes at different tints, so these are shapes-of-tissue, not colours.
WALLS = {
    "wall_esophagus_a": dict(
        p="human esophageal mucosa, longitudinal ribbed folds, wet ridges",
        seed=1101, tile=[4, 16]),
    "wall_esophagus_b": dict(
        p="human esophagus inner lining, tight circular muscular rings, contracted sphincter folds",
        seed=1102, tile=[3, 12]),
    "wall_stomach_a": dict(
        p="gastric mucosa, pitted craters, gastric pits, glistening mucus",
        seed=1201, tile=[4, 18]),
    "wall_stomach_b": dict(
        p="stomach rugae, thick heavy folded ridges, deep valleys between folds",
        seed=1202, tile=[3, 22]),
    "wall_smallint_a": dict(
        p="intestinal villi, dense finger-like projections, forest of fronds",
        seed=1301, tile=[5, 14]),
    "wall_smallint_b": dict(
        p="intestinal crypts of Lieberkuhn, honeycomb of round pits, packed glandular openings",
        seed=1302, tile=[6, 10]),
}
# matcap: NOT tileable, NOT seam-blended — a lit sphere cropped to a square (see derive_maps)
MATCAPS = {
    "matcap_tissue_wet": dict(
        p=("spherical material study ball, wet translucent organic tissue, subsurface glow, "
           "studio black background"),
        seed=1401, stem_override=True),
}
PROMPTS = {**{k: v for k, v in WALLS.items()}, **MATCAPS}

STEPS = 4          # flux2-klein-4b is a 4-step distilled model
GUIDANCE = 3.5
SIZE = 1024


def local_available():
    return os.path.exists(FLUX_PY) and os.path.exists(FLUX_RUN)


def prompt_of(slug):
    spec = PROMPTS[slug]
    if spec.get("stem_override"):          # matcaps want a lit sphere, not an SEM micrograph
        return spec["p"]
    return STEM + ", " + spec["p"]


def have(slug):
    return os.path.exists(os.path.join(RAW, f"{slug}.png"))


def gen_local(slug):
    """One GPU job at a time — MPS does not share cleanly (PROCITY scar)."""
    spec = PROMPTS[slug]
    outdir = os.path.join(RAW, "_job_" + slug)
    os.makedirs(outdir, exist_ok=True)
    params = json.dumps({"prompt": prompt_of(slug), "model": "flux2-klein-4b",
                         "steps": STEPS, "width": SIZE, "height": SIZE,
                         "seed": spec["seed"], "quantize": "8", "guidance": GUIDANCE})
    t0 = time.time()
    r = subprocess.run([FLUX_PY, FLUX_RUN, "--outdir", outdir, "--params", params],
                       capture_output=True, text=True, timeout=900)
    pngs = [p for p in os.listdir(outdir) if p.endswith(".png")]
    if not pngs:
        _sh.rmtree(outdir, ignore_errors=True)
        raise RuntimeError((r.stderr or r.stdout or "no png")[-200:])
    fn = os.path.join(RAW, f"{slug}.png")
    _sh.move(os.path.join(outdir, pngs[0]), fn)
    _sh.rmtree(outdir, ignore_errors=True)
    return fn, time.time() - t0


def record(slug, fn, secs):
    """Append/refresh provenance so any asset can be regenerated byte-for-byte."""
    log = {}
    if os.path.exists(BATCH):
        log = json.load(open(BATCH))
    spec = PROMPTS[slug]
    log[slug] = {"prompt": prompt_of(slug), "seed": spec["seed"], "model": "flux2-klein-4b",
                 "steps": STEPS, "guidance": GUIDANCE, "size": SIZE,
                 "gen_seconds": round(secs, 1), "kb": os.path.getsize(fn) // 1024,
                 "tile": spec.get("tile"), "kind": "matcap" if slug in MATCAPS else "wall"}
    json.dump(log, open(BATCH, "w"), indent=2, sort_keys=True)


if __name__ == "__main__":
    only = sys.argv[sys.argv.index("--only") + 1] if "--only" in sys.argv else ""
    todo = {s: prompt_of(s) for s in PROMPTS if (not only or only in s)}

    if "--harvest" in sys.argv:
        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
        import derive_maps
        derive_maps.harvest(only=only)
        sys.exit(0)

    pending = {s: p for s, p in todo.items() if not have(s)}
    print(f"{len(pending)}/{len(todo)} textures to generate "
          f"(MODELBEAST flux2-klein-4b, local·free·$0)" + (f"  filter:*{only}*" if only else ""))
    if "--dry-run" in sys.argv:
        for s in sorted(todo):
            mark = "have" if have(s) else " gen"
            print(f"  [{mark}] {s:22s} seed={PROMPTS[s]['seed']}  {prompt_of(s)[:78]}…")
        sys.exit(0)
    if not local_available():
        print(f"ERROR: MODELBEAST flux_local missing ({FLUX_RUN}).\n"
              "Set MB_HOME, or run this on the m3ultra box (PIPELINE.md).")
        sys.exit(2)

    fails = 0
    for i, slug in enumerate(sorted(pending), 1):
        try:
            fn, secs = gen_local(slug)
            record(slug, fn, secs)
            print(f"[{i}/{len(pending)}] {slug:22s} {secs:5.1f}s  {os.path.getsize(fn)//1024}KB")
        except Exception as e:
            fails += 1
            print(f"[{i}/{len(pending)}] {slug:22s} FAILED: {str(e)[:120]}")
    print(f"done, {fails} failures. Review {RAW}/, delete rejects, then --harvest.")