import platform
import shlex
import subprocess
import tarfile
import tempfile
import urllib.request
from pathlib import Path

import gradio as gr
import nvidia


def find_tool(name):
    path = next(p for d in nvidia.__path__ if (p := Path(d) / "cu13/bin" / name).exists())
    path.chmod(0o755)
    return str(path)


def fetch_legacy_nvdisasm(version):
    arch = "linux-sbsa" if platform.machine() == "aarch64" else "linux-x86_64"
    name = f"cuda_nvdisasm-{arch}-{version}-archive"
    path = Path(tempfile.gettempdir()) / name / "nvdisasm"
    if not path.exists():
        url = f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvdisasm/{arch}/{name}.tar.xz"
        with urllib.request.urlopen(url) as r, tarfile.open(fileobj=r, mode="r|xz") as tar:
            member = next(m for m in tar if m.name == f"{name}/bin/nvdisasm")
            path.parent.mkdir(exist_ok=True)
            path.write_bytes(tar.extractfile(member).read())
        path.chmod(0o755)
    return str(path)


TOOLS = {name: find_tool(name) for name in ("cuobjdump", "nvdisasm")}
# nvdisasm 13.x can't run flow analysis (CFGs) on old-format CUDA ELF (OSABI 0x33, pre-CUDA 13); 12.9 can.
LEGACY_NVDISASM = fetch_legacy_nvdisasm("12.9.88")
OUTPUTS = {
    "cuobjdump": ["--dump-sass", "--dump-ptx", "--dump-elf", "--dump-resource-usage", "--dump-elf-symbols",
                  "--list-elf", "--list-ptx"],
    "nvdisasm": ["--print-code", "--output-control-flow-graph", "--output-control-flow-graph-with-basic-blocks"],
}
MODIFIERS = {
    "cuobjdump": [],
    "nvdisasm": ["--print-instruction-encoding", "--print-life-ranges", "--print-raw", "--print-line-info",
                 "--print-line-info-inline"],
}
EXTRA_EXAMPLE = {"cuobjdump": "--function vecadd", "nvdisasm": "--emit-json"}
GRAPHVIZ_HEAD = "".join(f'<script src="https://cdn.jsdelivr.net/npm/{lib}"></script>' for lib in (
    "d3@7.9.0/dist/d3.min.js", "@hpcc-js/wasm@2.20.0/dist/graphviz.umd.js", "d3-graphviz@5.6.0/build/d3-graphviz.min.js"))
GRAPHVIZ_JS = """
const graph = element.querySelector('.graph');
d3.select(graph).graphviz({useWorker: false, zoom: true}).renderDot(props.value, () => {
  const svg = graph.querySelector('svg');
  svg.setAttribute('width', '100%');
  svg.setAttribute('height', Math.min(svg.viewBox.baseVal.height * 4 / 3, 600));
});
"""
EXAMPLES = sorted(str(p) for p in Path(__file__).with_name("examples").glob("*.cubin"))


def version(exe):
    return subprocess.run([exe, "--version"], capture_output=True, text=True).stdout.split("release ")[-1].splitlines()[0]


def is_old_elf(file):
    with open(file, "rb") as f:
        head = f.read(8)
    return head[:4] == b"\x7fELF" and head[7] == 0x33


def executable(tool, file):
    return LEGACY_NVDISASM if tool == "nvdisasm" and is_old_elf(file) else TOOLS[tool]


def run_one(tool, args, file):
    try:
        r = subprocess.run([executable(tool, file), *args, file], capture_output=True, text=True, timeout=60)
    except subprocess.TimeoutExpired:
        return "Timed out after 60s."
    return (r.stdout + r.stderr or "(no output)") + (f"\n[exit code {r.returncode}]" if r.returncode else "")


def run(file, tool, outputs, modifiers, extra):
    if file is None:
        return {"error": "Upload a cubin first."}
    common = [*modifiers, *shlex.split(extra or "")]
    return {" ".join(args): run_one(tool, args, file) for args in ([o, *common] for o in outputs)} or \
           {" ".join(common) or tool: run_one(tool, common, file)}


def select_tool(tool):
    return (gr.CheckboxGroup(OUTPUTS[tool], value=OUTPUTS[tool]),
            gr.CheckboxGroup(MODIFIERS[tool], value=[], visible=bool(MODIFIERS[tool])),
            gr.Textbox(placeholder=EXTRA_EXAMPLE[tool]))


with gr.Blocks(title="cuobjdump / nvdisasm") as demo:
    gr.Markdown("# cuobjdump / nvdisasm\n"
                "Disassemble CUDA cubins, fatbins and ELF objects with NVIDIA's tools ("
                + ", ".join(f"`{t}` {version(exe)}" for t, exe in TOOLS.items())
                + f"; old-format CUDA ELF goes to `nvdisasm` {version(LEGACY_NVDISASM)}).")
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            file = gr.File(label="cubin / fatbin / CUDA ELF", type="filepath")
        with gr.Column(scale=2):
            tool = gr.Radio(list(TOOLS), value="cuobjdump", label="Tool")
            outputs = gr.CheckboxGroup(OUTPUTS["cuobjdump"], value=OUTPUTS["cuobjdump"], label="Outputs (each runs separately)")
            modifiers = gr.CheckboxGroup(MODIFIERS["cuobjdump"], label="Display options (applied to every run)", visible=False)
            with gr.Row(equal_height=True):
                extra = gr.Textbox(label="Extra arguments", placeholder=EXTRA_EXAMPLE["cuobjdump"], scale=4)
                btn = gr.Button("Run", variant="primary", scale=1)
    gr.Examples(EXAMPLES, file)
    results = gr.State({})

    @gr.render(inputs=results)
    def show(results):
        for label, output in results.items():
            if output.startswith("digraph"):
                gr.Markdown(f"`{label}` (scroll to zoom, drag to pan)")
                gr.HTML(output.replace("\t", "    "), head=GRAPHVIZ_HEAD, js_on_load=GRAPHVIZ_JS,
                        html_template='<div class="graph" style="background:white;'
                                      'border:1px solid var(--border-color-primary)"></div>')
                with gr.Accordion(f"{label} (DOT source)", open=False):
                    gr.Code(output, max_lines=60)
            else:
                gr.Code(output, label=label, max_lines=60)

    tool.change(select_tool, tool, [outputs, modifiers, extra])
    (btn.click(lambda: gr.Button("Running…", interactive=False), None, btn)
        .then(run, [file, tool, outputs, modifiers, extra], results)
        .then(lambda: gr.Button("Run", interactive=True), None, btn))

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