import subprocess
import gradio as gr
import pandas as pd
from ansi2html import Ansi2HTMLConverter

ansi2html_converter = Ansi2HTMLConverter(inline=True)


def run_benchmark(kwargs):
    for key, value in kwargs.copy().items():
        if key.label == "experiment_name":
            experiment_name = value
            kwargs.pop(key)
        elif key.label == "model":
            model = value
            kwargs.pop(key)
        elif key.label == "task":
            task = value
            kwargs.pop(key)
        elif key.label == "device":
            device = value
            kwargs.pop(key)
        elif key.label == "backend":
            backend = value
            kwargs.pop(key)
        elif key.label == "benchmark":
            benchmark = value
            kwargs.pop(key)
        else:
            continue

    arguments = [
        "optimum-benchmark",
        "--config-dir",
        "./configs",
        "--config-name",
        "base_config",
        f"task={task}",
        f"model={model}",
        f"device={device}",
        f"backend={backend}",
        f"benchmark={benchmark}",
        f"experiment_name={experiment_name}",
    ]

    for component, value in kwargs.items():
        if f"{backend}." in component.label or f"{benchmark}." in component.label:
            label = component.label.replace(f"{backend}.", "backend.").replace(f"{benchmark}.", "benchmark.")

            if isinstance(component, gr.Dataframe):
                for sub_key, sub_value in zip(component.headers, value[0]):
                    arguments.append(f"++{label}.{sub_key}={sub_value}")
            else:
                arguments.append(f"{label}={value}")

    command = "<br>".join(arguments)
    html_text = f"<h3>Running command:</h3>{command}"
    yield gr.update(value=html_text), gr.update(interactive=False), gr.update(visible=False)

    # stream subprocess output
    process = subprocess.Popen(
        arguments,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True,
    )

    ansi_text = ""
    for ansi_line in iter(process.stdout.readline, ""):
        # stream process output to stdout
        print(ansi_line, end="")

        # skip torch.distributed.nn.jit.instantiator messages
        if "torch.distributed.nn.jit.instantiator" in ansi_line:
            continue
        # if the last message is a download message (contains "Downloading ") then remove it and replace it with a new one
        if "Downloading " in ansi_text and "Downloading " in ansi_line:
            ansi_text = ansi_text.split("\n")[:-2]
            print(ansi_text)
            ansi_text.append(ansi_line)
            ansi_text = "\n".join(ansi_text)
        else:
            # append line to ansi text
            ansi_text += ansi_line
        # convert ansi to html
        html_text = ansi2html_converter.convert(ansi_text)
        # stream html output to gradio
        yield gr.update(value=html_text), gr.update(interactive=False), gr.update(visible=False)

    if process.returncode != 0:
        table = pd.read_csv(f"runs/{experiment_name}/{benchmark}_results.csv", index_col=0)
        table_update = gr.update(visible=True, value={"headers": list(table.columns), "data": table.values.tolist()})
    else:
        table_update = gr.update(visible=False)

    yield gr.update(value=html_text), gr.update(interactive=True), table_update
    return