import re
import subprocess
import sys
from datetime import datetime
from importlib import import_module
from pathlib import Path

import gradio as gr
import pandas as pd

if Path("optimum-intel").is_dir():
    subprocess.run(["git", "pull"], cwd="optimum-intel")
else:
    subprocess.run(["git", "clone", "https://github.com/huggingface/optimum-intel.git"])
test_path = Path(__file__).parent / "optimum-intel" / "tests" / "openvino"
sys.path.append(str(test_path))

from test_modeling import *
from test_stable_diffusion import *

RESULT_FILE = "supported_models.md"

def generate_model_list():
    tests = []
    d = {}
    for item in globals().copy():
        match = re.match("(OVModelFor.*IntegrationTest)", item) or re.match("(OV.*PipelineTest)", item)
        if match:
            tests.append(match.group(1))

    for test in tests:
        task = test.replace("IntegrationTest", "").replace("Test", "")
        if "OVModelFor" in task:
            cls = getattr(import_module("test_modeling"), test)
        else:
            cls = getattr(import_module("test_stable_diffusion"), test)
        try:
            print(cls.SUPPORTED_ARCHITECTURES)
            d[task] = cls.SUPPORTED_ARCHITECTURES
        except AttributeError:
            print(cls.SUPPORTED_ARCHITECTURES_WITH_HIDDEN_STATES)
            print(cls.SUPPORTED_ARCHITECTURES_WITH_ATTENTION)
            d[task] = cls.SUPPORTED_ARCHITECTURES_WITH_ATTENTION + cls.SUPPORTED_ARCHITECTURES_WITH_HIDDEN_STATES

    with open(RESULT_FILE, "w") as f:
        f.write(f"Updated at {datetime.now().strftime('%d %B %Y')}\n\n")

        summary = []
        all_archs = []
        for archs in d.values():
            all_archs += archs
        for title, supported_models in d.items():
            f.write(f"## {title}\n\n")
            for item in supported_models:
                f.write(f" - {item}\n")
            f.write("\n")
            summary.append((title, len(supported_models)))
        md_summary = pd.DataFrame.from_records(summary, columns=["task", "number of architectures"]).to_markdown()
        f.write("# Summary\n\n")
        f.write(md_summary)
        f.write("\n\n")
        num_total_archs = len(set(all_archs))
        f.write(f"Total unique architectures: {num_total_archs}\n\n")
        f.write(f"Total validated architecture/task combinations: {len(all_archs)}\n\n")

    return Path(RESULT_FILE).read_text(), RESULT_FILE


demo = gr.Interface(
    fn=generate_model_list,
    title="List of validated architectures for optimum[openvino]",
    inputs=[],
    outputs=[gr.Markdown(), gr.File()],
    allow_flagging="never",
)
demo.launch(server_name="0.0.0.0")