"""cuda-kernels Space: two things in one app.

Tab 1 -- "Push to Hub": push any local folder to a Hugging Face Hub repo
(model/dataset/space) without the CLI.

Tab 2 -- "Sliding-window attention": a from-scratch CUDA kernel (JIT-
compiled with torch.utils.cpp_extension.load_inline) implementing
Longformer-style local attention with online softmax, benchmarked against
dense PyTorch attention, run live on this Space's GPU allocation.

Tab 3 -- "Kernel fusion compiler": fuses a small elementwise op graph
(y = gelu(x*w + b)) into one generated CUDA kernel via the `fusion_compiler`
package, compiles it, and benchmarks it against the naive 3-kernel-launch
version.

Tabs 2 and 3 need this Space's hardware set to ZeroGPU (or another CUDA
GPU) -- see README for the one-line command to flip that.
"""

from __future__ import annotations

import os
import sys
import time
import traceback

import gradio as gr
import pandas as pd
from huggingface_hub import HfApi, whoami

sys.path.insert(0, os.path.dirname(__file__))

try:
    import spaces
    _HAS_SPACES = True
except ImportError:
    _HAS_SPACES = False

    class _NoOpGPU:
        def __call__(self, *args, **kwargs):
            def decorator(fn):
                return fn
            return decorator

    class _SpacesShim:
        GPU = _NoOpGPU()

    spaces = _SpacesShim()

import torch

from fusion_compiler.graph import Graph
from fusion_compiler.fuser import fuse
from fusion_compiler.codegen import generate_cuda_source


def _use_torch_bundled_cuda_toolchain():
    """Point CUDA_HOME/PATH/LD_LIBRARY_PATH at the CUDA toolchain that ships
    as a pip dependency of the installed torch wheel (nvidia-cuda-nvcc-cu12,
    nvidia-cuda-runtime-cu12, etc.) instead of whatever CUDA toolkit happens
    to be installed system-wide in the container.

    Without this, torch.utils.cpp_extension.load_inline finds `nvcc` on
    PATH (which may be a different, newer CUDA toolkit than the one torch
    itself was built against) and compiles against that instead -- the
    resulting extension links against a libcudartXX.so version torch's own
    bundled runtime doesn't provide, which fails at import time with
    something like "libcudart.so.13: cannot open shared object file".
    Forcing both compile-time and runtime to use torch's own bundled CUDA
    libs keeps them consistent.

    We search for the actual nvcc binary rather than assuming a fixed path
    under nvidia-cuda-nvcc-cu12's package layout -- that assumption was
    wrong on at least one container (the bin/ dir existed but no nvcc file
    inside it), which produced a much worse failure ("nvcc: not found")
    than the problem we were trying to fix. If no working nvcc binary is
    found anywhere under the nvidia package tree, we leave PATH/CUDA_HOME
    untouched entirely and let torch fall back to its own default
    detection, rather than pointing at a path known not to exist.

    Returns a short diagnostic string describing what was actually found,
    so callers can surface it in error messages instead of us having to
    guess blind across another round of "it still doesn't work" reports.
    """
    try:
        import nvidia
    except ImportError:
        return "nvidia pip namespace package not importable (nvidia-cuda-runtime-cu12 / nvidia-cuda-nvcc-cu12 not installed?)"
    import glob
    from importlib import metadata as importlib_metadata

    nvidia_root = os.path.dirname(nvidia.__file__)
    subdirs = sorted(os.listdir(nvidia_root)) if os.path.isdir(nvidia_root) else []

    try:
        nvcc_pkg_version = importlib_metadata.version("nvidia-cuda-nvcc-cu12")
    except importlib_metadata.PackageNotFoundError:
        nvcc_pkg_version = None

    lib_dirs = glob.glob(os.path.join(nvidia_root, "*", "lib"))
    if lib_dirs:
        os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + [os.environ.get("LD_LIBRARY_PATH", "")])

    nvcc_candidates = [
        p for p in glob.glob(os.path.join(nvidia_root, "**", "nvcc"), recursive=True)
        if os.path.isfile(p) and os.access(p, os.X_OK)
    ]
    if nvcc_candidates:
        nvcc_dir = os.path.dirname(nvcc_candidates[0])
        os.environ["PATH"] = nvcc_dir + ":" + os.environ.get("PATH", "")
        os.environ["CUDA_HOME"] = os.path.dirname(nvcc_dir)
        return f"using nvcc at {nvcc_candidates[0]}"

    return (
        f"no nvcc binary found under {nvidia_root} (subdirs: {subdirs}); "
        f"nvidia-cuda-nvcc-cu12 pip package version: {nvcc_pkg_version!r}; "
        f"lib_dirs found: {lib_dirs}; falling back to system nvcc, which may "
        f"be a mismatched CUDA version"
    )


# ===========================================================================
# Tab 1: push-to-hub
# ===========================================================================


def _get_api(token: str) -> HfApi:
    token = (token or os.environ.get("HF_TOKEN") or "").strip()
    if not token:
        raise ValueError("No token provided and HF_TOKEN is not set in the environment.")
    return HfApi(token=token)


def check_token(token: str) -> str:
    try:
        api = _get_api(token)
        info = whoami(token=api.token)
        return f"✅ Authenticated as **{info['name']}**"
    except Exception as e:
        return f"❌ {e}"


def push_folder(
    token: str,
    repo_id: str,
    repo_type: str,
    local_folder: str,
    private: bool,
    commit_message: str,
    path_in_repo: str,
):
    log_lines = []

    def log(msg: str):
        log_lines.append(msg)
        return "\n".join(log_lines)

    try:
        if not repo_id or "/" not in repo_id:
            yield log("❌ repo_id must look like `username/repo-name`.")
            return
        if not local_folder or not os.path.isdir(local_folder):
            yield log(f"❌ '{local_folder}' is not a directory on this machine.")
            return

        api = _get_api(token)
        yield log(f"Authenticated. Ensuring repo '{repo_id}' ({repo_type}) exists...")

        api.create_repo(
            repo_id=repo_id,
            repo_type=repo_type,
            private=private,
            exist_ok=True,
        )
        yield log(f"Repo ready: https://huggingface.co/{'' if repo_type == 'model' else repo_type + 's/'}{repo_id}")

        yield log(f"Uploading contents of '{local_folder}'...")

        url = api.upload_folder(
            repo_id=repo_id,
            repo_type=repo_type,
            folder_path=local_folder,
            path_in_repo=path_in_repo or ".",
            commit_message=commit_message or "Push from hf-push-ui",
            ignore_patterns=["__pycache__/*", "*.pyc", ".git/*", "build/*", "*.so"],
        )
        yield log(f"✅ Done. {url}")

    except Exception:
        yield log("❌ Failed:\n" + traceback.format_exc())


# ===========================================================================
# Tab 2: sliding-window attention kernel
# ===========================================================================

_ATTN_CUDA_SRC = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#define MAX_HEAD_DIM 128

__global__ void sliding_window_attn_kernel(
    const float* __restrict__ Q, const float* __restrict__ K, const float* __restrict__ V,
    float* __restrict__ O, int batch_heads, int seq_len, int head_dim, int window, float scale)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_heads * seq_len;
    if (idx >= total) return;

    int bh = idx / seq_len;
    int q = idx % seq_len;

    const float* Qp = Q + ((long)bh * seq_len + q) * head_dim;
    const float* Kbase = K + (long)bh * seq_len * head_dim;
    const float* Vbase = V + (long)bh * seq_len * head_dim;
    float* Op = O + ((long)bh * seq_len + q) * head_dim;

    int k_lo = max(0, q - window);
    int k_hi = min(seq_len - 1, q + window);

    float m = -1e30f, l = 0.f;
    float acc[MAX_HEAD_DIM];
    for (int d = 0; d < head_dim; ++d) acc[d] = 0.f;

    for (int k = k_lo; k <= k_hi; ++k) {
        const float* Kp = Kbase + (long)k * head_dim;
        float score = 0.f;
        for (int d = 0; d < head_dim; ++d) score += Qp[d] * Kp[d];
        score *= scale;

        float m_new = fmaxf(m, score);
        float corr = expf(m - m_new);
        float p = expf(score - m_new);
        l = l * corr + p;

        const float* Vp = Vbase + (long)k * head_dim;
        for (int d = 0; d < head_dim; ++d) acc[d] = acc[d] * corr + p * Vp[d];
        m = m_new;
    }

    float inv_l = (l > 0.f) ? 1.f / l : 0.f;
    for (int d = 0; d < head_dim; ++d) Op[d] = acc[d] * inv_l;
}

torch::Tensor sliding_window_attention_naive(torch::Tensor Q, torch::Tensor K, torch::Tensor V, int64_t window, double scale) {
    TORCH_CHECK(Q.is_cuda() && K.is_cuda() && V.is_cuda(), "Q, K, V must be CUDA tensors");
    TORCH_CHECK(Q.scalar_type() == torch::kFloat32, "expected float32 tensors");
    auto Qc = Q.contiguous();
    auto Kc = K.contiguous();
    auto Vc = V.contiguous();

    int batch = Qc.size(0), heads = Qc.size(1), seq_len = Qc.size(2), head_dim = Qc.size(3);
    TORCH_CHECK(head_dim <= 128, "demo kernel supports head_dim <= 128");

    auto O = torch::empty_like(Qc);
    int batch_heads = batch * heads;
    int total = batch_heads * seq_len;
    int threads = 128;
    int blocks = (total + threads - 1) / threads;

    sliding_window_attn_kernel<<<blocks, threads>>>(
        Qc.data_ptr<float>(), Kc.data_ptr<float>(), Vc.data_ptr<float>(), O.data_ptr<float>(),
        batch_heads, seq_len, head_dim, (int)window, (float)scale);

    return O;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("sliding_window_attention_naive", &sliding_window_attention_naive, "sliding window attention (naive, demo)");
}
"""

_attn_module = None


def _load_cuda_extension(name: str, cuda_source: str):
    """Write `cuda_source` to a real .cu file and build it with
    torch.utils.cpp_extension.load(). This torch build requires ninja
    unconditionally (no use_ninja kwarg, no distutils fallback), so we rely
    on packages.txt installing a real "ninja-build" apt package rather than
    the pip "ninja" wheel, whose bundled binary was failing with exit code
    127 in this container. We use load() instead of load_inline for the
    real on-disk build directory, which is easier to debug if compilation
    fails again.
    """
    import tempfile
    from torch.utils.cpp_extension import load

    toolchain_diag = _use_torch_bundled_cuda_toolchain()
    build_dir = os.path.join(tempfile.gettempdir(), f"{name}_build")
    os.makedirs(build_dir, exist_ok=True)
    src_path = os.path.join(build_dir, f"{name}.cu")
    with open(src_path, "w") as f:
        f.write(cuda_source)

    try:
        return load(
            name=name,
            sources=[src_path],
            verbose=False,
            build_directory=build_dir,
        )
    except Exception as e:
        raise RuntimeError(f"{e}\n\n[CUDA toolchain diagnostic] {toolchain_diag}") from e


def _get_attn_module():
    global _attn_module
    if _attn_module is None:
        _attn_module = _load_cuda_extension("sliding_window_attention_demo", _ATTN_CUDA_SRC)
    return _attn_module


def _dense_reference(Q, K, V, window, scale):
    seq_len = Q.shape[2]
    idx = torch.arange(seq_len, device=Q.device)
    mask = (idx[:, None] - idx[None, :]).abs() <= window
    scores = torch.einsum("bhqd,bhkd->bhqk", Q, K) * scale
    scores = scores.masked_fill(~mask[None, None, :, :], float("-inf"))
    attn = torch.softmax(scores, dim=-1)
    return torch.einsum("bhqk,bhkd->bhqd", attn, V)


_EMPTY_CHART = pd.DataFrame({"kind": [], "time_ms": []})


@spaces.GPU(duration=120)
def run_attention_kernel(seq_len: int, window: int, heads: int, head_dim: int):
    try:
        if not torch.cuda.is_available():
            return "No CUDA GPU visible to this Space. Set hardware to ZeroGPU (or another GPU tier) in Settings.", _EMPTY_CHART

        mod = _get_attn_module()

        torch.manual_seed(0)
        device = "cuda"
        Q = torch.randn(1, heads, seq_len, head_dim, device=device, dtype=torch.float32)
        K = torch.randn(1, heads, seq_len, head_dim, device=device, dtype=torch.float32)
        V = torch.randn(1, heads, seq_len, head_dim, device=device, dtype=torch.float32)
        scale = 1.0 / (head_dim ** 0.5)

        for _ in range(3):
            out_kernel = mod.sliding_window_attention_naive(Q, K, V, window, scale)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(10):
            out_kernel = mod.sliding_window_attention_naive(Q, K, V, window, scale)
        torch.cuda.synchronize()
        t_kernel = (time.perf_counter() - t0) / 10 * 1000

        for _ in range(3):
            out_ref = _dense_reference(Q, K, V, window, scale)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(10):
            out_ref = _dense_reference(Q, K, V, window, scale)
        torch.cuda.synchronize()
        t_ref = (time.perf_counter() - t0) / 10 * 1000

        max_err = (out_kernel - out_ref).abs().max().item()
        gpu_name = torch.cuda.get_device_name(0)
        correct = max_err < 1e-3
        faster = t_kernel < t_ref

        verdict = (
            f"**{t_ref / t_kernel:.2f}x faster** than dense PyTorch attention"
            if faster else
            f"**{t_kernel / t_ref:.2f}x slower** than dense PyTorch attention on this run"
        )

        summary = (
            f"### {'✅' if correct else '⚠️'} Correctness: max error {max_err:.1e} "
            f"vs. dense reference ({'PASS' if correct else 'CHECK'})\n\n"
            f"### {'🚀' if faster else '🐢'} Speed: {verdict}\n\n"
            f"GPU: {gpu_name}  ·  seq_len={seq_len}, window=±{window}, heads={heads}, head_dim={head_dim}\n\n"
        )

        if not faster:
            summary += (
                "**Why slower here?** This demo kernel is written for clarity, not "
                "speed: one CUDA thread per query row, doing a plain serial loop over "
                "the window with no shared-memory tiling or warp-level reduction. "
                "PyTorch's dense attention path is backed by heavily hand-tuned cuBLAS "
                "kernels that have had years of optimization — beating that with a toy "
                "kernel isn't the point here. The [production tiled kernel]"
                "(https://github.com/data-geek-astronomy/long-context-attention-kernels) "
                "adds shared-memory blocking and online softmax across a query *tile* "
                "(not one row per thread), which is what actually closes this gap, "
                "and where the algorithmic O(n·w) vs O(n²) advantage starts to show up "
                "as real wall-clock speedup at long sequence lengths."
            )

        chart = pd.DataFrame(
            [
                {"kind": "custom CUDA kernel (this demo)", "time_ms": t_kernel},
                {"kind": "dense PyTorch attention", "time_ms": t_ref},
            ]
        )
        return summary, chart
    except Exception:
        return "ERROR:\n" + traceback.format_exc(), _EMPTY_CHART


# ===========================================================================
# Tab 3: kernel fusion compiler
# ===========================================================================


def _build_example_graph():
    g = Graph()
    g.add("mul", ["x", "w"], "t1")
    g.add("add", ["t1", "b"], "t2")
    g.add("gelu", ["t2"], "y")
    return g


@spaces.GPU(duration=120)
def run_fusion_kernel(n_elements: int):
    try:
        if not torch.cuda.is_available():
            return "No CUDA GPU visible to this Space. Set hardware to ZeroGPU (or another GPU tier) in Settings.", _EMPTY_CHART

        _use_torch_bundled_cuda_toolchain()
        from fusion_compiler.jit import compile_group

        graph = _build_example_graph()
        groups = fuse(graph)

        fn = compile_group(groups[0], kernel_name="fused_demo_kernel_live")

        device = "cuda"
        torch.manual_seed(0)
        x = torch.randn(n_elements, device=device, dtype=torch.float32)
        w = torch.randn(n_elements, device=device, dtype=torch.float32)
        b = torch.randn(n_elements, device=device, dtype=torch.float32)

        for _ in range(3):
            out_fused = fn(x, w, b)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(20):
            out_fused = fn(x, w, b)
        torch.cuda.synchronize()
        t_fused = (time.perf_counter() - t0) / 20 * 1000

        def unfused(x, w, b):
            t1 = x * w
            t2 = t1 + b
            return torch.nn.functional.gelu(t2, approximate="tanh")

        for _ in range(3):
            out_ref = unfused(x, w, b)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(20):
            out_ref = unfused(x, w, b)
        torch.cuda.synchronize()
        t_unfused = (time.perf_counter() - t0) / 20 * 1000

        max_err = (out_fused - out_ref).abs().max().item()
        gpu_name = torch.cuda.get_device_name(0)
        correct = max_err < 1e-3
        faster = t_fused < t_unfused

        verdict = (
            f"**{t_unfused / t_fused:.2f}x faster** than the naive 3-kernel-launch version"
            if faster else
            f"**{t_fused / t_unfused:.2f}x slower** on this run"
        )

        summary = (
            f"### {'✅' if correct else '⚠️'} Correctness: max error {max_err:.1e} "
            f"vs. PyTorch reference ({'PASS' if correct else 'CHECK'})\n\n"
            f"### {'🚀' if faster else '🐢'} Speed: {verdict}\n\n"
            f"GPU: {gpu_name}  ·  n_elements={n_elements:,}\n\n"
            f"1 auto-generated kernel launch instead of 3, fewer round trips to "
            f"global memory for the intermediate values — see the bar chart above "
            f"for the exact counts, computed straight from the fused graph."
        )

        chart = pd.DataFrame(
            [
                {"kind": "fused (1 launch)", "time_ms": t_fused},
                {"kind": "unfused (3 launches)", "time_ms": t_unfused},
            ]
        )
        return summary, chart
    except Exception:
        return "ERROR:\n" + traceback.format_exc(), _EMPTY_CHART


def preview_fused_source():
    graph = _build_example_graph()
    groups = fuse(graph)
    return generate_cuda_source(groups[0], kernel_name="fused_demo_kernel")


# ===========================================================================
# Showcase data -- illustrative complexity chart for attention (labeled as
# such, not presented as measured numbers), and an exactly-computed memory-
# traffic comparison for the fusion compiler (derived from the real graph,
# not synthetic).
# ===========================================================================

def _attention_complexity_chart(window: int = 64):
    """Relative compute cost, O(n^2) dense vs O(n*w) sparse. Arbitrary units,
    clearly not wall-clock timing -- for that, use 'Run live on GPU' below.
    Plotted as log10(cost) so both curves fit on one readable axis."""
    import math

    seq_lens = [512, 1024, 2048, 4096, 8192, 16384, 32768]
    rows = []
    for n in seq_lens:
        dense_cost = n * n
        sparse_cost = n * (2 * window + 1)
        rows.append({"seq_len": n, "kind": "dense O(n²)", "log10_relative_cost": math.log10(dense_cost)})
        rows.append({"seq_len": n, "kind": f"sliding-window O(n·w), w={window}", "log10_relative_cost": math.log10(sparse_cost)})
    return pd.DataFrame(rows)


def _fusion_memory_chart():
    """Exact counts derived from the actual graph/fuser, not made up: how
    many global-memory array reads+writes each version does."""
    graph = _build_example_graph()
    groups = fuse(graph)

    unfused_touches = sum(len(n.inputs) + 1 for n in graph)  # each node: reads + 1 write
    fused_touches = len(groups[0].inputs) + 1  # one read per external input, one final write

    df = pd.DataFrame(
        [
            {"metric": "kernel launches", "unfused": len(graph), "fused": len(groups)},
            {"metric": "global memory array touches", "unfused": unfused_touches, "fused": fused_touches},
        ]
    )
    return df.melt(id_vars="metric", var_name="version", value_name="count")


_ARCH_DIAGRAM_URL = "https://raw.githubusercontent.com/data-geek-astronomy/long-context-attention-kernels/main/assets/architecture.svg"

# ===========================================================================
# UI
# ===========================================================================

with gr.Blocks(title="cuda-kernels") as demo:
    gr.Markdown("# CUDA Kernels")

    with gr.Tab("Overview"):
        gr.Markdown(
            f"""
            Custom CUDA kernel work, compiled and run live on a free
            Hugging Face **ZeroGPU** allocation.

            - [long-context-attention-kernels](https://github.com/data-geek-astronomy/long-context-attention-kernels) —
              Longformer-style sliding-window + global-token attention, tiled kernel with online softmax.
            - [cuda-fusion-compiler](https://github.com/data-geek-astronomy/cuda-fusion-compiler) —
              a small compiler that auto-fuses chains of elementwise ops into a single generated CUDA kernel.
            - [cuda-ml-kernels](https://github.com/data-geek-astronomy/cuda-ml-kernels) —
              Flash Attention v2, fused LayerNorm+GELU, INT8 quantization, tiled GEMM.

            Open the tabs above to see each kernel's design and run it for real.
            """
        )
        gr.HTML(f'<img src="{_ARCH_DIAGRAM_URL}" alt="sliding window + global token attention architecture" style="width:100%;max-width:1000px;border-radius:8px;">')

    with gr.Tab("Sliding-window attention"):
        gr.Markdown(
            "Longformer-style local attention: a hand-written CUDA kernel (online softmax, "
            "one query per thread) vs. dense masked PyTorch attention."
        )
        gr.HTML(f'<img src="{_ARCH_DIAGRAM_URL}" alt="sparsity pattern" style="width:100%;max-width:700px;border-radius:8px;">')

        gr.Markdown("**Why this scales better** — algorithmic cost, dense O(n²) vs. sliding-window O(n·w). *Illustrative complexity comparison, not measured latency* — click below for real numbers.")
        complexity_plot = gr.LinePlot(
            _attention_complexity_chart(),
            x="seq_len", y="log10_relative_cost", color="kind",
            x_title="sequence length", y_title="log10(relative compute), arbitrary units",
            height=320,
        )

        gr.Markdown("**Run it for real** — compiles the kernel with `nvcc` and benchmarks it against dense PyTorch attention on this Space's GPU.")
        with gr.Row():
            seq_len_in = gr.Slider(128, 4096, value=1024, step=128, label="seq_len")
            window_in = gr.Slider(8, 512, value=64, step=8, label="window (±)")
        with gr.Row():
            heads_in = gr.Slider(1, 16, value=4, step=1, label="heads")
            head_dim_in = gr.Slider(16, 128, value=64, step=16, label="head_dim")
        attn_btn = gr.Button("Run on GPU", variant="primary")
        attn_summary = gr.Markdown()
        attn_chart = gr.BarPlot(
            _EMPTY_CHART, x="kind", y="time_ms",
            x_title="", y_title="measured time (ms, lower is better)",
            height=280,
        )
        attn_btn.click(
            run_attention_kernel,
            inputs=[seq_len_in, window_in, heads_in, head_dim_in],
            outputs=[attn_summary, attn_chart],
        )

    with gr.Tab("Kernel fusion compiler"):
        gr.Markdown(
            "`y = gelu(x*w + b)`: 3 elementwise ops, auto-fused into 1 CUDA kernel by "
            "`fusion_compiler`'s graph fuser + codegen — no hand-written fusion."
        )

        gr.Markdown("**Exact counts from the fuser** (not estimates — computed from the real graph):")
        fusion_plot = gr.BarPlot(
            _fusion_memory_chart(),
            x="metric", y="count", color="version",
            x_title="", y_title="count",
            height=300,
        )

        preview_btn = gr.Button("Preview generated source (no GPU needed)")
        preview_out = gr.Code(label="Generated CUDA", language="cpp")
        preview_btn.click(preview_fused_source, outputs=[preview_out])

        gr.Markdown("**Run it for real** — compiles the fused kernel and benchmarks it against the naive 3-kernel-launch PyTorch version on this Space's GPU.")
        n_elements_in = gr.Slider(1024, 1 << 22, value=1 << 20, step=1024, label="n_elements")
        fusion_btn = gr.Button("Compile + run on GPU", variant="primary")
        fusion_summary = gr.Markdown()
        fusion_chart = gr.BarPlot(
            _EMPTY_CHART, x="kind", y="time_ms",
            x_title="", y_title="measured time (ms, lower is better)",
            height=280,
        )
        fusion_btn.click(run_fusion_kernel, inputs=[n_elements_in], outputs=[fusion_summary, fusion_chart])

    with gr.Tab("Push to Hub (utility)"):
        gr.Markdown(
            """
            Utility tab, unrelated to the kernels above: point this at any
            local project folder and push it to a Hugging Face Hub repo.
            """
        )
        with gr.Row():
            token_box = gr.Textbox(
                label="HF Access Token",
                type="password",
                placeholder="hf_... (or leave blank to use HF_TOKEN env var)",
            )
            check_btn = gr.Button("Check token", scale=0)
        token_status = gr.Markdown()
        check_btn.click(check_token, inputs=[token_box], outputs=[token_status])

        with gr.Row():
            repo_id_box = gr.Textbox(label="Repo ID", placeholder="your-username/cuda-fusion-compiler")
            repo_type_box = gr.Radio(["model", "dataset", "space"], value="model", label="Repo type")

        local_folder_box = gr.Textbox(
            label="Local folder to push",
            placeholder="/absolute/path/to/long-context-attention-kernels",
        )
        path_in_repo_box = gr.Textbox(
            label="Destination path in repo (optional)",
            placeholder="leave blank to push into repo root",
        )
        with gr.Row():
            private_box = gr.Checkbox(label="Private repo", value=False)
            commit_msg_box = gr.Textbox(label="Commit message", value="Push from hf-push-ui")

        push_btn = gr.Button("🚀 Push to Hugging Face", variant="primary")
        log_box = gr.Textbox(label="Log", lines=10, interactive=False)

        push_btn.click(
            push_folder,
            inputs=[token_box, repo_id_box, repo_type_box, local_folder_box, private_box, commit_msg_box, path_in_repo_box],
            outputs=[log_box],
        )

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