src/server/bpy_server.py

"""
File-based IPC bpy server.

Instead of listening on a TCP port (which triggers HF Space proxy issues),
this server communicates with the main process via the filesystem:

  - Requests  are written to  {tmpdir}/req_{uuid}.torch  +  .ready flag
  - Responses are written to  {tmpdir}/resp_{uuid}.torch +  .done  flag

All data (potentially large meshes / assets) lives on disk, avoiding
HTTP body-size limits and proxy interference entirely.

bpy MUST run in its own *process* (not thread) — this server is launched
as a standalone subprocess via `bpy_server.py` at the repo root.

**Crash recovery**: if any load/export/transfer operation raises an
exception (or hangs beyond ``OP_TIMEOUT`` seconds), the server writes an
error response to the client and then **exits immediately**.  The client
detects the dead server via the PID file and launches a fresh instance.
This guarantees a clean bpy state for every request.
"""

import glob
import os
import signal
import sys
import time

from .spec import bytes_to_object, object_to_bytes, BPY_IPC_DIR

from ..rig_package.parser.bpy import BpyParser, transfer_rigging

# --- per-operation watchdog (seconds) ---------------------------------------
OP_TIMEOUT = 50  # max time for a single load / export / transfer (must be < client timeout of 60s)


def _on_op_timeout(signum, frame):
    """Called by SIGALRM when an operation exceeds OP_TIMEOUT."""
    print(f"[bpy_server] FATAL: operation timed out after {OP_TIMEOUT}s", flush=True)
    # _exit bypasses Python cleanup (bpy may be stuck); the parent will
    # detect the dead server via the missing PID and restart.
    os._exit(1)


def run():
    """
    Main loop of the bpy worker process.

    It polls ``{BPY_IPC_DIR}/*.ready`` files, reads the companion
    ``.torch`` request, dispatches to the appropriate handler, writes
    the result to a ``resp_{uuid}.torch`` file, and signals completion
    with a ``resp_{uuid}.done`` flag.

    A ``.ready`` file inside ``BPY_IPC_DIR`` signals that the server has
    finished importing bpy and is ready to accept work.

    **On any dispatch error the process exits** so the next request
    always starts with a fresh bpy interpreter.
    """
    # --- ensure the IPC directory exists -----------------------------------
    os.makedirs(BPY_IPC_DIR, exist_ok=True)

    # --- write PID so other processes can quickly detect us ----------------
    _write_pid()

    # --- signal readiness --------------------------------------------------
    ready_flag = os.path.join(BPY_IPC_DIR, ".ready")
    with open(ready_flag, "w") as f:
        f.write("ready")
    print(f"[bpy_server] ready  (pid={os.getpid()}, ipc_dir={BPY_IPC_DIR})", flush=True)

    # --- main loop ---------------------------------------------------------
    while True:
        # Look for any pending request (a `.ready` flag file).
        ready_files = sorted(glob.glob(os.path.join(BPY_IPC_DIR, "req_*.ready")))
        if not ready_files:
            time.sleep(0.05)  # 50 ms poll – negligible overhead vs bpy ops
            continue

        ready_file = ready_files[0]
        base = ready_file[:-len(".ready")]  # strip .ready suffix
        req_id = os.path.basename(base)      # e.g. "req_abc123"

        req_torch = base + ".torch"
        resp_torch = os.path.join(BPY_IPC_DIR, req_id.replace("req_", "resp_", 1) + ".torch")
        resp_done = os.path.join(BPY_IPC_DIR, req_id.replace("req_", "resp_", 1) + ".done")

        # --- read request --------------------------------------------------
        if not os.path.exists(req_torch):
            # Stale / partial – clean up and move on.
            _safe_remove(ready_file)
            _safe_remove(req_torch)
            continue

        try:
            with open(req_torch, "rb") as f:
                raw = f.read()
            op, data = bytes_to_object(raw)
        except Exception as exc:
            print(f"[bpy_server] ERROR reading request {req_id}: {exc}", flush=True)
            _safe_remove(ready_file)
            _safe_remove(req_torch)
            continue

        # --- dispatch ------------------------------------------------------
        fatal = False
        result = None  # type: ignore
        try:
            # Arm watchdog – SIGALRM fires if the op hangs (e.g. bpy deadlock).
            signal.signal(signal.SIGALRM, _on_op_timeout)
            signal.alarm(OP_TIMEOUT)

            if op == "ping":
                result = "pong"
                print("[bpy_server] ping", flush=True)
            elif op == "load":
                print(f"[bpy_server] load: {data}", flush=True)
                asset = BpyParser.load(data)
                result = asset
            elif op == "export":
                print(f"[bpy_server] export: {data.get('filepath', '?')}", flush=True)
                BpyParser.export(**data)
                result = "ok"
            elif op == "transfer":
                print(f"[bpy_server] transfer: {data.get('target_path', '?')}", flush=True)
                transfer_rigging(**data)
                result = "ok"
            else:
                result = f"unsupported op: {str(op)}"

            # Disarm watchdog
            signal.alarm(0)

        except Exception as exc:
            signal.alarm(0)  # disarm
            # Include context (file path for load, export path for export)
            ctx = ""
            if op == "load":
                ctx = f" (file: {data})"
            elif op == "export":
                ctx = f" (file: {data.get('filepath', '?') if isinstance(data, dict) else data})"
            elif op == "transfer":
                ctx = f" (target: {data.get('target_path', '?') if isinstance(data, dict) else data})"
            result = f"error: [{op}]{ctx} — {exc}"
            print(f"[bpy_server] ERROR {result}", flush=True)
            # bpy state may be corrupted — exit after writing response so
            # the next request gets a clean interpreter.
            fatal = True

        # --- write response ------------------------------------------------
        try:
            with open(resp_torch, "wb") as f:
                print(f"[bpy_server] writing response to {resp_torch}", flush=True)
                f.write(object_to_bytes(result))
            with open(resp_done, "w") as f:
                print(f"[bpy_server] writing done to {resp_done}", flush=True)
                f.write("done")
        except Exception as exc:
            print(f"[bpy_server] ERROR writing response {req_id}: {exc}", flush=True)

        # --- clean up request ----------------------------------------------
        _safe_remove(ready_file)
        _safe_remove(req_torch)

        # --- exit on fatal error so client restarts us ---------------------
        if fatal:
            print("[bpy_server] exiting after error — restart for clean state", flush=True)
            _cleanup_ipc_state()
            sys.exit(1)


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def _write_pid():
    pid_file = os.path.join(BPY_IPC_DIR, ".pid")
    with open(pid_file, "w") as f:
        f.write(str(os.getpid()))


def _cleanup_ipc_state():
    """Remove our PID file so clients don't think we're still alive."""
    pid_file = os.path.join(BPY_IPC_DIR, ".pid")
    _safe_remove(pid_file)
    # Also remove .ready since it's no longer valid
    ready_flag = os.path.join(BPY_IPC_DIR, ".ready")
    _safe_remove(ready_flag)


def _safe_remove(path: str):
    """Remove a file if it exists; never raise."""
    try:
        os.remove(path)
    except OSError:
        pass