__init__.py

"""H3 Studio: serves the quick-generation page and manages ephemeral outputs.

Generations are treated as disposable. A clip survives only if the user clicks
Download; anything else (tab closed, new generation, server restart, or a stale
entry) gets overwritten and unlinked.

Wipe limits, stated plainly: overwriting in place defeats undelete and
file-carving at the filesystem level, but on an SSD the flash translation layer
does wear-levelling, so the original NAND pages are not necessarily the ones
rewritten. Only full-disk encryption, ATA secure erase, or writing to a
RAM-backed filesystem gives a hard guarantee.
"""

import asyncio
import json
import os
import threading
import time
import urllib.error
import uuid

from aiohttp import web

import folder_paths
from server import PromptServer

from . import promptwriter

NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
WEB_DIRECTORY = "./web"

PREFIX = "h3_studio"          # only files starting with this are ever touched
INPUT_PREFIX = "h3_studio_kf"  # uploaded keyframes, same disposable treatment
INPUT_EXTS = (".png", ".jpg", ".jpeg", ".webp")
SUBFOLDER = "video"
STALE_AFTER = 6 * 60 * 60     # wipe unclaimed entries after 6h
SWEEP_EVERY = 600
# A reload fires the same teardown as a close, and the browser cannot tell you
# which happened. So a departing page only *schedules* its wipe; a page that
# comes back within the grace period cancels it. A real close never comes back,
# so the wipe lands a minute later instead of instantly.
SESSION_GRACE = 60
TICK = 15                     # how often pending wipes are checked

# filename -> {"kept": bool, "ts": float}
_tracked = {}
# session id -> {"deadline": float, "files": [str]}
_pending = {}
_lock = threading.Lock()


def _video_dir():
    return os.path.join(folder_paths.get_output_directory(), SUBFOLDER)


def _scoped(filename, base_dir, prefix, exts):
    """Resolve a caller-supplied name, or None if it escapes our own files."""
    if not filename or os.path.basename(filename) != filename:
        return None
    if not filename.startswith(prefix) or not filename.lower().endswith(exts):
        return None
    base = os.path.realpath(base_dir)
    path = os.path.realpath(os.path.join(base, filename))
    if os.path.dirname(path) != base:
        return None
    return path


def _safe_path(filename):
    return _scoped(filename, _video_dir(), PREFIX, (".mp4",))


def _safe_input_path(filename):
    """Keyframes we uploaded ourselves, in ComfyUI's input directory.

    Scoped the same way as outputs: basename only, our own prefix, image
    extensions, and the resolved path must sit directly in input/.
    """
    return _scoped(filename, folder_paths.get_input_directory(),
                   INPUT_PREFIX, INPUT_EXTS)


def _resolve(filename):
    """A tracked file is either a generated clip or a keyframe we uploaded."""
    return _safe_path(filename) or _safe_input_path(filename)


def secure_wipe(path, passes=3):
    """Overwrite, rename, then unlink. See module docstring for SSD caveats."""
    try:
        size = os.path.getsize(path)
        with open(path, "r+b", buffering=0) as f:
            for p in range(passes):
                f.seek(0)
                left = size
                while left > 0:
                    chunk = min(left, 1 << 20)
                    # final pass writes zeros so the tail is not random noise
                    f.write(b"\0" * chunk if p == passes - 1 else os.urandom(chunk))
                    left -= chunk
                f.flush()
                os.fsync(f.fileno())
            f.truncate(0)
            f.flush()
            os.fsync(f.fileno())
        # rename first so the old name does not linger in the directory entry
        scratch = os.path.join(os.path.dirname(path), uuid.uuid4().hex + ".tmp")
        os.replace(path, scratch)
        os.unlink(scratch)
        dir_fd = os.open(os.path.dirname(path), os.O_RDONLY)
        try:
            os.fsync(dir_fd)
        finally:
            os.close(dir_fd)
        return True
    except FileNotFoundError:
        return False
    except Exception as e:
        print(f"[h3_studio] wipe failed for {os.path.basename(path)}: {e!r}")
        return False


def _discard(filename):
    path = _resolve(filename)
    if not path:
        return False
    wiped = secure_wipe(path)
    with _lock:
        _tracked.pop(filename, None)
    if wiped:
        print(f"[h3_studio] wiped {filename}")
    return wiped


@PromptServer.instance.routes.post("/h3_studio/track")
async def track(request):
    """Register a freshly generated clip as unsaved."""
    data = await request.json()
    name = data.get("filename")
    if not _resolve(name):
        return web.json_response({"error": "rejected"}, status=400)
    with _lock:
        _tracked[name] = {"kept": False, "ts": time.time()}
    return web.json_response({"ok": True})


@PromptServer.instance.routes.post("/h3_studio/keep")
async def keep(request):
    """Mark a clip as saved so it survives."""
    data = await request.json()
    name = data.get("filename")
    if not _safe_path(name):
        return web.json_response({"error": "rejected"}, status=400)
    with _lock:
        _tracked.setdefault(name, {"ts": time.time()})["kept"] = True
    print(f"[h3_studio] keeping {name}")
    return web.json_response({"ok": True})


@PromptServer.instance.routes.post("/h3_studio/discard")
async def discard(request):
    """Wipe a clip. Also reachable via sendBeacon on tab close."""
    try:
        data = await request.json()
    except Exception:
        data = {}
    if not data:
        try:
            data = {"filename": (await request.text()).strip()}
        except Exception:
            return web.json_response({"error": "no filename"}, status=400)
    name = data.get("filename")
    if not _resolve(name):
        return web.json_response({"error": "rejected"}, status=400)
    with _lock:
        entry = _tracked.get(name)
    if entry and entry.get("kept"):
        return web.json_response({"ok": True, "skipped": "saved"})
    return web.json_response({"ok": _discard(name)})


@PromptServer.instance.routes.post("/h3_studio/session_end")
async def session_end(request):
    """A page is going away. Wipe its clips unless it returns shortly."""
    try:
        data = await request.json()
    except Exception:
        return web.json_response({"error": "bad request"}, status=400)
    sid = str(data.get("session") or "")[:64]
    files = [f for f in (data.get("filenames") or []) if _resolve(f)]
    if not sid:
        return web.json_response({"error": "no session"}, status=400)
    with _lock:
        _pending[sid] = {"deadline": time.time() + SESSION_GRACE, "files": files}
    return web.json_response({"ok": True, "grace": SESSION_GRACE,
                              "scheduled": len(files)})


@PromptServer.instance.routes.post("/h3_studio/session_resume")
async def session_resume(request):
    """The page came back — it was a reload, not a close. Cancel the wipe."""
    try:
        data = await request.json()
    except Exception:
        return web.json_response({"error": "bad request"}, status=400)
    sid = str(data.get("session") or "")[:64]
    with _lock:
        entry = _pending.pop(sid, None)
    # Anything still on disk is the page's again; re-arm its stale timer.
    kept = []
    if entry:
        with _lock:
            for name in entry["files"]:
                if _resolve(name) and os.path.exists(_resolve(name)):
                    _tracked.setdefault(name, {"kept": False})["ts"] = time.time()
                    kept.append(name)
    return web.json_response({"ok": True, "recovered": kept})


def _run_pending_wipes(now):
    """Wipe sessions whose grace period elapsed without them coming back."""
    with _lock:
        due = [(sid, m) for sid, m in _pending.items() if now >= m["deadline"]]
        for sid, _ in due:
            _pending.pop(sid, None)
    for sid, meta in due:
        for name in meta["files"]:
            with _lock:
                entry = _tracked.get(name)
            if entry and entry.get("kept"):
                continue
            _discard(name)
        if meta["files"]:
            print(f"[h3_studio] session {sid[:8]} did not return; "
                  f"wiped {len(meta['files'])} file(s)")


@PromptServer.instance.routes.get("/h3_studio/settings")
async def get_settings(request):
    s = promptwriter.settings()
    return web.json_response({
        "llm_base": s.get("llm_base") or promptwriter.DEFAULT_LLM_BASE,
        "llm_model": s.get("llm_model") or "",
        "resolved_model": promptwriter.model_name(),
        "env_locked": bool(os.environ.get("H3_LLM_BASE")),
    })


@PromptServer.instance.routes.post("/h3_studio/settings")
async def set_settings(request):
    data = await request.json()
    base = (data.get("llm_base") or "").strip().rstrip("/")
    if base and not base.startswith(("http://", "https://")):
        return web.json_response({"error": "endpoint must start with http:// or https://"},
                                 status=400)
    promptwriter.save_settings(llm_base=base or None,
                               llm_model=(data.get("llm_model") or "").strip())
    return web.json_response({"ok": True, "resolved_model": promptwriter.model_name()})


@PromptServer.instance.routes.get("/h3_studio/llm_status")
async def llm_status(request):
    return web.json_response(promptwriter.llm_status())


@PromptServer.instance.routes.post("/h3_studio/unload_llm")
async def unload_llm(request):
    """Hand the GPU back before an H3 run — they cannot both be resident."""
    return web.json_response({"ok": promptwriter.unload_llm()})


def comfy_busy():
    running, pending = PromptServer.instance.prompt_queue.get_current_queue()
    return bool(running or pending)


def llm_blocked_reason():
    """Explain an LLM failure the user can actually act on.

    The common case is that H3 is mid-generation holding the card, so the LLM
    cannot load. A bare 500 from the upstream tells the user nothing.
    """
    if comfy_busy():
        return ("H3 is generating right now, so the LLM cannot load — they "
                "cannot share the GPU. Wait for it to finish, or press Cancel.")
    return None


def free_comfy_vram():
    """Drop H3's weights so the LLM can load.

    The card cannot hold both. /free only sets a flag for the main loop, which
    is too slow to help a request that is about to start the LLM, so unload
    directly — but never while a job is running, or we would evict weights out
    from under the sampler.
    """
    q = PromptServer.instance.prompt_queue
    running, pending = q.get_current_queue()
    if running or pending:
        return False
    try:
        import comfy.model_management as mm
        mm.unload_all_models()
        mm.soft_empty_cache()
        return True
    except Exception as e:
        print(f"[h3_studio] could not free VRAM: {e!r}")
        return False


@PromptServer.instance.routes.post("/h3_studio/free_comfy")
async def free_comfy(request):
    return web.json_response({"freed": free_comfy_vram()})


@PromptServer.instance.routes.get("/h3_studio/gpu")
async def gpu(request):
    """Device-wide VRAM, so the UI can show what H3 and the LLM are competing for."""
    try:
        import torch
        free, total = torch.cuda.mem_get_info()
        return web.json_response({
            "used_mb": round((total - free) / 1048576),
            "total_mb": round(total / 1048576),
        })
    except Exception as e:
        return web.json_response({"error": repr(e)[:120]}, status=503)


@PromptServer.instance.routes.post("/h3_studio/write_prompt_stream")
async def write_prompt_stream(request):
    """Server-sent events so the UI can show tokens arriving live."""
    data = await request.json()
    seconds = float(data.get("seconds") or 5)
    brief = data.get("brief", "")
    anchors = promptwriter.anchors_from(brief, data.get("anchors", ""))
    current = (data.get("current") or "").strip()
    editing = data.get("mode") == "edit" and bool(current)
    converting = data.get("mode") == "convert" and bool(current)
    # "first" / "last" / "both" / None — changes how the body must be written.
    task = promptwriter.task_mode(data.get("keyframes"))
    # Reference mode is a different checkpoint and a different six-section
    # schema, so it takes its own prompt path rather than a flag on this one.
    subjects = int(data.get("subjects") or 0)
    # Attached photos, as data URLs, so the model describes the actual pixels
    # instead of inventing attributes from the brief. Bounded hard: two images,
    # ~4 MB each, images only — this arrives from the network.
    photos = [u for u in (data.get("photos") or [])
              if isinstance(u, str) and u.startswith("data:image/")
              and len(u) < 4_000_000][:2]

    if converting:
        # Reshape an existing prompt into the other family's schema. The
        # words are the user's and worth keeping; only the structure is wrong.
        msgs = promptwriter.convert_messages(current, subjects, seconds,
                                             photos)
        temp, budget = 0.3, None
    elif editing:
        # Edits are a diff against the existing prompt, not a regeneration:
        # smaller output, and everything untouched stays byte-identical.
        # Schema-agnostic — find/replace works on either shape.
        msgs = promptwriter.edit_messages(current, brief,
                                          data.get("history") or [], task,
                                          subjects, photos)
        temp, budget = 0.3, promptwriter.EDIT_MAX_TOKENS
    elif subjects:
        msgs = promptwriter._ref_messages(brief, seconds, subjects, anchors,
                                          photos)
        temp, budget = 0.8, None
    else:
        msgs = promptwriter._messages(brief, seconds, anchors, task, photos)
        temp, budget = 0.8, None

    resp = web.StreamResponse(headers={
        "Content-Type": "text/event-stream",
        "Cache-Control": "no-cache",
        "X-Accel-Buffering": "no",
    })
    await resp.prepare(request)

    async def send(obj):
        # A closed tab mid-stream is normal, not an error worth a traceback.
        try:
            await resp.write(f"data: {json.dumps(obj)}\n\n".encode())
            return True
        except (ConnectionResetError, RuntimeError, Exception) as e:
            if "clos" in repr(e).lower() or isinstance(e, ConnectionResetError):
                return False
            raise

    blocked = llm_blocked_reason()
    if blocked:
        await send({"type": "error", "error": blocked})
        await resp.write_eof()
        return resp

    # The LLM needs the card that H3's weights are sitting on.
    if free_comfy_vram():
        await send({"type": "freed_vram"})

    loop = asyncio.get_running_loop()
    queue = asyncio.Queue()

    no_vision = {"hit": False}

    def produce():
        use = msgs
        while True:
            try:
                for kind, payload in promptwriter.stream(use, temp, budget):
                    loop.call_soon_threadsafe(queue.put_nowait, (kind, payload))
                break
            except urllib.error.HTTPError as e:
                try:
                    detail = e.read().decode("utf-8", "replace")
                except Exception:
                    detail = ""
                # An endpoint without a vision projector rejects the whole
                # request. The words are still good, so retry without the
                # pixels — which is exactly the pre-vision behaviour — and say
                # so rather than failing the write. Errors surface before any
                # tokens stream, so the retry cannot duplicate output.
                if use is msgs and photos and (
                        "image input" in detail or "mmproj" in detail):
                    no_vision["hit"] = True
                    use = promptwriter.strip_photos(msgs)
                    continue
                loop.call_soon_threadsafe(
                    queue.put_nowait, ("error", detail[:200] or repr(e)[:200]))
                break
            except Exception as e:
                loop.call_soon_threadsafe(
                    queue.put_nowait, ("error", repr(e)[:200]))
                break
        loop.call_soon_threadsafe(queue.put_nowait, None)

    threading.Thread(target=produce, daemon=True).start()

    try:
        while True:
            item = await queue.get()
            if item is None:
                break
            kind, payload = item
            if kind == "done":
                note, ops_json = None, None
                if editing:
                    ops = promptwriter.parse_edits(payload)
                    if ops is None:
                        # model ignored the format and rewrote it wholesale
                        result = promptwriter._clean(payload)
                        note = "model rewrote the prompt instead of editing it"
                    elif not ops:
                        result, note = current, "no change needed"
                        ops_json = "[]"
                    else:
                        result, applied, failed = promptwriter.apply_edits(current, ops)
                        note = f"{len(applied)} edit(s) applied"
                        if failed:
                            note += f", {len(failed)} could not be located"
                        # only what landed goes into history — replaying a
                        # failed op as context teaches the model bad `find`s
                        ops_json = promptwriter.ops_digest(ops, applied)
                elif subjects:
                    result = promptwriter._clean_ref(payload)
                    if converting:
                        note = "converted to the reference format"
                else:
                    result = promptwriter._clean(payload)

                # an edited prompt keeps whichever schema it already had
                is_ref = bool(subjects) and "detailed_description:" in result
                warn = (promptwriter.validate_ref(result, seconds, subjects)
                        if is_ref else promptwriter.validate(result, seconds))
                lost = promptwriter.missing_anchors(result, anchors)
                if lost:
                    warn.append("dropped from the brief: " + ", ".join(lost))

                # Hand the failures back to the model together with the rule
                # behind each one. A validator message on its own ("missing
                # [Shot 1] marker") is a complaint; the remedy text is what
                # makes it actionable.
                # Refinements get this too: an edit that introduces a format
                # fault should not be the user's problem to notice and undo.
                if warn:
                    await send({"type": "repairing", "problems": warn})
                    fixed, warn, rounds = await loop.run_in_executor(
                        None, promptwriter.repair, result, warn, task,
                        subjects if is_ref else 0, seconds)
                    if rounds:
                        result = fixed
                        done_note = (f"self-corrected in {rounds} pass"
                                     f"{'es' if rounds > 1 else ''}")
                        note = f"{note}; {done_note}" if note else done_note

                if no_vision["hit"]:
                    fb = ("LLM endpoint has no vision (mmproj) — photos were "
                          "not shown to it")
                    note = f"{note}; {fb}" if note else fb
                await send({"type": "done", "prompt": result, "warnings": warn,
                            "anchors": anchors, "note": note, "ops": ops_json})
            elif kind == "error":
                await send({"type": "error",
                            "error": llm_blocked_reason() or payload})
            else:
                await send({"type": kind, "n": len(payload)})
    finally:
        try:
            await resp.write_eof()
        except Exception:
            pass          # client already gone
    return resp


@PromptServer.instance.routes.post("/h3_studio/write_prompt")
async def write_prompt(request):
    data = await request.json()
    mode = data.get("mode", "new")
    seconds = float(data.get("seconds") or 5)
    free_comfy_vram()
    try:
        if mode == "edit":
            text = promptwriter.edit(data.get("current", ""),
                                     data.get("brief", ""), seconds)
        else:
            text = promptwriter.write(data.get("brief", ""), seconds)
    except Exception as e:
        return web.json_response(
            {"error": llm_blocked_reason() or repr(e)[:300]}, status=502)
    if not text:
        return web.json_response(
            {"error": "model returned nothing (token budget too small?)"},
            status=502)
    return web.json_response({"prompt": text,
                              "warnings": promptwriter.validate(text, seconds)})


def _untracked_debris(now):
    """Files on disk that no live session claims.

    A job abandoned mid-run — browser killed, machine slept, beacon lost —
    still finishes and writes its clip, and no client is left to register it.
    Without this such a file would outlive every other wipe path and sit there
    until the next restart. Age-gated by STALE_AFTER so an in-flight run is
    never a candidate.
    """
    found = []
    for d, resolve in ((_video_dir(), _safe_path),
                       (folder_paths.get_input_directory(), _safe_input_path)):
        if not os.path.isdir(d):
            continue
        for name in os.listdir(d):
            path = resolve(name)
            if not path:
                continue
            with _lock:
                if name in _tracked:
                    continue          # a session owns it; handled below
            try:
                if now - os.path.getmtime(path) > STALE_AFTER:
                    found.append(name)
            except OSError:
                pass
    return found


def _sweep_stale():
    """Catch clips whose tab died before its beacon landed."""
    last_deep = time.time()
    while True:
        time.sleep(TICK)
        now = time.time()
        _run_pending_wipes(now)
        if now - last_deep < SWEEP_EVERY:
            continue
        last_deep = now
        with _lock:
            doomed = [n for n, m in _tracked.items()
                      if not m.get("kept") and now - m["ts"] > STALE_AFTER]
        for name in doomed + _untracked_debris(now):
            _discard(name)


def _wipe_orphans_at_startup():
    """Anything left from a previous session was never saved by definition.

    Covers uploaded keyframes too: they are written just before a run and
    wiped as soon as it ends, so one surviving a restart is by definition
    debris from a session that died mid-generation.
    """
    for d, resolve in ((_video_dir(), _safe_path),
                       (folder_paths.get_input_directory(), _safe_input_path)):
        if not os.path.isdir(d):
            continue
        for name in os.listdir(d):
            path = resolve(name)
            if path:
                secure_wipe(path)
                print(f"[h3_studio] wiped orphan {name}")


_wipe_orphans_at_startup()
threading.Thread(target=_sweep_stale, daemon=True).start()

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]