"""MOSS-Image as a gr.Workflow app — node-canvas text-to-image pipeline.

A reimagining of https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image as a
Gradio Workflow (gradio 6.30.0): generation is a bound Python-function node on
a drag-and-drop canvas, fed by reference nodes (prompt, size, and every
sampling knob of the original demo) and writing into an output-gallery subject.
Edit the graph on the canvas or in workflow.json; the App view and REST API
stay in sync automatically.

Both checkpoints are loaded once at module scope (ZeroGPU packs and transfers
them when a GPU is actually allocated) and shared by every run of the node.
"""

import os
import secrets
import tempfile
import time
from collections import OrderedDict
from pathlib import Path

# `spaces` must be imported before anything that initializes CUDA.
import spaces  # ty: ignore[unresolved-import]
from gradio_client import utils as client_utils

import gradio as gr
from gradio import workflow_api as _wf_api

# Workaround for gradio 6.30.0: the App view hands a subject's value to its
# output component via `_from_output`, which unwraps a top-level file dict to
# its path but leaves dicts nested in a LIST untouched — and gr.Gallery only
# accepts path strings, not file dicts. A gallery subject fed by an fn node
# (which must return file dicts for the canvas) therefore crashes the App
# view. Unwrap list items the same way so both views work.
_orig_from_output = _wf_api._from_output


def _from_output_unwrapping_lists(value):
    if isinstance(value, list):
        return [_orig_from_output(item) for item in value]
    return _orig_from_output(value)


_wf_api._from_output = _from_output_unwrapping_lists


def _install_component_server_keepalive():
    """Keep `/gradio_api/component_server` connections alive during long fn calls.

    Canvas fn nodes call back through the component_server route — a plain
    POST that awaits the bound function synchronously, unlike App-view
    endpoints, which run through the queue with SSE heartbeats. A MOSS-Image
    generation takes 45-130s, and the silent connection is dropped by the
    proxy chain, surfacing on the canvas as "Connection errored out" (the
    App view is unaffected). Wrap the route so it streams whitespace
    heartbeat bytes — legal leading whitespace in JSON — until the real
    result is ready.
    """
    import asyncio
    import logging

    from fastapi import Request
    from fastapi.dependencies.utils import get_dependant
    from fastapi.routing import APIRoute
    from starlette.responses import StreamingResponse
    from starlette.routing import request_response
    from gradio.routes import App as _RoutesApp

    original_create_app = _RoutesApp.create_app

    def _wrap(endpoint):
        async def component_server_with_keepalive(request: Request):
            task = asyncio.create_task(endpoint(request))

            async def body():
                try:
                    while not task.done():
                        yield b" "
                        await asyncio.sleep(10)
                    payload = json.dumps(task.result(), default=str).encode()
                except Exception as exc:  # keep the canvas error path readable
                    payload = json.dumps({"error": str(exc)}).encode()
                yield payload

            return StreamingResponse(body(), media_type="application/json")

        return component_server_with_keepalive

    @staticmethod
    def create_app(blocks, *args, **kwargs):
        app = original_create_app(blocks, *args, **kwargs)
        for route in app.routes:
            if isinstance(route, APIRoute) and route.path.rstrip("/").endswith(
                "/component_server"
            ):
                route.endpoint = _wrap(route.endpoint)
                dependant = get_dependant(
                    path=route.path_format, endpoint=route.endpoint
                )
                # Keep the route-level dependencies (e.g. Depends(login_check)).
                dependant.dependencies = (
                    route.dependant.dependencies + dependant.dependencies
                )
                route.dependant = dependant
                route.app = request_response(route.get_route_handler())
                logging.getLogger(__name__).info(
                    "component_server keep-alive wrapper installed"
                )
        return app

    _RoutesApp.create_app = create_app


_install_component_server_keepalive()

MODEL_IDS = {
    "base": "OpenMOSS-Team/MOSS-Image-Base",
    "sft": "OpenMOSS-Team/MOSS-Image-SFT",
}

# Set MOSS_WORKFLOW_PREVIEW=1 to open the canvas without downloading weights
# (useful for local UI work on a CPU-only machine).
PREVIEW = os.environ.get("MOSS_WORKFLOW_PREVIEW") == "1"


def _load_pipeline(model_id, shared_pipe=None):
    """Load one MOSS-Image checkpoint with the bundled inference classes.

    Text encoder, tokenizer and VAE are shared with `shared_pipe` when their
    Hub file hashes match — this is what keeps both checkpoints inside a
    ZeroGPU worker's GPU memory (two full copies do not fit).
    """
    import torch
    from huggingface_hub import HfApi, snapshot_download
    from pipeline_moss_image import (
        MOSSImagePipeline,
        MOSSImageScheduler,
        MOSSImageTransformer,
    )

    info = HfApi().model_info(model_id, files_metadata=True)
    revision = info.sha
    shared_hashes = {
        item.rfilename: item.lfs.sha256 if item.lfs else item.blob_id
        for item in info.siblings
        if item.rfilename.startswith(("text_encoder/", "tokenizer/", "vae/"))
    }
    components = {}
    if shared_pipe is not None:
        for name in ("text_encoder", "tokenizer", "vae"):
            current = {k: v for k, v in shared_hashes.items() if k.startswith(f"{name}/")}
            previous = {
                k: v for k, v in shared_pipe._shared_component_hashes.items()
                if k.startswith(f"{name}/")
            }
            if current and all(current.values()) and current == previous:
                components[name] = getattr(shared_pipe, name)
    snapshot = snapshot_download(
        model_id, revision=revision, ignore_patterns=[f"{name}/*" for name in components]
    )
    pipe = MOSSImagePipeline.from_pretrained(
        snapshot,
        custom_pipeline=str(Path(__file__).with_name("pipeline_moss_image")),
        trust_remote_code=True,
        transformer=MOSSImageTransformer.from_pretrained(
            snapshot, subfolder="transformer", torch_dtype=torch.bfloat16
        ),
        scheduler=MOSSImageScheduler.from_pretrained(snapshot, subfolder="scheduler"),
        local_files_only=True,
        torch_dtype={"default": torch.bfloat16, "vae": torch.float32},
        **components,
    )
    # On ZeroGPU the `spaces` package intercepts this move, packs the weights,
    # and transfers them when a GPU is allocated — module scope costs no GPU time.
    pipe.to("cuda")
    pipe._prompt_cache = OrderedDict()
    pipe._shared_component_hashes = shared_hashes
    return pipe


PIPELINES = {}
if not PREVIEW:
    PIPELINES["base"] = _load_pipeline(MODEL_IDS["base"])
    PIPELINES["sft"] = _load_pipeline(MODEL_IDS["sft"], shared_pipe=PIPELINES["base"])


def _as_media(path: str) -> dict:
    """The dict shape the workflow canvas renders media from."""
    return {
        "path": path,
        "url": f"/gradio_api/file={client_utils.encode_file_path(path)}",
        "is_file": True,
    }


def _parse_layers(value: str, name: str) -> tuple:
    """'22' -> (22,); '' -> (). Same rules as the original demo."""
    if not (value or "").strip():
        return ()
    try:
        layers = tuple(int(item.strip()) for item in value.split(","))
    except ValueError:
        raise gr.Error(
            f"{name} must be comma-separated integer indices, e.g. 22."
        ) from None
    if any(layer < 0 for layer in layers) or len(set(layers)) != len(layers):
        raise gr.Error(f"{name} indices must be nonnegative and unique.")
    return layers


@spaces.GPU(duration=80)
def generate(
    prompt: str,
    negative_prompt: str,
    model: str,
    width: float,
    height: float,
    steps: float,
    cfg: float,
    shift: float,
    seed: float,
    num_images: float,
    max_sequence_length: float,
    uncond_skip_layers: str,
    slg_layers: str,
    slg_scale: float,
    slg_start: float,
    slg_stop: float,
) -> list:
    """Generate images with MOSS-Image (Base or SFT). Full sampling controls."""
    import torch

    if PREVIEW:
        raise gr.Error("Preview mode: models are not loaded (MOSS_WORKFLOW_PREVIEW=1).")
    if not (prompt or "").strip():
        raise gr.Error("Enter a prompt.")
    key = (model or "").strip().lower().replace("moss-image-", "")
    if key not in PIPELINES:
        raise gr.Error(f"Unknown model {model!r}. Choose 'base' or 'sft'.")
    width, height, steps, num_images, max_sequence_length = (
        round(width), round(height), round(steps), round(num_images),
        round(max_sequence_length),
    )
    seed = round(seed)
    if width <= 0 or height <= 0 or width % 16 or height % 16:
        raise gr.Error("Width and height must be positive multiples of 16.")
    if steps < 1 or num_images < 1:
        raise gr.Error("Steps and image count must be positive.")
    if not 0 <= cfg or shift <= 0:
        raise gr.Error("CFG must be nonnegative; shift must be positive.")
    if not 1 <= max_sequence_length <= 4096:
        raise gr.Error("Maximum text length must be between 1 and 4096.")
    if slg_scale < 0 or not 0 <= slg_start <= slg_stop <= 1:
        raise gr.Error("Require SLG scale >= 0 and 0 <= SLG start <= SLG stop <= 1.")
    if seed == -1:
        seed = secrets.randbelow(2**32)
    if not 0 <= seed < 2**53:
        raise gr.Error("Seed must be -1 (random) or an integer from 0 to 2^53 - 1.")

    pipe = PIPELINES[key]
    depth = pipe.transformer.config.num_layers
    skip = _parse_layers(uncond_skip_layers, "Skip layers")
    slg = _parse_layers(slg_layers, "SLG layers")
    if any(layer >= depth for layer in skip + slg):
        raise gr.Error(f"Skip-layer indices must be below {depth}.")

    pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config, shift=shift)
    start = time.monotonic()
    try:
        images = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt or "",
            width=width,
            height=height,
            num_inference_steps=steps,
            guidance_scale=float(cfg),
            num_images_per_prompt=num_images,
            max_sequence_length=max_sequence_length,
            uncond_skip_layers=skip,
            slg_layers=slg,
            slg_scale=float(slg_scale),
            slg_start=float(slg_start),
            slg_stop=float(slg_stop),
            generator=torch.Generator("cpu").manual_seed(seed),
        ).images
    except torch.cuda.OutOfMemoryError as exc:
        torch.cuda.empty_cache()
        raise gr.Error("GPU memory exhausted. Try fewer or smaller images.") from exc

    results = []
    batch = os.urandom(8).hex()
    for i, image in enumerate(images):
        path = os.path.join(tempfile.gettempdir(), f"moss_{batch}_{i + 1:02d}.png")
        image.save(path)
        results.append(_as_media(path))
    elapsed = time.monotonic() - start
    status = (
        f"{len(results)} image(s) · {MODEL_IDS[key].rsplit('/', 1)[-1]} "
        f"· seed {seed} · {width}×{height} · {steps} steps "
        f"· CFG {cfg:g} · shift {shift:g} · {elapsed:.1f}s"
    )
    # call_fn maps a fn's returned list/tuple onto the node's output PORTS
    # (one element per port), so the gallery — one port holding many images —
    # must come back as a single element containing the list: [results].
    # Returning `results` bare would hand the gallery port only image #1.
    return [results, status]


demo = gr.Workflow(
    graph=str(Path(__file__).with_name("workflow.json")),
    bind={"generate": generate},
)
demo.title = "MOSS-Image"

# --- Landing-page styling for the auto-generated App view -------------------
# gr.Workflow builds its App view inside #workflow-app; launch() threads
# theme/css/head straight through to Blocks.launch(). The hero uses the
# official lockup SVGs from the original Space (hot-linked, Apache-2.0).

_LOCKUP_LIGHT = (
    "https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image"
    "/resolve/main/assets/lockup-light.svg"
)
_LOCKUP_DARK = (
    "https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image"
    "/resolve/main/assets/lockup-dark.svg"
)

HEAD = """
<link rel="preconnect" href="https://fonts.googleapis.com">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap" rel="stylesheet">
"""

# Custom JS for launch(js=...): the gradio frontend executes this function once
# the page is up. A <script> tag in head= is unreliable for this under gradio 6
# SSR — the frontend can inject head HTML after the window load event, and
# innerHTML-inserted <script> tags never execute.
#
# Two pieces: the top nav, and the examples strip — the original Space's
# showcase, hot-linked. Clicking a thumbnail fetches its parameters.json and
# restores every field of the generated App-view form, replicating the
# original's click-to-restore.
JS = """
() => {
  const ORIGIN = "https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image/resolve/main/assets/images";
  const IDS = Array.from({length: 20}, (_, i) => String(i + 1).padStart(3, "0"));

  function setField(prefix, value) {
    if (value === undefined || value === null) return;
    const labels = Array.from(document.querySelectorAll('#workflow-app label'));
    const match = labels.find(l => (l.textContent || "").trim().toLowerCase().startsWith(prefix));
    if (!match) return;
    let root = match;
    for (let i = 0; i < 6 && root; i++) {
      const input = root.querySelector('input:not([type="hidden"]), textarea');
      if (input) {
        input.value = value;
        input.dispatchEvent(new Event("input", {bubbles: true}));
        input.dispatchEvent(new Event("change", {bubbles: true}));
        return;
      }
      root = root.parentElement;
    }
  }

  function applyRecord(rec) {
    const p = rec.parameters || {};
    let uncond = p.uncond_skip_layers ?? "";
    let slgLayers = p.slg_layers ?? "";
    let slgScale = p.slg_scale ?? 0;
    if ((rec.schema_version ?? 2) === 1) {
      // v1 records called negative-branch skipping "slg_layers".
      uncond = (p.slg ?? true) ? (p.slg_layers ?? "") : "";
      slgLayers = "";
      slgScale = 0;
    }
    const model = (rec.model_repo || rec.model || "").toLowerCase().includes("sft") ? "sft" : "base";
    setField("negative", p.negative_prompt ?? "");
    setField("prompt", p.prompt ?? "");
    setField("model", model);
    setField("width", p.width);
    setField("height", p.height);
    setField("steps", p.steps);
    setField("cfg", p.cfg);
    setField("shift", p.shift);
    setField("seed", p.seed ?? -1);
    setField("images", p.num_images ?? 1);
    setField("max text", p.max_sequence_length ?? 4096);
    setField("skip start", p.slg_start ?? 0);
    setField("skip stop", p.slg_stop ?? 1);
    setField("skip layers", uncond);
    setField("slg layers", slgLayers);
    setField("slg scale", slgScale);
  }

  const recordCache = {};
  function fetchRecord(id) {
    if (!recordCache[id]) {
      recordCache[id] = fetch(`${ORIGIN}/${id}/parameters.json`)
        .then(r => (r.ok ? r.json() : null))
        .catch(() => null);
    }
    return recordCache[id];
  }

  function buildStrip(app) {
    if (document.getElementById('moss-examples')) return true;
    const wrap = document.createElement('div');
    wrap.id = 'moss-examples';
    const title = document.createElement('div');
    title.className = 'mx-title';
    title.textContent = 'Examples · click to restore prompt & settings';
    wrap.appendChild(title);
    const row = document.createElement('div');
    row.className = 'mx-row';
    IDS.forEach((id, i) => {
      const fig = document.createElement('button');
      fig.className = 'mx-item';
      fig.type = 'button';
      fig.title = 'Load this example into the form';
      const img = document.createElement('img');
      img.loading = 'lazy';
      img.alt = `Example ${i + 1}`;
      img.src = `${ORIGIN}/${id}/image.webp`;
      img.onerror = () => fig.remove();
      fig.appendChild(img);
      fig.addEventListener('click', () => {
        fetchRecord(id).then(rec => { if (rec) applyRecord(rec); });
      });
      row.appendChild(fig);
    });
    wrap.appendChild(row);
    app.prepend(wrap);
    return true;
  }

  if (!document.getElementById('moss-nav')) {
    const nav = document.createElement('nav');
    nav.id = 'moss-nav';
    nav.innerHTML = `
      <span class="moss-brand">MOSS-Image</span>
      <span class="moss-links">
        <a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base" target="_blank" rel="noopener">Base</a>
        <a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT" target="_blank" rel="noopener">SFT</a>
        <a href="https://huggingface.co/spaces/OpenMOSS-Team/MOSS-Image" target="_blank" rel="noopener">Original Demo</a>
        <a href="https://github.com/OpenMOSS/MOSS-Image" target="_blank" rel="noopener">GitHub &#8599;</a>
      </span>`;
    document.body.prepend(nav);
  }

  // #workflow-app appears when the frontend mounts; poll briefly for it.
  let tries = 0;
  const timer = setInterval(() => {
    const app = document.getElementById('workflow-app');
    if ((app && buildStrip(app)) || ++tries > 75) clearInterval(timer);
  }, 400);
}
"""

CSS = f"""
body, .gradio-container {{
  font-family: 'Inter', ui-sans-serif, system-ui, sans-serif !important;
}}
.gradio-container {{
  background: linear-gradient(180deg, #eef3fc 0%, #ffffff 460px) !important;
}}
.dark .gradio-container {{
  background: linear-gradient(180deg, #0d1526 0%, #0b0f19 460px) !important;
}}

#moss-nav {{
  display: flex;
  align-items: center;
  justify-content: space-between;
  padding: 14px 24px 0;
  max-width: 1200px;
  margin: 0 auto;
  font-size: 14px;
}}
#moss-nav .moss-brand {{
  font-weight: 800;
  font-size: 16px;
  letter-spacing: -0.01em;
  background: linear-gradient(90deg, #3b7fe0, #4aa5f2, #7fe6fb);
  -webkit-background-clip: text;
  background-clip: text;
  color: transparent;
}}
#moss-nav .moss-links a {{
  color: #475569;
  text-decoration: none;
  font-weight: 500;
  margin-left: 22px;
}}
#moss-nav .moss-links a:hover {{ color: #2563eb; }}
.dark #moss-nav .moss-links a {{ color: #94a3b8; }}
.dark #moss-nav .moss-links a:hover {{ color: #7fe6fb; }}

#workflow-app::before {{
  content: "MOSS-Image: Text-to-Image Generation\\Awith Native Layout Control";
  white-space: pre-wrap;
  display: block;
  text-align: center;
  font-weight: 800;
  font-size: clamp(24px, 4.2vw, 40px);
  line-height: 1.15;
  letter-spacing: -0.02em;
  color: #0b1220;
  max-width: 860px;
  margin: 0 auto 12px;
  padding: 158px 16px 26px;
  background: url("{_LOCKUP_LIGHT}") no-repeat top 26px center / min(480px, 78%) auto;
}}
.dark #workflow-app::before {{
  color: #f1f5f9;
  background-image: url("{_LOCKUP_DARK}");
}}

#workflow-app::after {{
  content: "Models by the OpenMOSS Team · MOSS-Image-Base & MOSS-Image-SFT (Apache-2.0) · Rebuilt as a Gradio Workflow";
  display: block;
  text-align: center;
  font-size: 13px;
  color: #64748b;
  padding: 8px 16px 36px;
}}
.dark #workflow-app::after {{ color: #94a3b8; }}

/* Examples strip (injected by the head script; hot-linked from the original Space). */
#moss-examples {{
  max-width: 960px;
  margin: 0 auto 18px;
  padding: 0 16px;
}}
#moss-examples .mx-title {{
  text-align: center;
  font-size: 13px;
  font-weight: 600;
  letter-spacing: 0.02em;
  color: #64748b;
  margin-bottom: 10px;
}}
.dark #moss-examples .mx-title {{ color: #94a3b8; }}
#moss-examples .mx-row {{
  display: flex;
  gap: 10px;
  overflow-x: auto;
  padding-bottom: 6px;
  scrollbar-width: thin;
}}
#moss-examples .mx-item {{
  flex: 0 0 auto;
  height: 128px;
  padding: 0;
  border: 2px solid transparent;
  border-radius: 12px;
  overflow: hidden;
  cursor: pointer;
  background: #e2e8f0;
  transition: border-color 0.15s ease, transform 0.15s ease;
}}
#moss-examples .mx-item:hover {{
  border-color: #2563eb;
  transform: translateY(-2px);
}}
#moss-examples .mx-item img {{
  height: 128px;
  width: auto;
  display: block;
}}

/* Soften the generated form into landing-page cards. */
#workflow-app .form, #workflow-app .panel, #workflow-app .block {{
  border-radius: 14px !important;
}}
#workflow-app button.primary, #workflow-app button.lg {{
  border-radius: 10px !important;
}}
"""

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(primary_hue="blue", neutral_hue="slate"),
        css=CSS,
        head=HEAD,
        js=JS,
    )