from optimum.exporters.tasks import TasksManager

from optimum.exporters.onnx import OnnxConfigWithPast, export, validate_model_outputs

from tempfile import TemporaryDirectory

from transformers import AutoConfig, is_torch_available

from transformers import AutoConfig

from pathlib import Path

import os
import shutil
import argparse

from typing import Optional

from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, get_repo_discussions
from huggingface_hub.file_download import repo_folder_name

def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:
    try:
        discussions = api.get_repo_discussions(repo_id=model_id)
    except Exception:
        return None
    for discussion in discussions:
        if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title:
            return discussion

def convert_onnx(model_id: str, task: str, folder: str):
        model_class = TasksManager.get_model_class_for_task(task)
        config = AutoConfig.from_pretrained(model_id)
        model = model_class.from_config(config)

        device = "cpu"  # ?

        # Dynamic axes aren't supported for YOLO-like models. This means they cannot be exported to ONNX on CUDA devices.
        # See: https://github.com/ultralytics/yolov5/pull/8378
        if model.__class__.__name__.startswith("Yolos") and device != "cpu":
            return

        onnx_config_class_constructor = TasksManager.get_exporter_config_constructor(model_type=config.model_type, exporter="onnx", task=task, model_name=model_id)
        onnx_config = onnx_config_class_constructor(model.config)

        # We need to set this to some value to be able to test the outputs values for batch size > 1.
        if (
            isinstance(onnx_config, OnnxConfigWithPast)
            and getattr(model.config, "pad_token_id", None) is None
            and task == "sequence-classification"
        ):
            model.config.pad_token_id = 0

            if is_torch_available():
                from optimum.exporters.onnx.utils import TORCH_VERSION

                if not onnx_config.is_torch_support_available:
                    print(
                        "Skipping due to incompatible PyTorch version. Minimum required is"
                        f" {onnx_config.MIN_TORCH_VERSION}, got: {TORCH_VERSION}"
                    )

                onnx_inputs, onnx_outputs = export(
                    model, onnx_config, onnx_config.DEFAULT_ONNX_OPSET, Path(folder), device=device
                )
                atol = onnx_config.ATOL_FOR_VALIDATION
                if isinstance(atol, dict):
                    atol = atol[task.replace("-with-past", "")]
                validate_model_outputs(
                    onnx_config,
                    model,
                    Path(folder),
                    onnx_outputs,
                    atol,
                )
            
            # TODO: iterate in folder and add all
            operations = [CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames]

            return operations


def convert(api: "HfApi", model_id: str, task:str, force: bool=False) -> Optional["CommitInfo"]:
    pr_title = "Adding ONNX file of this model"
    info = api.model_info(model_id)
    filenames = set(s.rfilename for s in info.siblings)

    with TemporaryDirectory() as d:
        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
        os.makedirs(folder)
        new_pr = None
        try:
            pr = previous_pr(api, model_id, pr_title)
            if "model.onnx" in filenames and not force:
                raise Exception(f"Model {model_id} is already converted, skipping..")
            elif pr is not None and not force:
                url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
                new_pr = pr
                raise Exception(f"Model {model_id} already has an open PR check out {url}")
            else:
                convert_onnx(model_id, task, folder)
        finally:
            shutil.rmtree(folder)
        return new_pr


if __name__ == "__main__":
    DESCRIPTION = """
    Simple utility tool to convert automatically a model on the hub to onnx format.
    It is PyTorch exclusive for now.
    It works by downloading the weights (PT), converting them locally, and uploading them back
    as a PR on the hub.
    """
    parser = argparse.ArgumentParser(description=DESCRIPTION)
    parser.add_argument(
        "model_id",
        type=str,
        help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",
    )
    parser.add_argument(
        "task",
        type=str,
        help="The task the model is performing",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Create the PR even if it already exists of if the model was already converted.",
    )
    args = parser.parse_args()
    api = HfApi()
    convert(api, args.model_id, task=args.task, force=args.force)