import pprint
import subprocess
import gradio as gr
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",
        "./",
        "--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}")

    pprint.pprint(arguments)

    # 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, ""):
        if "torch.distributed.nn.jit.instantiator" in ansi_line:
            continue
        # stream process output
        print(ansi_line, end="")
        # append line to ansi text
        ansi_text += ansi_line
        # convert ansi to html
        html_text = ansi2html_converter.convert(ansi_text)
        # stream html output
        yield html_text

    return html_text