import os
import time
from huggingface_hub import create_repo, whoami
import gradio as gr
from config_store import (
    get_process_config,
    get_inference_config,
    get_onnxruntime_config,
    get_openvino_config,
    get_pytorch_config,
    get_ipex_config,
)
from optimum_benchmark.launchers.base import Launcher  # noqa
from optimum_benchmark.backends.openvino.utils import TASKS_TO_OVMODEL
from optimum_benchmark.backends.transformers_utils import TASKS_TO_MODEL_LOADERS
from optimum_benchmark.backends.onnxruntime.utils import TASKS_TO_ORTMODELS
from optimum_benchmark.backends.ipex.utils import TASKS_TO_IPEXMODEL
from optimum_benchmark import (
    BenchmarkConfig,
    PyTorchConfig,
    OVConfig,
    ORTConfig,
    IPEXConfig,
    ProcessConfig,
    InferenceConfig,
    Benchmark,
)
from optimum_benchmark.logging_utils import setup_logging


DEVICE = "cpu"
LAUNCHER = "process"
SCENARIO = "inference"
BACKENDS = ["onnxruntime", "openvino", "pytorch", "ipex"]
MODELS = [
    "hf-internal-testing/tiny-random-bert",
    "google-bert/bert-base-uncased",
    "openai-community/gpt2",
]
TASKS = (
    set(TASKS_TO_OVMODEL.keys())
    & set(TASKS_TO_ORTMODELS.keys())
    & set(TASKS_TO_IPEXMODEL.keys())
    & set(TASKS_TO_MODEL_LOADERS.keys())
)


def run_benchmark(kwargs, oauth_token: gr.OAuthToken):
    if oauth_token.token is None:
        return "You must be logged in to use this space"

    username = whoami(oauth_token.token)["name"]
    create_repo(
        f"{username}/benchmarks",
        token=oauth_token.token,
        repo_type="dataset",
        exist_ok=True,
    )

    configs = {
        "process": {},
        "inference": {},
        "onnxruntime": {},
        "openvino": {},
        "pytorch": {},
        "ipex": {},
    }

    for key, value in kwargs.items():
        if key.label == "model":
            model = value
        elif key.label == "task":
            task = value
        elif key.label == "backends":
            backends = value
        elif "." in key.label:
            backend, argument = key.label.split(".")
            configs[backend][argument] = value
        else:
            continue

    configs["process"] = ProcessConfig(**configs.pop("process"))
    configs["inference"] = InferenceConfig(**configs.pop("inference"))

    configs["onnxruntime"] = ORTConfig(
        task=task,
        model=model,
        device=DEVICE,
        **configs["onnxruntime"],
    )
    configs["openvino"] = OVConfig(
        task=task,
        model=model,
        device=DEVICE,
        **configs["openvino"],
    )
    configs["pytorch"] = PyTorchConfig(
        task=task,
        model=model,
        device=DEVICE,
        **configs["pytorch"],
    )
    configs["ipex"] = IPEXConfig(
        task=task,
        model=model,
        device=DEVICE,
        **configs["ipex"],
    )

    md_output = (
        f"<h3>Running benchmark for model {model} on task {task} with {backends}</h3>"
    )

    yield md_output

    timestamp = time.strftime("%Y-%m-%d-%H-%M-%S")

    for backend in backends:
        md_output += f"<br>🚀 Launching benchmark for {backend}"
        yield md_output

        try:
            benchmark_name = f"{timestamp}/{backend}"
            benchmark_config = BenchmarkConfig(
                name=benchmark_name,
                backend=configs[backend],
                launcher=configs[LAUNCHER],
                scenario=configs[SCENARIO],
            )
            benchmark_config.push_to_hub(
                repo_id=f"{username}/benchmarks",
                subfolder=benchmark_name,
                token=oauth_token.token,
            )
            benchmark_report = Benchmark.launch(benchmark_config)
            benchmark_report.push_to_hub(
                repo_id=f"{username}/benchmarks",
                subfolder=benchmark_name,
                token=oauth_token.token,
            )
            benchmark = Benchmark(config=benchmark_config, report=benchmark_report)
            benchmark.push_to_hub(
                repo_id=f"{username}/benchmarks",
                subfolder=benchmark_name,
                token=oauth_token.token,
            )

            md_output += (
                f"<br>✅ Benchmark for {backend} backend completed successfully"
            )
            yield md_output
        except Exception as e:
            md_output += (
                f"<br>❌ Error while running benchmark for {backend} backend: {e}"
            )
            yield md_output


def build_demo():
    with gr.Blocks() as demo:
        # add login button
        gr.LoginButton(min_width=250)

        # add image
        gr.Markdown(
            """<img src="https://huggingface.co/spaces/optimum/optimum-benchmark-ui/resolve/main/huggy_bench.png" style="display: block; margin-left: auto; margin-right: auto; width: 30%;">"""
        )

        # title text
        gr.Markdown(
            "<h1 style='text-align: center'>🤗 Optimum-Benchmark Interface 🏋️</h1>"
        )

        # explanation text
        gr.HTML(
            "<h3 style='text-align: center'>"
            "Zero code Gradio interface of "
            "<a href='https://github.com/huggingface/optimum-benchmark.git'>"
            "Optimum-Benchmark"
            "</a>"
            "<br>"
            "</h3>"
        )

        model = gr.Dropdown(
            label="model",
            choices=MODELS,
            value=MODELS[0],
            info="Model to run the benchmark on.",
        )
        task = gr.Dropdown(
            label="task",
            choices=TASKS,
            value="feature-extraction",
            info="Task to run the benchmark on.",
        )
        backends = gr.CheckboxGroup(
            interactive=True,
            label="backends",
            choices=BACKENDS,
            value=BACKENDS,
            info="Backends to run the benchmark on.",
        )

        with gr.Row():
            with gr.Accordion(label="Process Config", open=False, visible=True):
                process_config = get_process_config()

        with gr.Row():
            with gr.Accordion(label="Scenario Config", open=False, visible=True):
                inference_config = get_inference_config()

        with gr.Row() as backend_configs:
            with gr.Accordion(label="OnnxRuntime Config", open=False, visible=True):
                onnxruntime_config = get_onnxruntime_config()
            with gr.Accordion(label="OpenVINO Config", open=False, visible=True):
                openvino_config = get_openvino_config()
            with gr.Accordion(label="PyTorch Config", open=False, visible=True):
                pytorch_config = get_pytorch_config()
            with gr.Accordion(label="IPEX Config", open=False, visible=True):
                ipex_config = get_ipex_config()

        backends.change(
            inputs=backends,
            outputs=backend_configs.children,
            fn=lambda values: [
                gr.update(visible=value in values) for value in BACKENDS
            ],
        )

        with gr.Row():
            button = gr.Button(value="Run Benchmark", variant="primary")

        with gr.Row():
            md_output = gr.Markdown(label="Output", value="")

        button.click(
            fn=run_benchmark,
            inputs={
                task,
                model,
                backends,
                *process_config.values(),
                *inference_config.values(),
                *onnxruntime_config.values(),
                *openvino_config.values(),
                *pytorch_config.values(),
                *ipex_config.values(),
            },
            outputs=[md_output],
            concurrency_limit=1,
        )

    return demo


if __name__ == "__main__":
    os.environ["LOG_TO_FILE"] = "0"
    os.environ["LOG_LEVEL"] = "INFO"
    setup_logging(level="INFO", prefix="MAIN-PROCESS")

    demo = build_demo()
    demo.queue(max_size=10).launch()