import glob
import os
import shutil
import subprocess
import sys
import tempfile
import gradio as gr

ARCHITECTURES = ["sm_70", "sm_75", "sm_80", "sm_86", "sm_89", "sm_90"]

DEFAULT_MODEL = """\
import torch

class Model(torch.nn.Module):
    def forward(self, x, y):
        return torch.relu(x + y)

model = Model()
inputs = (torch.randn(4, 4), torch.randn(4, 4))
torch.onnx.export(model, inputs, 'model.onnx', opset_version=17)
"""


def _run(cmd: list[str], label: str, timeout: int) -> tuple[str, bool]:
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        if r.returncode != 0:
            return f"{label} failed:\n\n{r.stderr.strip()}", False
        return r.stdout, True
    except subprocess.TimeoutExpired:
        return f"{label} timed out ({timeout}s).", False


def compile_pipeline(code: str, arch: str) -> tuple[str, str, str, str, str | None, str | None]:
    if not code.strip():
        return "Error: no code provided.", "", "", "", None, None
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            src        = os.path.join(tmpdir, "model_def.py")
            onnx_path  = os.path.join(tmpdir, "model.onnx")
            mlir_path  = os.path.join(tmpdir, "model.mlir")
            vmfb_path  = os.path.join(tmpdir, "model.vmfb")
            exec_dir   = os.path.join(tmpdir, "exec")
            cubin_path = os.path.join(tmpdir, "kernel.cubin")
            os.makedirs(exec_dir)

            with open(src, "w") as f:
                f.write(code)

            py_cmd      = [sys.executable, src]
            import_cmd  = ["iree-import-onnx", onnx_path, "-o", mlir_path]
            compile_cmd = [
                "iree-compile", mlir_path,
                "--iree-hal-target-backends=cuda",
                f"--iree-cuda-target={arch}",
                f"--iree-hal-dump-executable-files-to={exec_dir}",
                "-o", vmfb_path,
            ]

            cmds = "\n".join(" ".join(c) for c in [py_cmd, import_cmd, compile_cmd])

            # PyTorch → ONNX
            r = subprocess.run(py_cmd, capture_output=True, text=True, timeout=60, cwd=tmpdir)
            if r.returncode != 0:
                err = f"PyTorch export failed:\n\n{r.stderr.strip()}"
                return cmds, err, err, err, None, None

            # ONNX → MLIR
            out, ok = _run(import_cmd, "iree-import-onnx", 60)
            if not ok:
                return cmds, out, out, out, None, None
            mlir_text = open(mlir_path).read() if os.path.exists(mlir_path) else "(no MLIR output)"

            # MLIR → VMFB, PTX dumped to exec_dir
            out, ok = _run(compile_cmd, "iree-compile", 120)
            if not ok:
                return cmds, mlir_text, out, out, None, None

            ptx_files = sorted(glob.glob(os.path.join(exec_dir, "**", "*.ptx"), recursive=True))
            if not ptx_files:
                return cmds, mlir_text, "No PTX files produced.", "No PTX to disassemble.", None, None

            ptx_text = "\n\n".join(open(p).read() for p in ptx_files)

            # PTX → CUBIN → SASS (first PTX file)
            ptxas_cmd     = ["ptxas", f"--gpu-name={arch}", ptx_files[0], "-o", cubin_path]
            cuobjdump_cmd = ["cuobjdump", "--dump-sass", cubin_path]
            cmds += "\n" + " ".join(ptxas_cmd) + "\n" + " ".join(cuobjdump_cmd)

            _, ok = _run(ptxas_cmd, "ptxas", 30)
            sass_text = _run(cuobjdump_cmd, "cuobjdump", 30)[0] if ok else "ptxas failed; SASS unavailable."

            ptx_tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".ptx")
            ptx_tmp.write(ptx_text.encode())
            ptx_tmp.close()

            cubin_tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".cubin")
            cubin_tmp.close()
            if ok and os.path.exists(cubin_path):
                shutil.copy(cubin_path, cubin_tmp.name)
                cubin_out = cubin_tmp.name
            else:
                cubin_out = None

            return cmds, mlir_text, ptx_text, sass_text, ptx_tmp.name, cubin_out

    except FileNotFoundError as e:
        return "", "", f"Tool not found: {e}\n\nEnsure iree-compiler is installed.", "", None, None
    except Exception as e:
        return "", "", f"Unexpected error: {e}", "", None, None


with gr.Blocks(title="PyTorch → SASS Explorer") as demo:
    gr.Markdown(
        "# PyTorch → SASS Explorer\n"
        "Write a PyTorch model that exports to `model.onnx`, pick a GPU architecture, "
        "and trace the full compilation pipeline: PyTorch → ONNX → IREE MLIR → PTX → SASS."
    )
    with gr.Row():
        with gr.Column(scale=2):
            code_input = gr.Code(
                label="PyTorch Model", language="python",
                value=DEFAULT_MODEL, lines=20,
            )
        with gr.Column(scale=1):
            arch = gr.Dropdown(label="GPU Architecture", choices=ARCHITECTURES, value="sm_80")
            btn  = gr.Button("Compile", variant="primary")

    cmds_out = gr.Code(label="Commands", language="shell", lines=4, interactive=False)
    mlir_out = gr.Code(label="MLIR",     language=None, lines=10, max_lines=10, interactive=False)

    with gr.Row():
        ptx_out  = gr.Code(label="PTX",  language=None, lines=10, max_lines=10, interactive=False)
        sass_out = gr.Code(label="SASS", language=None, lines=10, max_lines=10, interactive=False)

    with gr.Row():
        ptx_file  = gr.File(label="Download PTX",  interactive=False)
        cubin_file = gr.File(label="Download CUBIN", interactive=False)

    btn.click(compile_pipeline, [code_input, arch], [cmds_out, mlir_out, ptx_out, sass_out, ptx_file, cubin_file])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)