backend/services/blender_processor.py

# backend/services/blender_processor.py
"""
S4: OMEGA Blender pipeline — one pipeline, not two drifting halves.

model in → headless Blender (inside the Space container or local PC)
         → storage/models3d/*.glb (served at /static/models3d, immediate)
         → huggingface_hub upload to the OMEGA models dataset (durable — the Space
           container FS is ephemeral per current HF docs, and pushing to the SPACE
           repo would trigger a rebuild, so a private DATASET repo is the store)
         → restore-on-demand pulls a missing GLB back after a Space restart.

Forge reuse-before-regenerate (server mirror of the client logic):
match_template() scores saved templates (storage/forge_templates/*.json) against
the part names actually inside a GLB — read with a pure-python GLB parse, no
Blender launch needed for the check. Blender only runs on a miss, and the fresh
result is saved as a new reusable template.
"""

import asyncio
import json
import logging
import os
import re
import shutil
import struct
import subprocess
import sys
import tempfile
import uuid

logger = logging.getLogger(__name__)

MODELS3D_DIR = os.path.join("storage", "models3d")
TEMPLATES_DIR = os.path.join("storage", "forge_templates")
PROCESS_SCRIPT = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                              "blender", "process_model.py")
BLENDER_TIMEOUT_S = int(os.environ.get("OMEGA_BLENDER_TIMEOUT_S", "300"))
MODELS_DATASET = os.environ.get("OMEGA_MODELS_DATASET", "Jarvis2345/omega-models3d")

os.makedirs(MODELS3D_DIR, exist_ok=True)
os.makedirs(TEMPLATES_DIR, exist_ok=True)


def slug(value: str) -> str:
    return re.sub(r"^_+|_+$", "", re.sub(r"[^a-z0-9]+", "_", str(value or "").lower()))


# ─────────────────────────────────────────────────────────────────────────────
# Blender discovery + subprocess
# ─────────────────────────────────────────────────────────────────────────────

def find_blender() -> str:
    """env override → Space install (/opt/blender) → PATH → Windows default."""
    candidates = [os.environ.get("BLENDER_PATH", ""), "/opt/blender/blender",
                  shutil.which("blender") or ""]
    if sys.platform == "win32":
        import glob as _glob
        candidates += sorted(_glob.glob(r"C:\Program Files\Blender Foundation\Blender *\blender.exe"),
                             reverse=True)
    for path in candidates:
        if path and os.path.isfile(path):
            return path
    return ""


def blender_available() -> bool:
    return bool(find_blender())


async def run_blender(script_args: list[str]) -> dict:
    """blender --background --factory-startup --python process_model.py -- <args>.
    Parses the OMEGA_BLENDER_RESULT json line the script prints."""
    blender = find_blender()
    if not blender:
        raise RuntimeError("Blender not found (BLENDER_PATH, /opt/blender, PATH all empty).")

    cmd = [blender, "--background", "--factory-startup",
           "--python", PROCESS_SCRIPT, "--", *script_args]

    def _run():
        return subprocess.run(cmd, capture_output=True, text=True,
                              timeout=BLENDER_TIMEOUT_S, encoding="utf-8", errors="replace")

    proc = await asyncio.to_thread(_run)
    marker = "OMEGA_BLENDER_RESULT "
    for line in (proc.stdout or "").splitlines():
        if line.startswith(marker):
            result = json.loads(line[len(marker):])
            if result.get("error"):
                raise RuntimeError(f"Blender script error: {result['error']}")
            return result
    tail = (proc.stderr or proc.stdout or "")[-400:]
    raise RuntimeError(f"Blender exited rc={proc.returncode} without a result marker: {tail}")


# ─────────────────────────────────────────────────────────────────────────────
# GLB inspection (pure python — the reuse check must not need a Blender launch)
# ─────────────────────────────────────────────────────────────────────────────

def read_glb_part_names(path: str) -> list[str]:
    """Node + mesh names from the GLB's JSON chunk (glTF 2.0 binary container)."""
    with open(path, "rb") as fh:
        header = fh.read(12)
        if len(header) < 12 or header[:4] != b"glTF":
            return []
        chunk_header = fh.read(8)
        if len(chunk_header) < 8:
            return []
        chunk_len, chunk_type = struct.unpack("<I4s", chunk_header)
        if chunk_type != b"JSON":
            return []
        doc = json.loads(fh.read(chunk_len).decode("utf-8", errors="replace"))
    # Nodes only: the client's discoverPartsFromObject walks the scene graph, which
    # maps to glTF NODES — mesh data-block names (Cube.001 etc.) are noise for matching.
    names = []
    for node in doc.get("nodes") or []:
        name = node.get("name")
        if name:
            names.append(name)
    return names


# ─────────────────────────────────────────────────────────────────────────────
# Template store — server mirror of client findMatchingBlueprintTemplate
# ─────────────────────────────────────────────────────────────────────────────

def load_templates() -> list[dict]:
    templates = []
    try:
        for fname in os.listdir(TEMPLATES_DIR):
            if not fname.endswith(".json"):
                continue
            try:
                with open(os.path.join(TEMPLATES_DIR, fname), encoding="utf-8") as fh:
                    templates.append(json.load(fh))
            except Exception as e:
                logger.warning(f"forge template {fname} unreadable: {e}")
    except OSError:
        pass
    return templates


def save_template(template: dict) -> str:
    template_id = template.get("id") or f"template_{uuid.uuid4().hex[:10]}"
    template["id"] = template_id
    path = os.path.join(TEMPLATES_DIR, f"{slug(template_id)}.json")
    with open(path, "w", encoding="utf-8") as fh:
        json.dump(template, fh, indent=1)
    return path


def match_template(kit: str, part_names: list[str], min_score: float = 0.5, min_hits: int = 2):
    """Same scoring contract as the client: kit equality + behavior-target overlap.

    S4 fix: templates built from models with no rig-able behaviors were saved with
    `behaviors: []` and then skipped here forever (`if not targets: continue`), so
    the reuse path was dead for every non-interactive model — each re-import
    relaunched Blender. Such templates now fall back to a near-exact structural
    match on the SOURCE part names (`matchParts`), so an identical re-import
    correctly skips the redundant Blender normalization pass.
    """
    clean_kit = str(kit or "").lower()
    names = {slug(n) for n in part_names if n}
    if not names:
        return None
    best, best_score = None, 0.0
    for template in load_templates():
        if str(template.get("kit") or "").lower() != clean_kit:
            continue
        targets = {slug(b.get("target")) for b in template.get("behaviors") or [] if b.get("target")}
        if targets:
            hits = len(targets & names)
            score = hits / len(targets)
            if hits >= min_hits and score >= min_score and score > best_score:
                best, best_score = template, score
        else:
            # No interactive behaviors: reuse only on a strong structural match of
            # the source part names, so we never falsely reuse a different model.
            sig = {slug(n) for n in (template.get("matchParts") or []) if n}
            if not sig:  # templates saved before matchParts existed
                sig = {slug(p.get("name")) for p in template.get("parts") or [] if p.get("name")}
            if not sig:
                continue
            score = len(sig & names) / len(sig)
            if score >= 0.9 and score > best_score:
                best, best_score = template, score
    return {"template": best, "score": best_score} if best else None


# ─────────────────────────────────────────────────────────────────────────────
# Durable push-back + restore-on-demand (HF dataset repo)
# ─────────────────────────────────────────────────────────────────────────────

def _hf_token() -> str:
    try:
        from backend.services.usb_vault import get_secret
        return get_secret("HF_TOKEN") or get_secret("HF_API_TOKEN") or ""
    except Exception:
        return os.environ.get("HF_TOKEN", "")


async def push_model_to_hub(local_path: str) -> str:
    """Upload a processed GLB to the OMEGA models dataset. Never raises — durability
    is best-effort on top of the immediate local serve; failures are logged."""
    token = _hf_token()
    if not token:
        logger.warning("push_model_to_hub: no HF token in vault/env — skipping durable push.")
        return ""

    def _push():
        from huggingface_hub import HfApi
        api = HfApi(token=token)
        api.create_repo(MODELS_DATASET, repo_type="dataset", private=True, exist_ok=True)
        filename = os.path.basename(local_path)
        api.upload_file(path_or_fileobj=local_path,
                        path_in_repo=f"models3d/{filename}",
                        repo_id=MODELS_DATASET, repo_type="dataset")
        return f"https://huggingface.co/datasets/{MODELS_DATASET}/resolve/main/models3d/{filename}"

    try:
        url = await asyncio.to_thread(_push)
        logger.info(f"push_model_to_hub: {os.path.basename(local_path)} → {MODELS_DATASET}")
        return url
    except Exception as e:
        logger.error(f"push_model_to_hub failed: {e}")
        return ""


async def restore_model_from_hub(filename: str) -> str:
    """Pull a GLB back from the dataset after a Space restart (ephemeral FS)."""
    token = _hf_token()

    def _pull():
        from huggingface_hub import hf_hub_download
        cached = hf_hub_download(repo_id=MODELS_DATASET, repo_type="dataset",
                                 filename=f"models3d/{filename}", token=token or None)
        dest = os.path.join(MODELS3D_DIR, filename)
        shutil.copy(cached, dest)
        return dest

    try:
        return await asyncio.to_thread(_pull)
    except Exception as e:
        logger.warning(f"restore_model_from_hub({filename}) failed: {e}")
        return ""


# ─────────────────────────────────────────────────────────────────────────────
# Public pipeline entry points
# ─────────────────────────────────────────────────────────────────────────────

async def forge_process_glb(input_path: str, name: str, kit: str,
                            template: dict | None = None) -> dict:
    """Forge handoff: reuse-check FIRST (pure-python part scan), Blender on miss.
    Fresh Blender results are saved as new reusable templates."""
    part_names = read_glb_part_names(input_path)
    if template is None:
        matched = match_template(kit, part_names)
        if matched:
            return {
                "reused": True,
                "template": matched["template"],
                "score": matched["score"],
                "model_path": input_path,
                "message": f"Existing template fits ({matched['score']:.0%} target match) — no Blender pass needed.",
            }
        template = None  # explicit: full Blender normalization pass below

    out_name = f"{slug(name) or 'model'}_{uuid.uuid4().hex[:8]}.glb"
    out_path = os.path.join(MODELS3D_DIR, out_name)
    args = ["--input", os.path.abspath(input_path), "--output", os.path.abspath(out_path),
            "--name", name, "--target-size", "1.0"]
    template_file = ""
    if template:
        fd, template_file = tempfile.mkstemp(suffix=".json")
        with os.fdopen(fd, "w", encoding="utf-8") as fh:
            json.dump(template, fh)
        args += ["--template", template_file]
    try:
        report = await run_blender(args)
    finally:
        if template_file and os.path.exists(template_file):
            os.unlink(template_file)

    new_template = {
        "id": f"template_{uuid.uuid4().hex[:10]}",
        "name": name,
        "kit": kit,
        "source": "blender",
        "parts": [{"name": n} for n in read_glb_part_names(out_path)],
        # S4: the SOURCE node names are the reuse signature — matched against the
        # next import's source names so structurally-identical models reuse this
        # template instead of relaunching Blender (see match_template fallback).
        "matchParts": part_names,
        "behaviors": (template or {}).get("behaviors") or [],
        "fromModel": out_name,
    }
    save_template(new_template)
    hub_url = await push_model_to_hub(out_path)
    return {"reused": False, "model_path": out_path, "model_url": f"/static/models3d/{out_name}",
            "hub_url": hub_url, "template": new_template, "blender_report": report}


async def build_from_blueprint(blueprint: dict, name: str) -> dict:
    """Builder: procedural Blender build from a Forge blueprint JSON."""
    out_name = f"{slug(name) or 'builder'}_{uuid.uuid4().hex[:8]}.glb"
    out_path = os.path.join(MODELS3D_DIR, out_name)
    fd, blueprint_file = tempfile.mkstemp(suffix=".json")
    with os.fdopen(fd, "w", encoding="utf-8") as fh:
        json.dump(blueprint, fh)
    try:
        report = await run_blender(["--blueprint", blueprint_file,
                                    "--output", os.path.abspath(out_path),
                                    "--name", name, "--target-size", "1.0"])
    finally:
        os.unlink(blueprint_file)
    hub_url = await push_model_to_hub(out_path)
    return {"model_path": out_path, "model_url": f"/static/models3d/{out_name}",
            "hub_url": hub_url, "blender_report": report}