import os
import shutil
import torch
import gradio as gr
from huggingface_hub import HfApi, whoami, ModelCard
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from textwrap import dedent
from pathlib import Path

from tempfile import TemporaryDirectory

from huggingface_hub.file_download import repo_folder_name
from optimum.intel.utils.constant import _TASK_ALIASES
from optimum.exporters import TasksManager

from optimum.intel.utils.modeling_utils import _find_files_matching_pattern
from optimum.intel import (
    OVModelForAudioClassification,
    OVModelForCausalLM,
    OVModelForFeatureExtraction,
    OVModelForImageClassification,
    OVModelForMaskedLM,
    OVModelForQuestionAnswering,
    OVModelForSeq2SeqLM,
    OVModelForSequenceClassification,
    OVModelForTokenClassification,
    OVStableDiffusionPipeline,
    OVStableDiffusionXLPipeline,
    OVLatentConsistencyModelPipeline,
    OVModelForPix2Struct,
    OVWeightQuantizationConfig,
)
from diffusers import ConfigMixin

_HEAD_TO_AUTOMODELS = {
    "feature-extraction": "OVModelForFeatureExtraction",
    "fill-mask": "OVModelForMaskedLM",
    "text-generation": "OVModelForCausalLM",
    "text-classification": "OVModelForSequenceClassification",
    "token-classification": "OVModelForTokenClassification",
    "question-answering": "OVModelForQuestionAnswering",
    "image-classification": "OVModelForImageClassification",
    "audio-classification": "OVModelForAudioClassification",
    "stable-diffusion": "OVStableDiffusionPipeline",
    "stable-diffusion-xl": "OVStableDiffusionXLPipeline",
    "latent-consistency": "OVLatentConsistencyModelPipeline",
}


def export(model_id: str, private_repo: bool, overwritte: bool, oauth_token: gr.OAuthToken):
    if oauth_token.token is None:
        return "You must be logged in to use this space"

    if not model_id:
        return f"### Invalid input 🐞 Please specify a model name, got {model_id}"

    try:
        model_name = model_id.split("/")[-1]
        username = whoami(oauth_token.token)["name"]
        new_repo_id = f"{username}/{model_name}-openvino"
        library_name = TasksManager.infer_library_from_model(model_id, token=oauth_token.token)

        if library_name == "diffusers":
            ConfigMixin.config_name = "model_index.json"
            class_name = ConfigMixin.load_config(model_id, token=oauth_token.token)["_class_name"].lower()
            if "xl" in class_name:
                task = "stable-diffusion-xl"
            elif "consistency" in class_name:
                task = "latent-consistency"
            else:
                task = "stable-diffusion"
        else:
            task = TasksManager.infer_task_from_model(model_id, token=oauth_token.token)

        if task == "text2text-generation":
            return "Export of Seq2Seq models is currently disabled"

        if task not in _HEAD_TO_AUTOMODELS:
            return f"The task '{task}' is not supported, only {_HEAD_TO_AUTOMODELS.keys()} tasks are supported"

        auto_model_class = _HEAD_TO_AUTOMODELS[task]
        ov_files = _find_files_matching_pattern(
            model_id,
            pattern=r"(.*)?openvino(.*)?\_model.xml",
            use_auth_token=oauth_token.token,
        )

        if len(ov_files) > 0:
            return f"Model {model_id} is already converted, skipping.."

        api = HfApi(token=oauth_token.token)
        if api.repo_exists(new_repo_id) and not overwritte:
            return f"Model {new_repo_id} already exist, please set overwritte=True to push on an existing repo"

        with TemporaryDirectory() as d:
            folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
            os.makedirs(folder)
            try:
                api.snapshot_download(repo_id=model_id, local_dir=folder, allow_patterns=["*.json"])
                ov_model = eval(auto_model_class).from_pretrained(model_id, export=True, cache_dir=folder, token=oauth_token.token)
                ov_model.save_pretrained(folder)
                new_repo_url = api.create_repo(repo_id=new_repo_id, exist_ok=True, private=private_repo)
                new_repo_id = new_repo_url.repo_id
                print("Repo created successfully!", new_repo_url)

                folder = Path(folder)
                for dir_name in (
                    "",
                    "vae_encoder",
                    "vae_decoder",
                    "text_encoder",
                    "text_encoder_2",
                    "unet",
                    "tokenizer",
                    "tokenizer_2",
                    "scheduler",
                    "feature_extractor",
                ):
                    if not (folder / dir_name).is_dir():
                        continue
                    for file_path in (folder / dir_name).iterdir():
                        if file_path.is_file():
                            try:
                                api.upload_file(
                                    path_or_fileobj=file_path,
                                    path_in_repo=os.path.join(dir_name, file_path.name),
                                    repo_id=new_repo_id,
                                )
                            except Exception as e:
                                return f"Error uploading file {file_path}: {e}"

                try:
                    card = ModelCard.load(model_id, token=oauth_token.token)
                except:
                    card = ModelCard("")

                if card.data.tags is None:
                    card.data.tags = []
                card.data.tags.append("openvino")
                card.data.base_model = model_id
                card.text = dedent(
                    f"""
                    This model was converted to OpenVINO from [`{model_id}`](https://huggingface.co/{model_id}) using [optimum-intel](https://github.com/huggingface/optimum-intel)
                    via the [export](https://huggingface.co/spaces/echarlaix/openvino-export) space.

                    First make sure you have optimum-intel installed:

                    ```bash
                    pip install optimum[openvino]
                    ```

                    To load your model you can do as follows:

                    ```python
                    from optimum.intel import {auto_model_class}

                    model_id = "{new_repo_id}"
                    model = {auto_model_class}.from_pretrained(model_id)
                    ```
                    """
                )
                card_path = os.path.join(folder, "README.md")
                card.save(card_path)

                api.upload_file(
                    path_or_fileobj=card_path,
                    path_in_repo="README.md",
                    repo_id=new_repo_id,
                )
                return f"This model was successfully exported, find it under your repo {new_repo_url}'"
            finally:
                shutil.rmtree(folder, ignore_errors=True)
    except Exception as e:
        return f"### Error: {e}"

DESCRIPTION = """
This Space uses [Optimum Intel](https://huggingface.co/docs/optimum/main/en/intel/openvino/export) to automatically export a model from the Hub to the [OpenVINO format](https://docs.openvino.ai/2024/documentation/openvino-ir-format.html).

The resulting model will then be pushed under your HF user namespace.

The list of the supported architectures can be found in the [documentation](https://huggingface.co/docs/optimum/main/en/intel/openvino/models)
"""

model_id = HuggingfaceHubSearch(
    label="Hub Model ID",
    placeholder="Search for model id on the hub",
    search_type="model",
)
private_repo = gr.Checkbox(
    value=False,
    label="Private Repo",
    info="Create a private repo under your username",
)
overwritte = gr.Checkbox(
    value=False,
    label="Overwrite repo content",
    info="Push files on existing repo potentially overwriting existing files",
)
interface = gr.Interface(
    fn=export,
    inputs=[
        model_id,
        private_repo,
        overwritte,
    ],
    outputs=[
        gr.Markdown(label="output"),
    ],
    title="Export your model to OpenVINO",
    description=DESCRIPTION,
    api_name=False,
)

with gr.Blocks() as demo:
    gr.Markdown("You must be logged in to use this space")
    gr.LoginButton(min_width=250)
    interface.render()

demo.launch()