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


def get_functions(file):

    with tempfile.TemporaryDirectory() as TEMP_DIR:

        subprocess.run(
            f"/ghidra/support/analyzeHeadless {TEMP_DIR} Project -import {file} -postscript /home/user/app/scripts/dump_functions.py {TEMP_DIR}/funcs.json",
            shell=True,
        )

        json_funcs = json.load(open(f"{TEMP_DIR}/funcs.json"))

        return json_funcs


with gr.Blocks() as demo:

    state = gr.State()

    intro = gr.Markdown(
        """
    # DIRTY-Ghidra Inference Demo
    Welcome!  This is a demo of DIRTY-Ghidra, a tool that predict names and types for variables for Ghidra's decompiler.

    To get started, upload a binary.
    """
    )

    file_widget = gr.File(label="Executable file")

    with gr.Column(visible=False) as col:
        # output = gr.Textbox("Output")

        gr.Markdown(
            """
        Great, you selected an executable!  Now pick the function you would like
        to analyze.  Simple functions (without variables) will probably fail, so
        you may have to try a few before you find one that works.
        """
        )

        fun_dropdown = gr.Dropdown(
            label="Select a function", choices=["Woohoo!"], interactive=True
        )

        gr.Markdown(
            """
        Below you can find some information.
                    """
        )

        with gr.Row(visible=True) as result:
            disassembly = gr.Textbox(
                label="Disassembly", value="Please wait...", lines=20
            )
            original_decompile = gr.Textbox(
                label="Original Decompilation", value="Please wait...", lines=20
            )
            decompile = gr.Textbox(
                label="Renamed and retyped Decompilation",
                value="Please wait...",
                lines=20,
            )
            model_output = gr.Textbox(
                label="Model Output", value="Please wait...", lines=4
            )
            # with gr.Column():
            #     clazz = gr.Label()
            # interpret_button = gr.Button("Interpret (very slow)")
            # interpretation = gr.components.Interpretation(disassembly)

    def file_change_fn(file):

        if file is None:
            return {col: gr.update(visible=False), state: {"file": None}}
        else:

            try:
                progress = gr.Progress()
                progress(
                    0,
                    desc=f"Analyzing binary  {os.path.basename(file.name)} with Ghidra...",
                )
                fun_data = get_functions(file.name)
                # print(fun_data)

                addrs = [
                    (f"{name} ({hex(int(addr))})", int(addr))
                    for addr, name in fun_data.items()
                ]
            except Exception as e:
                raise gr.Error(f"Unable to analyze binary with Ghidra: {e}")

            return {
                col: gr.Column(visible=True),
                fun_dropdown: gr.Dropdown(choices=addrs, value=addrs[0][1]),
                state: {"file": file},
            }

    def function_change_fn(selected_fun, state, progress=gr.Progress()):

        # disassembly_str = fun_data[int(selected_fun, 16)].decode("utf-8")
        # load_results = model.fn(disassembly_str)
        # top_k = {e['label']: e['confidence'] for e in load_results['confidences']}

        with tempfile.TemporaryDirectory() as TEMP_DIR:

            print(selected_fun)

            progress(0, desc=f"Running DIRTY Ghidra on {hex(selected_fun)}...")

            try:
                subprocess.run(
                    f"/ghidra/support/analyzeHeadless {TEMP_DIR} Project -import {state['file'].name} -postscript /DIRTY/scripts/DIRTY_infer.py {TEMP_DIR}/funcs.json {selected_fun}",
                    shell=True,
                )

                json_info = json.load(open(f"{TEMP_DIR}/funcs.json"))

                if "exception" in json_info:
                    raise gr.Error(f"DIRTY Ghidra failed: {json_info['exception']}")

            except Exception as e:
                raise gr.Error(f"Unable to run DIRTY Ghidra: {e}")

            #print(json_info)

            return {
                disassembly: gr.Textbox(value=json_info["disassembly"]),
                original_decompile: gr.Textbox(value=json_info["original_decompile"]),
                decompile: gr.Textbox(value=json_info["decompile"]),
                model_output: gr.Textbox(value=json_info["model_output"]),
            }

    # Need to put intro as output to get progress to work!
    file_widget.change(
        file_change_fn, file_widget, outputs=[intro, state, col, fun_dropdown]
    )

    fun_dropdown.change(
        function_change_fn,
        inputs=[fun_dropdown, state],
        outputs=[disassembly, original_decompile, decompile, model_output],
    )

# spaces only shows stderr..
os.dup2(sys.stdout.fileno(), sys.stderr.fileno())

demo.queue()
demo.launch(server_name="0.0.0.0", server_port=7860)