#!/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. DO NOT add framing words ──
# Tried and rejected (round 1, Lane D — see LANE_D_NOTES §prompt-kit): appending PROCITY's
# "uniform coverage, no single central subject, filling the frame edge to edge" to the STEM.
# It cured the centred-subject failures but flattened the whole pack — esophagus_a lost its
# ribbed drama for bland wood-grain, and stomach/smallint collapsed into the same honeycomb.
# Uniformity is not the goal; VARIETY across biomes is. So the stem stays pure style, and the
# two prompts that actually drew a hero subject carry their own framing words instead.
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]),
    # "circular muscular rings" drew ONE ring, centred -> polka dots when tiled. Ask for the
    # FIELD, not the feature, and carry the framing clause in the subject.
    "wall_esophagus_b": dict(
        p="human esophageal mucosa, rows of transverse circular ridges repeating across the "
          "whole frame, wet muscular bands, filling the entire frame edge to edge",
        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]),
    # "forest of fronds" drew one radial urchin burst. Same fix; keep "finger-like" so this
    # stays VILLI and doesn't drift into smallint_b's honeycomb crypts.
    "wall_smallint_a": dict(
        p="dense field of intestinal villi, soft finger-like projections packed across the "
          "whole frame, filling the entire frame edge to edge",
        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]),

    # ── oral + colon: added round 2 by LANE A, on D's behalf, using D's pipeline unchanged ──
    # Why A: the six-biome tint contact sheet (web/dev/laneA_world.html?lvl=biomes) could only
    # eyeball 3 of 6 biomes for real — oral and large_intestine had no pack and fell back to
    # procedural, and you cannot ratify a tint against a stand-in. These complete the set that
    # ART_BIBLE §Biome palettes promises. **D owns them: ratify, retune or bin them.**
    # Seeds follow D's anatomical-order convention (11xx esophagus … 13xx smallint); oral sits
    # before esophagus at 10xx, and colon takes 15xx because 14xx was already the matcap.
    #
    # NOTE for D + F: ART_BIBLE §FLUX prompt kit has NO oral entry (it lists esophageal,
    # stomach, small intestine, colon). colon_a below is the bible's line VERBATIM; the other
    # three subjects are new text, and the mouth ones are therefore an ART_BIBLE *gap* I filled
    # rather than a spec I followed — see LANE_A_NOTES §→ Lane D for the proposed amendment.
    "wall_oral_a": dict(
        # Third subject, and the reason is worth more than the texture (PIPELINE §Textures.2:
        # "≤2 attempts, then pick a different prompt" — so this is a different prompt, not a
        # third tweak):
        #   v1 "…soft wet TONGUE surface"        -> a perfect SEM dome. Naming an ORGAN summons
        #      its silhouette: "tongue" is a thing with a shape, so FLUX drew the thing and
        #      arranged the papillae on it. That's D's "one urchin of villi" in a field's
        #      clothing, and it's why the edge-to-edge clause alone couldn't save it.
        #   v2 "…mucosa … flat top-down macro"   -> full coverage, but the papillae still
        #      radiated from a centre. The real culprit was never the framing: **papillae are
        #      3D projections, so any honest render of them has perspective, and perspective
        #      has a vanishing point** — a radial rosette that tiles into starbursts.
        #   v3 (this): ask for a subject that is FLAT BY CONSTRUCTION. Squamous epithelium is
        #      what oral mucosa actually is at SEM scale — a pavement of polygonal cells, no
        #      projections, nothing to foreshorten. It also reads clinical/bright (ART_BIBLE
        #      "tutorial-readable") and collides with nothing else in the pack.
        # Rule for the kit, generalising D's: prompt the tissue's SURFACE, not its FEATURES.
        p="oral mucosa, cobblestone pavement of flat polygonal squamous epithelial cells, "
          "tightly packed cell borders, wet glistening surface, filling the entire frame edge to edge",
        seed=1001, tile=[5, 12]),
    "wall_oral_b": dict(
        # Deliberately NOT tissue: enamel is the one hard, mineral, non-organic surface in the
        # whole game, it reads instantly as "mouth", and it gives L1's molar crush-cycles
        # (GDD/Lane C) a wall of their own. No framing clause — a surface, not a subject.
        p="tooth enamel surface, hard mineralized crystalline texture, fine cracks and wear pits",
        seed=1002, tile=[3, 20]),
    "wall_colon_a": dict(
        p="colonic mucosa, smooth undulating folds, scattered biofilm patches",   # bible, verbatim
        seed=1501, tile=[3, 20]),
    "wall_colon_b": dict(
        # A COATING, not folds — that's what keeps it clear of stomach_b's rugae and of
        # smallint_b's crypt honeycomb, and it is what "dark swamp" (ART_BIBLE) actually is.
        p="colonic wall coated in thick bacterial biofilm, lumpy microbial mats and mucus "
          "sheets covering the whole surface, filling the entire frame edge to edge",
        seed=1502, tile=[4, 15]),
}
# 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),

    # ── molecule material families (Lane A, round 2) ─────────────────────────────────────
    # GUTS has no real-time lights (ART_BIBLE), so a matcap is the ONLY way to make an object
    # look *made of* something. These are the molecule pickups' materials — see docs/MOLECULES.md.
    #
    # Why families and not one per element: derive_maps greyscales every matcap anyway
    # (`load_gray`), and Lane D already proved the pattern on the walls — author the LUMINANCE,
    # let the shader apply the tint. So one greyscale "glass" ball serves oxygen AND nitrogen at
    # their own CPK hues, and five balls dress the whole periodic table we care about. The CPK
    # colour keeps the science legible; the matcap decides whether it reads as a gem, a metal or
    # a lump of soot.
    "matcap_mol_glass": dict(
        p=("spherical material study ball, polished translucent glass marble, deep internal "
           "glow, single sharp specular highlight, studio black background"),
        seed=1410, stem_override=True),
    "matcap_mol_matte": dict(
        p=("spherical material study ball, matte graphite, fine powdery surface, soft diffuse "
           "shading, no shine, studio black background"),
        seed=1411, stem_override=True),
    "matcap_mol_chrome": dict(
        p=("spherical material study ball, polished chrome metal, mirror finish, bright hard "
           "specular highlights, studio black background"),
        seed=1412, stem_override=True),
    "matcap_mol_molten": dict(
        # v1 asked for an "incandescent white hot CORE" and got exactly that: a black ball with
        # a hot spot. A matcap is a LUMINANCE lookup that gets multiplied by the CPK tint, so a
        # dark-dominant ball makes a dark-dominant atom — ATP's phosphates came out as tiny
        # white specks with no orange in them. Ask for the whole sphere to glow, not a core.
        p=("spherical material study ball, glowing hot lava, the entire sphere uniformly "
           "incandescent and bright, even emissive glow, studio black background"),
        seed=1413, stem_override=True),
    "matcap_mol_pearl": dict(
        p=("spherical material study ball, white pearl, soft satin subsurface sheen, gentle "
           "highlight, studio black background"),
        seed=1414, 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)


def verify_provenance():
    """Assert the recorded prompt still equals what this script emits, per asset.

    Added round 2 by Lane A after finding 4 of 6 walls had drifted (see LANE_A_NOTES §-> Lane D).
    The failure is invisible and the script's own docstring is the thing it breaks
    ("regeneration must always be possible"): record() only runs when you GENERATE, so if you
    change a prompt and keep the old image — exactly what D's rejected framing experiment
    correctly did, keeping four winners rather than churning them — the JSON keeps the prompt
    that was never shipped. Nothing warns you. Re-generating from it silently yields a
    different texture (proved: docs/shots/laneA/round2_provenance_probe.png).

    Fixing a drift is a judgement call, so this reports and never rewrites:
      - the image is the one you want  => the RECORD is stale: re-point it at the code prompt.
      - the record is what you want    => regenerate the image (--local) and it self-heals.
    """
    log = json.load(open(BATCH)) if os.path.exists(BATCH) else {}
    drift, missing = [], []
    for slug in sorted(PROMPTS):
        if slug not in log:
            missing.append(slug)
        elif log[slug].get("prompt") != prompt_of(slug):
            drift.append(slug)
    for s in missing:
        print(f"  no record  {s:22s} (never generated, or provenance lost)")
    for s in drift:
        print(f"  DRIFT      {s:22s} recorded prompt != prompt this script would send")
    ok = not drift
    print(f"provenance: {len(PROMPTS) - len(drift) - len(missing)}/{len(PROMPTS)} verified, "
          f"{len(drift)} drifted, {len(missing)} unrecorded")
    return ok


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 "--verify-provenance" in sys.argv:
        sys.exit(0 if verify_provenance() else 1)

    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.")