"""
Screenshot -> React, with Gemma 3 27B IT VLM + a LoRA adapter.

Upload a UI screenshot; the model writes a React component, and the app renders it live
in a sandboxed iframe. The training data (Reubencf/frontend-react-dataset) leans on
Tailwind, framer-motion and lucide-react, so the preview resolves those imports to UMD
globals instead of stripping them.
"""

import html
import json
import os
import re
import tempfile
import threading

import gradio as gr
import spaces
import torch
from peft import PeftModel
from transformers import AutoProcessor, Gemma3ForConditionalGeneration, TextIteratorStreamer

# `unsloth/gemma-3-27b-it` is an ungated mirror, so the Space runs without a token.
# Set BASE_MODEL_ID=google/gemma-3-27b-it (plus an HF_TOKEN secret) to use the canonical repo.
#
# Both are Space variables, so dropping back to the 4B pair is a settings change rather than
# a code change: BASE_MODEL_ID=unsloth/gemma-3-4b-it and
# ADAPTER_ID=Reubencf/gemma-3-4b-it-vlm-react-screenshot-to-code.
BASE_MODEL_ID = os.environ.get("BASE_MODEL_ID", "unsloth/gemma-3-27b-it")
ADAPTER_ID = os.environ.get(
    "ADAPTER_ID", "Reubencf/gemma-3-27b-it-vlm-react-screenshot-to-code"
)
MODEL_LABEL = os.environ.get("MODEL_LABEL", "Gemma 3 27B IT VLM")

# 27B in bfloat16 is ~55GB of weights, which does not fit ZeroGPU's default 48GB slice.
# `xlarge` is a full RTX Pro 6000 Blackwell at 96GB — and costs 2x quota per call.
GPU_SIZE = os.environ.get("GPU_SIZE", "xlarge")

DEFAULT_INSTRUCTION = (
    "Convert this screenshot into a single self-contained React component.\n"
    "Use Tailwind CSS utility classes for styling. Return only the component code."
)

# --------------------------------------------------------------------------------------
# Model
# --------------------------------------------------------------------------------------

processor = AutoProcessor.from_pretrained(BASE_MODEL_ID)

# Build entirely on CPU, then move once.
#
# torch_device="cpu" is load-bearing on ZeroGPU: PEFT picks the adapter's load device via
# infer_device(), which returns "cuda" because torch.cuda.is_available() reports True at
# startup even though no GPU is attached yet. safetensors then materialises straight onto
# CUDA and dies with "No CUDA GPUs are available". Module .to("cuda") is patched by
# `spaces` and replayed when a GPU attaches; a direct safetensors CUDA load is not.
#
# The same constraint rules out loading 4-bit to save memory: bitsandbytes quantises during
# from_pretrained and needs a real GPU to do it, which does not exist at module scope here.
model = Gemma3ForConditionalGeneration.from_pretrained(
    BASE_MODEL_ID, dtype=torch.bfloat16, device_map="cpu"
)
model = PeftModel.from_pretrained(model, ADAPTER_ID, torch_device="cpu")

# Deliberately not merged. merge_and_unload() would rebuild all 62 layers' q/v projections
# on CPU, briefly holding a second copy of weights that are already ~55GB. The adapter is
# rank 4 on two projections, so leaving it live costs a pair of tiny matmuls per layer.
model = model.eval().to("cuda")
DEVICE = "cuda"

_eot = processor.tokenizer.convert_tokens_to_ids("<end_of_turn>")
EOS_IDS = [i for i in {processor.tokenizer.eos_token_id, _eot} if i is not None]

# --------------------------------------------------------------------------------------
# Pulling code out of the raw generation
# --------------------------------------------------------------------------------------

FENCED = re.compile(
    r"```[ \t]*(?:jsx|tsx|js|javascript|typescript|react)?[ \t]*\n(.*?)```", re.S | re.I
)
UNCLOSED = re.compile(
    r"```[ \t]*(?:jsx|tsx|js|javascript|typescript|react)?[ \t]*\n(.*)\Z", re.S | re.I
)


def extract_code(raw: str) -> str:
    """Return just the code from a generation that may be wrapped in prose and fences."""
    raw = raw.strip()
    blocks = FENCED.findall(raw)
    if blocks:
        return max(blocks, key=len).strip()
    truncated = UNCLOSED.search(raw)  # generation cut off mid-block
    if truncated:
        return truncated.group(1).strip()
    return raw


# --------------------------------------------------------------------------------------
# Rewriting ESM into something that runs in a plain page
# --------------------------------------------------------------------------------------

IMPORT_STMT = re.compile(
    r"^[ \t]*import\s+(?:(?P<clause>[\s\S]*?)\s+from\s+)?"
    r"['\"](?P<mod>[^'\"]+)['\"][ \t]*;?[ \t]*$",
    re.M,
)
REQUIRE = re.compile(
    r"^[ \t]*(?:const|let|var)\s+(?P<clause>[\s\S]*?)\s*=\s*require\("
    r"['\"](?P<mod>[^'\"]+)['\"]\)[ \t]*;?[ \t]*$",
    re.M,
)

EXPORT_DEFAULT_DECL = re.compile(
    r"^[ \t]*export\s+default\s+(?=(?:async\s+)?(?:function|class)\b)", re.M
)
EXPORT_DEFAULT_DECL_NAME = re.compile(
    r"^[ \t]*export\s+default\s+(?:async\s+)?(?:function|class)\s+([A-Za-z_$][\w$]*)", re.M
)
EXPORT_DEFAULT_ANON_FN = re.compile(
    r"^[ \t]*export\s+default\s+((?:async\s+)?function\s*\()", re.M
)
EXPORT_DEFAULT_NAME = re.compile(
    r"^[ \t]*export\s+default\s+([A-Za-z_$][\w$]*)[ \t]*;?[ \t]*$", re.M
)
EXPORT_DEFAULT_EXPR = re.compile(r"^[ \t]*export\s+default\s+", re.M)
EXPORT_NAMED = re.compile(r"^[ \t]*export\s+(?=(?:const|let|var|function|class|async)\b)", re.M)
EXPORT_LIST = re.compile(r"^[ \t]*export\s*\{[^}]*\}[ \t]*;?[ \t]*$", re.M)

DECL_NAME = re.compile(
    r"^[ \t]*(?:export\s+default\s+|export\s+)?"
    r"(?:async\s+)?(?:function|class|const|let|var)\s+([A-Z][\w$]*)",
    re.M,
)

IDENT = re.compile(r"^[A-Za-z_$][\w$]*$")
ANON_NAME = "__GeneratedComponent"


def _parse_clause(clause: str):
    """Split an import clause into (default_name, namespace_name, [(imported, local)])."""
    default_name = namespace = None
    named: list[tuple[str, str]] = []

    braces = re.search(r"\{([\s\S]*)\}", clause)
    if braces:
        for part in braces.group(1).split(","):
            part = part.strip()
            if not part:
                continue
            part = re.sub(r"^type\s+", "", part)  # `{ type Foo }`
            alias = re.match(r"^([A-Za-z_$][\w$]*)\s+as\s+([A-Za-z_$][\w$]*)$", part)
            if alias:
                named.append((alias.group(1), alias.group(2)))
            elif IDENT.match(part):
                named.append((part, part))
        clause = clause[: braces.start()] + clause[braces.end():]

    for part in clause.split(","):
        part = part.strip()
        if not part:
            continue
        ns = re.match(r"^\*\s+as\s+([A-Za-z_$][\w$]*)$", part)
        if ns:
            namespace = ns.group(1)
        elif IDENT.match(part):
            default_name = part

    return default_name, namespace, named


def _rebind(clause: str | None, mod: str, bound: set[str]) -> str:
    """Turn one import/require into `const` bindings backed by CDN globals."""
    if clause is None:  # side-effect import, e.g. `import './styles.css'`
        return ""
    clause = clause.strip()
    if clause.startswith("type "):  # `import type { Props } from ...`
        return ""

    default_name, namespace, named = _parse_clause(clause)
    mod_js = json.dumps(mod)
    lines = []

    if namespace:
        lines.append(f"const {namespace} = window.__ns({mod_js});")
        bound.add(namespace)
    if default_name:
        lines.append(f"const {default_name} = window.__default({mod_js});")
        bound.add(default_name)
    if named:
        spec = ", ".join(f"{src}: {dst}" if src != dst else src for src, dst in named)
        lines.append(f"const {{ {spec} }} = window.__ns({mod_js});")
        bound.update(dst for _, dst in named)

    return "\n".join(lines)


def to_browser_module(code: str) -> tuple[str, str]:
    """Rewrite ESM-flavoured React so Babel-standalone can run it. -> (code, component)."""
    bound: set[str] = set()
    code = IMPORT_STMT.sub(lambda m: _rebind(m.group("clause"), m.group("mod"), bound), code)
    code = REQUIRE.sub(lambda m: _rebind(m.group("clause"), m.group("mod"), bound), code)
    code = EXPORT_LIST.sub("", code)

    # Resolve the default export's name before `export` prefixes are stripped — afterwards
    # there's no way to tell the default component from any other exported one.
    name = None
    named_default = EXPORT_DEFAULT_NAME.search(code)
    if named_default:
        name = named_default.group(1)
        code = EXPORT_DEFAULT_NAME.sub("", code)
    else:
        decl = EXPORT_DEFAULT_DECL_NAME.search(code)
        if decl:
            name = decl.group(1)

    if name is None and EXPORT_DEFAULT_ANON_FN.search(code):
        code = EXPORT_DEFAULT_ANON_FN.sub(rf"const {ANON_NAME} = \1", code, count=1)
        name = ANON_NAME

    code = EXPORT_DEFAULT_DECL.sub("", code)
    code = EXPORT_NAMED.sub("", code)

    # Anything still `export default <expr>` is an anonymous arrow/class expression.
    if EXPORT_DEFAULT_EXPR.search(code):
        code = EXPORT_DEFAULT_EXPR.sub(f"const {ANON_NAME} = ", code, count=1)
        name = ANON_NAME

    if name is None:
        # Skip the `const React = window.__default(...)` lines we just generated.
        name = next(
            (m.group(1) for m in DECL_NAME.finditer(code) if m.group(1) not in bound),
            "App",
        )

    # Bare `useState(...)` with no import resolves off `window` (see the preview preamble),
    # rather than a `const` here that could redeclare something the code already binds.
    return code.strip(), name


# --------------------------------------------------------------------------------------
# Live preview
# --------------------------------------------------------------------------------------

PREVIEW_DOC = """<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<style>
  html, body { margin: 0; background: #fff; }
  #err {
    display: none; margin: 0; padding: 14px 18px; white-space: pre-wrap;
    font: 12.5px/1.6 ui-monospace, SFMono-Regular, Menlo, monospace;
    color: #b91c1c; background: #fef2f2; border-bottom: 1px solid #fecaca;
  }
</style>

<!-- Registered before the CDN tags so a library that fails to initialise is visible. -->
<script>
  function showError(msg) {
    var el = document.getElementById('err');
    if (!el) return;
    el.style.display = 'block';
    el.textContent = (el.textContent ? el.textContent + '\\n\\n' : '') + String(msg);
  }
  window.__showError = showError;
  window.addEventListener('error', function (e) { showError(e.message); });
</script>

<script src="https://cdn.tailwindcss.com/3.4.16"></script>
<script crossorigin src="https://unpkg.com/react@18/umd/react.production.min.js"></script>
<script crossorigin src="https://unpkg.com/react-dom@18/umd/react-dom.production.min.js"></script>
<script src="https://unpkg.com/framer-motion@11/dist/framer-motion.js"></script>
<!-- The lucide-react UMD build throws on load; the vanilla `lucide` build ships the same
     PascalCase icon names as plain node data, which we turn into components below. -->
<script src="https://unpkg.com/lucide@0.544.0/dist/umd/lucide.js"></script>
<script src="https://unpkg.com/@babel/standalone@7/babel.min.js"></script>
</head>
<body>
<pre id="err"></pre>
<div id="root"></div>

<script>
  // Props that only make sense to framer-motion / lucide — never forward them to the DOM.
  var DROP = ['initial','animate','exit','transition','variants','layout','layoutId',
              'whileHover','whileTap','whileFocus','whileDrag','whileInView','viewport',
              'drag','dragConstraints','custom','onAnimationComplete','absoluteStrokeWidth'];

  function passthrough(tag) {
    var C = React.forwardRef(function (props, ref) {
      var p = {};
      for (var k in props) { if (DROP.indexOf(k) === -1 && k !== 'children') p[k] = props[k]; }
      p.ref = ref;
      return React.createElement(tag, p, props.children);
    });
    C.displayName = 'stub(' + String(tag) + ')';
    return C;
  }

  // If framer-motion fails to load, `motion.div` still degrades to a plain <div>.
  var motionFallback = new Proxy({}, {
    get: function (_t, tag) { return typeof tag === 'string' ? passthrough(tag) : undefined; }
  });

  // Build a real <svg> component from lucide's icon node data, matching lucide-react's API.
  var iconCache = {};
  function camel(a) { return a.replace(/-([a-z])/g, function (_m, c) { return c.toUpperCase(); }); }

  function lucideIcon(name) {
    if (iconCache[name] !== undefined) return iconCache[name];
    var L = window.lucide || {};
    var node = (L.icons && L.icons[name]) || L[name];
    if (!Array.isArray(node)) { iconCache[name] = null; return null; }

    var C = React.forwardRef(function (props, ref) {
      props = props || {};
      var size = props.size == null ? 24 : props.size;
      var attrs = { ref: ref, xmlns: 'http://www.w3.org/2000/svg', width: size, height: size,
                    viewBox: '0 0 24 24', fill: 'none', stroke: 'currentColor',
                    strokeWidth: props.strokeWidth == null ? 2 : props.strokeWidth,
                    strokeLinecap: 'round', strokeLinejoin: 'round' };
      for (var k in props) {
        if (k !== 'size' && k !== 'strokeWidth' && k !== 'children' &&
            DROP.indexOf(k) === -1) attrs[k] = props[k];
      }
      var kids = node.map(function (child, i) {
        var tag = child[0], raw = child[1] || {}, p = { key: i };
        for (var a in raw) p[camel(a)] = raw[a];
        return React.createElement(tag, p);
      });
      return React.createElement('svg', attrs, kids);
    });
    C.displayName = name;
    iconCache[name] = C;
    return C;
  }

  var MODULES = {
    'react': window.React,
    'react-dom': window.ReactDOM,
    'react-dom/client': window.ReactDOM,
    'framer-motion': window.Motion,
    'motion/react': window.Motion
  };

  window.__ns = function (mod) {
    var base = MODULES[mod] || {};
    return new Proxy(base, {
      get: function (t, k) {
        if (k in t) return t[k];
        if (typeof k !== 'string') return undefined;
        if (k === 'motion') return motionFallback;
        if (/^[A-Z]/.test(k)) {
          var icon = lucideIcon(k);
          if (icon) return icon;
          // Unknown capitalised binding -> render children, so layout survives.
          return passthrough(React.Fragment);
        }
        return undefined;
      }
    });
  };

  window.__default = function (mod) {
    var m = MODULES[mod];
    if (m && m.default) return m.default;
    return m || window.__ns(mod);
  };

  // Let bare `useState(...)` work when the code never imported it. These land on `window`
  // instead of being declared in the eval scope, so code that *does* declare them shadows
  // these rather than colliding with a duplicate `const`.
  ['useState','useEffect','useRef','useMemo','useCallback','useReducer','useContext',
   'useLayoutEffect','createContext','Fragment','memo','forwardRef','cloneElement',
   'Children','createRef','Suspense'].forEach(function (k) {
    if (window[k] === undefined && React[k] !== undefined) window[k] = React[k];
  });

  class __ErrorBoundary extends React.Component {
    constructor(p) { super(p); this.state = { crashed: false }; }
    static getDerivedStateFromError() { return { crashed: true }; }
    componentDidCatch(err) { showError(err && err.stack ? err.stack : String(err)); }
    render() { return this.state.crashed ? null : this.props.children; }
  }
  window.__ErrorBoundary = __ErrorBoundary;
</script>

<script type="text/plain" id="__src">/*__CODE__*/

/*__MOUNT__*/
</script>

<script>
(function () {
  var src = document.getElementById('__src').textContent;

  function compile(code) {
    return Babel.transform(code, {
      filename: 'Component.tsx',
      sourceType: 'unambiguous',
      presets: [
        ['typescript', { isTSX: true, allExtensions: true }],
        ['react', { runtime: 'classic' }]
      ]
    }).code;
  }

  // The model transcribes on-screen text like "<1 Years" literally, and a bare `<` is
  // invalid in JSX. Only applied after a real failure, so valid code is never touched.
  function repair(code) {
    return code.replace(/<<+/g, '&lt;').replace(/<(?=\\s*\\d)/g, '&lt;');
  }

  var out, firstErr = null;
  try {
    out = compile(src);
  } catch (e) {
    firstErr = e && e.message ? e.message : String(e);
    try {
      out = compile(repair(src));
      showError('The generated code had invalid JSX; the preview auto-repaired it to render.\\n\\n'
                + firstErr);
    } catch (e2) {
      showError('Could not compile the generated code:\\n\\n' + firstErr);
      return;
    }
  }
  try {
    (0, eval)(out);
  } catch (e) {
    showError(e && e.stack ? e.stack : String(e));
  }
})();
</script>
</body>
</html>
"""

MOUNT = """window.ReactDOM.createRoot(document.getElementById('root')).render(
  window.React.createElement(
    window.__ErrorBoundary, null, window.React.createElement(/*__NAME__*/)
  )
);"""

# Code that already calls ReactDOM...render() mounts itself; don't mount it twice.
SELF_MOUNTING = re.compile(r"ReactDOM(?:Client)?\s*\.\s*(?:createRoot|render)\s*\(")

PLACEHOLDER = """
<div style="display:flex;align-items:center;justify-content:center;height:520px;
            border:1px dashed var(--border-color-primary,#d0d5dd);border-radius:10px;
            color:var(--body-text-color-subdued,#667085);
            font:14px/1.5 var(--font,system-ui,sans-serif);text-align:center;padding:24px;">
  Upload a screenshot and hit <strong>&nbsp;Generate&nbsp;</strong> — the rendered React lands here.
</div>
"""


def build_preview(code: str) -> str:
    if not code.strip():
        return PLACEHOLDER

    body, name = to_browser_module(code)
    mount = "" if SELF_MOUNTING.search(body) else MOUNT.replace("/*__NAME__*/", name)

    # A literal </script> in the code would close the host <script> tag early.
    body = body.replace("</script", "<\\/script")

    doc = PREVIEW_DOC.replace("/*__CODE__*/", body).replace("/*__MOUNT__*/", mount)

    return (
        f'<iframe sandbox="allow-scripts" srcdoc="{html.escape(doc, quote=True)}" '
        'style="width:100%;height:660px;border:1px solid var(--border-color-primary,#d0d5dd);'
        'border-radius:10px;background:#fff;"></iframe>'
    )


def save_jsx(code: str) -> str | None:
    if not code.strip():
        return None
    fd, path = tempfile.mkstemp(suffix=".jsx", prefix="component_")
    with os.fdopen(fd, "w", encoding="utf-8") as fh:
        fh.write(code)
    return path


# --------------------------------------------------------------------------------------
# Generation
# --------------------------------------------------------------------------------------


def gpu_duration(image, instruction, max_new_tokens, temperature, top_p, repetition_penalty):
    """Ask for GPU time proportional to the token budget.

    27B decodes far slower than 4B did, and a flat 120s cut long components off mid-JSX.
    Asking for less when the budget is small keeps queue priority up for everyone else.
    """
    return int(45 + int(max_new_tokens) * 0.075)


@spaces.GPU(duration=gpu_duration, size=GPU_SIZE)
def stream_generation(image, instruction, max_new_tokens, temperature, top_p, repetition_penalty):
    messages = [{
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": instruction.strip() or DEFAULT_INSTRUCTION},
        ],
    }]

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(DEVICE, dtype=torch.bfloat16)

    streamer = TextIteratorStreamer(
        processor.tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=int(max_new_tokens),
        do_sample=temperature > 0,
        temperature=float(temperature) if temperature > 0 else None,
        top_p=float(top_p),
        repetition_penalty=float(repetition_penalty),
        eos_token_id=EOS_IDS,
    )

    worker = threading.Thread(target=model.generate, kwargs=kwargs)
    worker.start()

    acc = ""
    pending = 0
    for chunk in streamer:
        acc += chunk
        pending += len(chunk)
        if pending >= 32:  # throttle: token-by-token updates flood the socket
            pending = 0
            yield acc
    worker.join()
    yield acc


def run(image, instruction, max_new_tokens, temperature, top_p, repetition_penalty):
    if image is None:
        raise gr.Error("Upload a screenshot first.")

    raw = ""
    for raw in stream_generation(
        image, instruction, max_new_tokens, temperature, top_p, repetition_penalty
    ):
        yield extract_code(raw), gr.update(), gr.update(), raw

    code = extract_code(raw)
    yield code, build_preview(code), gr.update(value=save_jsx(code), visible=True), raw


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------

CSS = """
.app-title { text-align: center; }
footer { visibility: hidden; }
"""

with gr.Blocks(title="Screenshot → React", theme=gr.themes.Soft(), css=CSS) as demo:
    gr.Markdown(
        f"""
        <div class="app-title">

        # 🖼️ → ⚛️ Screenshot to React

        Upload a screenshot of a UI. [{MODEL_LABEL}](https://huggingface.co/{BASE_MODEL_ID})
        with the [`{ADAPTER_ID.split('/')[-1]}`](https://huggingface.co/{ADAPTER_ID}) LoRA
        adapter writes a React component, and it renders live below.

        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=4):
            image_in = gr.Image(type="pil", label="Screenshot", height=300)
            instruction = gr.Textbox(label="Instruction", value=DEFAULT_INSTRUCTION, lines=3)
            generate_btn = gr.Button("Generate React", variant="primary", size="lg")

            with gr.Accordion("Generation settings", open=False):
                max_new_tokens = gr.Slider(256, 4096, value=2048, step=128, label="Max new tokens")
                temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature (0 = greedy)")
                top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
                repetition_penalty = gr.Slider(
                    1.0, 1.5, value=1.1, step=0.01, label="Repetition penalty",
                    info="This adapter can fall into loops of nested <div>s. Raise this if that happens.",
                )

        with gr.Column(scale=6):
            with gr.Tabs():
                with gr.Tab("Preview"):
                    preview = gr.HTML(value=PLACEHOLDER)
                with gr.Tab("React code"):
                    code_out = gr.Code(language="javascript", label="Component", lines=26)
                    download = gr.DownloadButton("Download .jsx", visible=False)
                with gr.Tab("Raw output"):
                    raw_out = gr.Textbox(
                        label="Unparsed model output", lines=26, show_copy_button=True
                    )

    gr.Markdown(
        """
        <sub>The preview runs the generated code in a sandboxed iframe with React 18, Babel
        standalone, Tailwind 3.4, framer-motion and lucide-react. `import` statements are rewritten
        onto those globals, and the default export is mounted behind an error boundary — compile and
        runtime errors show up in a red banner. The adapter was trained for only 26 steps, so treat
        output as a scaffold to edit, not finished code.<br><br>
        The 27B base is ~55&nbsp;GB, so the first request after the Space wakes has to load it before
        anything generates. Each run also uses a full-size ZeroGPU slice, which draws 2× daily quota
        — signed-out visitors get roughly one generation per day.</sub>
        """
    )

    generate_btn.click(
        fn=run,
        inputs=[image_in, instruction, max_new_tokens, temperature, top_p, repetition_penalty],
        outputs=[code_out, preview, download, raw_out],
    )

if __name__ == "__main__":
    demo.queue(max_size=16).launch()