from __future__ import annotations

import html

import gradio as gr
from model_unfolder import unfold


DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"

CSS = """
.gradio-container {
  max-width: 1180px !important;
}
.unfold-shell {
  min-height: 72vh;
}
.unfold-frame {
  width: 100%;
  min-height: 860px;
  border: 1px solid #d9e1e8;
  border-radius: 8px;
  background: #f8fafc;
}
.unfold-status {
  color: #385047;
  font-size: 0.95rem;
}
.unfold-note {
  color: #5f7c73;
  font-size: 1rem;
  margin-top: -0.35rem;
}
.unfold-input-panel {
  background: #f6fbf8;
  border: 1px solid #b6ddcb;
  border-radius: 8px;
  padding: 18px;
  box-shadow: 0 1px 3px rgba(15, 23, 42, 0.04);
}
.unfold-input-panel .form {
  border: 0 !important;
  background: transparent !important;
}
.unfold-input-panel textarea,
.unfold-input-panel input {
  border-color: #b6ddcb !important;
}
.unfold-input-panel textarea:focus,
.unfold-input-panel input:focus {
  border-color: #0f6e56 !important;
  box-shadow: 0 0 0 1px #0f6e56 !important;
}
.unfold-button,
.unfold-button button {
  background: #0f6e56 !important;
  border-color: #0f6e56 !important;
  color: #ffffff !important;
}
.unfold-button:hover,
.unfold-button button:hover,
.unfold-button:focus,
.unfold-button button:focus {
  background: #0b5f49 !important;
  border-color: #0b5f49 !important;
}
"""

EMPTY_HTML = """
<div class="unfold-shell" style="display:flex;align-items:center;justify-content:center;min-height:420px;color:#64748b;border:1px dashed #cbd5e1;border-radius:8px;background:#f8fafc;">
  <span>Enter a Hugging Face model ID and unfold it.</span>
</div>
"""


def _clean(value: str | None) -> str:
    return (value or "").strip()


def _as_iframe(document: str) -> str:
    srcdoc = html.escape(document, quote=True)
    return (
        '<iframe class="unfold-frame" '
        'sandbox="allow-scripts" '
        f'srcdoc="{srcdoc}"></iframe>'
    )


def _status_from_diagram(model_id: str, diagram) -> str:
    ir = diagram.to_ir()
    params = ir.get("params") or {}
    bits = [f"Rendered `{model_id}`"]

    total_h = params.get("total_h")
    active_h = params.get("active_h")
    if total_h:
        if params.get("is_sparse") and active_h:
            bits.append(f"~{total_h} params, ~{active_h} active")
        else:
            bits.append(f"~{total_h} params")

    warnings = diagram.warnings
    if warnings:
        bits.append("Warnings: " + "; ".join(warnings))

    return " · ".join(bits)


def render_model(model_id: str | None, hf_token: str | None):
    model_id = _clean(model_id)
    if not model_id:
        raise gr.Error("Enter a Hugging Face model ID.")

    token = _clean(hf_token) or None

    try:
        diagram = unfold(model_id, token=token)
        return (
            _status_from_diagram(model_id, diagram),
            _as_iframe(diagram.to_html(standalone=True)),
        )
    except Exception as exc:
        raise gr.Error("That model card isn't accessible.") from exc


with gr.Blocks(title="Unfold") as demo:
    gr.Markdown("# Unfold")
    gr.Markdown(
        "Generated purely from the Hugging Face config map. No model weights are downloaded.\nBuilt on : [Unfold](https://github.com/SoumilB7/unfold)",
        elem_classes=["unfold-note"],
    )

    with gr.Group(elem_classes=["unfold-input-panel"]):
        with gr.Row(equal_height=False):
            model_id = gr.Textbox(
                label="Model card from HF",
                value=DEFAULT_MODEL,
                placeholder="organization/model-name",
                lines=1,
            )
            hf_token = gr.Textbox(
                label="Your HF token (optional)",
                info="Use models your token can access.",
                type="password",
                placeholder="hf_...",
                lines=1,
            )

        render_button = gr.Button("Unfold", variant="primary", elem_classes=["unfold-button"])
    status = gr.Markdown("", elem_classes=["unfold-status"])
    output = gr.HTML(value=EMPTY_HTML)

    render_button.click(
        fn=render_model,
        inputs=[model_id, hf_token],
        outputs=[status, output],
        show_progress="full",
    )
    model_id.submit(
        fn=render_model,
        inputs=[model_id, hf_token],
        outputs=[status, output],
        show_progress="full",
    )


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