"""
ComfyUI Metadata Viewer
========================
Drag & drop a file produced by ComfyUI (PNG, MP3, FLAC, MP4...) and it
displays, as readable tables:
  - the model and LoRAs used
  - the main parameters (steps, cfg, sampler, seed, duration...)
  - the final prompt actually used, and the "dynamic prompt" { | }
    template that generated it, when there is one.
"""

import html
import re
import traceback

import gradio as gr
import spaces

from comfyui_meta import analyze, MetadataNotFound, StructuredView

BRACE_RE = re.compile(r"(\{[^{}]*\})")


@spaces.GPU(duration=1)
def _zerogpu_startup_probe():
    """This app never needs a GPU — it only reads metadata out of a
    file. But Spaces running on the free ZeroGPU tier require at least
    one @spaces.GPU-decorated function to exist, or the Space refuses
    to start. This function exists only to satisfy that check; it is
    never called."""
    return None


def _highlight_template(text: str) -> str:
    """Escapes HTML and highlights {a | b | c} blocks."""
    escaped = html.escape(text)
    # re-apply the pattern on the escaped text (braces are unaffected
    # by html.escape)
    parts = BRACE_RE.split(escaped)
    out = []
    for part in parts:
        if part.startswith("{") and part.endswith("}"):
            out.append(f'<span class="wildcard">{part}</span>')
        else:
            out.append(part)
    return "".join(out)


def _models_table(view: StructuredView):
    rows = [[m.kind, m.name, m.detail, m.node] for m in view.models]
    if not rows:
        rows = [["—", "No model/LoRA detected", "", ""]]
    return rows


def _params_table(view: StructuredView):
    rows = [[p.node, p.param, p.value] for p in view.params]
    if not rows:
        rows = [["—", "No parameter detected", ""]]
    return rows


def _prompts_html(view: StructuredView) -> str:
    if not view.prompts:
        return "<p><em>No prompt detected in this file.</em></p>"

    cards = []
    for p in view.prompts:
        template_block = ""
        if p.template:
            template_block = f"""
            <div class="field-box template-box">
                <div class="label">🎲 Dynamic template (raw)</div>
                <div class="text mono">{_highlight_template(p.template)}</div>
            </div>
            """
        cards.append(f"""
        <div class="prompt-card">
            <div class="card-title">{html.escape(p.node)} — {html.escape(p.role)}</div>
            {template_block}
            <div class="field-box resolved-box">
                <div class="label">✅ Final prompt (resolved)</div>
                <div class="text">{html.escape(p.resolved)}</div>
            </div>
        </div>
        """)

    style = """
    <style>
      .prompt-card { margin-bottom: 20px; }
      .prompt-card .card-title { font-weight: 600; margin-bottom: 8px; opacity: 0.85; }
      /* Same "gray frame" look as the Gradio file drop zone, reusing the
         theme's own tokens so it matches light/dark mode automatically. */
      .field-box {
          border: var(--block-border-width, 1px) solid var(--block-border-color, rgba(128,128,128,0.35));
          border-radius: var(--block-radius, 8px);
          background: var(--block-background-fill, transparent);
          padding: 12px 14px;
          margin-bottom: 10px;
      }
      .field-box .label { font-size: 0.85em; text-transform: uppercase;
                           letter-spacing: 0.03em; opacity: 0.6; margin-bottom: 6px; }
      .field-box .text { white-space: pre-wrap; line-height: 1.4; }
      .field-box .mono { font-family: ui-monospace, SFMono-Regular, Menlo, monospace;
                          font-size: 0.92em; }
      .wildcard { background: rgba(63, 92, 168, 0.28); border-radius: 4px;
                  padding: 0 2px; font-weight: normal; }
    </style>
    """
    return style + "\n".join(cards)


def _warnings_html(view: StructuredView) -> str:
    if not view.warnings:
        return ""
    items = "".join(f"<li>{html.escape(w)}</li>" for w in view.warnings)
    return f"<div style='opacity:0.75; font-size:0.9em;'>⚠️ <ul>{items}</ul></div>"


def process_file(file_path):
    if not file_path:
        return [], [], "", "", "Drop a file above to get started."

    try:
        view = analyze(file_path)
    except MetadataNotFound as e:
        return [], [], "", "", f"❌ {e}"
    except Exception:
        return [], [], "", "", (
            "❌ Unexpected error during analysis:\n\n" + traceback.format_exc()
        )

    status = "✅ ComfyUI metadata found and analyzed."
    return (
        _models_table(view),
        _params_table(view),
        _prompts_html(view),
        _warnings_html(view),
        status,
    )


with gr.Blocks(title="ComfyUI Metadata Viewer") as demo:
    gr.Markdown(
        "# 🔍 ComfyUI Metadata Viewer\n"
        "Drag & drop a file generated by ComfyUI (PNG, MP3, FLAC, MP4...) "
        "to see the model, LoRAs, parameters and prompt (with its original "
        "dynamic prompt, if there is one)."
    )

    file_input = gr.File(
        label="ComfyUI file (drag & drop)",
        file_types=[".png", ".webp", ".jpg", ".jpeg",
                    ".mp3", ".flac", ".m4a", ".ogg",
                    ".mp4", ".mov", ".webm", ".mkv"],
        type="filepath",
    )
    clear_btn = gr.Button("🗑️ Clear", size="sm", variant="secondary")

    status_box = gr.Markdown("Drop a file above to get started.")

    with gr.Tab("🧩 Model & LoRAs"):
        models_df = gr.Dataframe(
            headers=["Type", "Name", "Detail", "Node"],
            datatype=["str", "str", "str", "str"],
            row_count=(0, "dynamic"),
            wrap=True,
        )

    with gr.Tab("⚙️ Main Parameters"):
        params_df = gr.Dataframe(
            headers=["Node", "Parameter", "Value"],
            datatype=["str", "str", "str"],
            row_count=(0, "dynamic"),
            wrap=True,
        )

    with gr.Tab("📝 Prompt & Dynamic Prompt"):
        prompts_html = gr.HTML()

    warnings_html = gr.HTML()

    file_input.change(
        fn=process_file,
        inputs=[file_input],
        outputs=[models_df, params_df, prompts_html, warnings_html, status_box],
    )

    def _clear_all():
        return None, [], [], "", "", "Drop a file above to get started."

    clear_btn.click(
        fn=_clear_all,
        inputs=[],
        outputs=[file_input, models_df, params_df, prompts_html, warnings_html, status_box],
    )

if __name__ == "__main__":
    demo.launch()