import argparse
import json
import os
import torch
import shutil
from tempfile import TemporaryDirectory
from typing import List, Optional
from diffusers import DiffusionPipeline

from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download
from huggingface_hub.file_download import repo_folder_name


class AlreadyExists(Exception):
    pass


def is_index_stable_diffusion_like(config_dict):
    if "_class_name" not in config_dict:
        return False

    compatible_classes = [
        "AltDiffusionImg2ImgPipeline",
        "AltDiffusionPipeline",
        "CycleDiffusionPipeline",
        "StableDiffusionImageVariationPipeline",
        "StableDiffusionImg2ImgPipeline",
        "StableDiffusionInpaintPipeline",
        "StableDiffusionInpaintPipelineLegacy",
        "StableDiffusionPipeline",
        "StableDiffusionPipelineSafe",
        "StableDiffusionUpscalePipeline",
        "VersatileDiffusionDualGuidedPipeline",
        "VersatileDiffusionImageVariationPipeline",
        "VersatileDiffusionPipeline",
        "VersatileDiffusionTextToImagePipeline",
        "OnnxStableDiffusionImg2ImgPipeline",
        "OnnxStableDiffusionInpaintPipeline",
        "OnnxStableDiffusionInpaintPipelineLegacy",
        "OnnxStableDiffusionPipeline",
        "StableDiffusionOnnxPipeline",
        "FlaxStableDiffusionPipeline",
    ]
    return config_dict["_class_name"] in compatible_classes


def convert_single(model_id: str, folder: str) -> List["CommitOperationAdd"]:
    pipe = DiffusionPipeline.from_pretrained(model_id, cache_dir="/home/patrick/cache_to_delete")

    try:
        pipe.to(torch_dtype=torch.float16)
        pipe.save_pretrained(folder, variant="fp16")
        pipe.save_pretrained(folder, variant="fp16", safe_serialization=True)

        all_files = []
        def find_files_in_dir(directory):
            for root, dirs, files in os.walk(directory):
                for file in files:
                    all_files.append(os.path.join(root, file))

        find_files_in_dir(folder)
        files = [f for f in all_files if ".fp16." in f]

        operations = [CommitOperationAdd(path_in_repo='/'.join(f.split("/")[-2:]), path_or_fileobj=f) for f in files]
        return operations
    except Exception as e:
        print(e)
        return False

def convert_file(
    old_config: str,
    new_config: str,
):
    with open(old_config, "r") as f:
        old_dict = json.load(f)

    old_dict["feature_extractor"][-1] = "CLIPImageProcessor"
    # if "clip_sample" not in old_dict:
    #     print("Make scheduler DDIM compatible")
    #     old_dict["clip_sample"] = False
    # else:
    #     print("No matching config")
    #     return False

    with open(new_config, 'w') as f:
        json_str = json.dumps(old_dict, indent=2, sort_keys=True) + "\n"
        f.write(json_str)

    return "Stable Diffusion"


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(api: "HfApi", model_id: str, force: bool = False) -> Optional["CommitInfo"]:
    pr_title = "Fix deprecated float16/fp16 variant loading through new `version` API."

    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:
            operations = None
            pr = previous_pr(api, model_id, pr_title)
            if pr is not None and not force:
                url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
                new_pr = pr
                raise AlreadyExists(f"Model {model_id} already has an open PR check out {url}")
            else:
                operations = convert_single(model_id, folder)

            if operations:
                contributor = model_id.split("/")[0]
                pr_description = (
                        f"Hey {contributor} 👋, \n\n Your model repository seems to contain a [`fp16` branch](https://huggingface.co/{model_id}/tree/fp16) to load the model in float16 precision. "
                        "Loading `fp16` versions from a branch instead of the main branch is deprecated and will eventually be forbidden. "
                        "Instead, we strongly recommend to save `fp16` versions of the model under `.fp16.` version files directly on the 'main' branch as enabled through this PR."
                        f"This PR makes sure that your model repository allows the user to correctly download float16 precision model weights by adding `fp16` model weights in both safetensors and PyTorch bin format:"
                        "\n\n"
                        "```py\n"
                        f"pipe = DiffusionPipeline.from_pretrained({model_id}, torch_dtype=torch.float16, variant='fp16')"
                        "\n```"
                        "\n\n"
                        "For more information please have a look at: https://huggingface.co/docs/diffusers/using-diffusers/loading#checkpoint-variants."
                        "\nWe made sure you that you can safely merge this pull request. \n\n Best, the 🧨 Diffusers team."
                )
                new_pr = api.create_commit(
                    repo_id=model_id,
                    operations=operations,
                    commit_message=pr_title,
                    commit_description=pr_description,
                    create_pr=True,
                )
                print(f"Pr created at {new_pr.pr_url}")
            else:
                print(f"No files to convert for {model_id}")
        finally:
            shutil.rmtree(folder)
        return new_pr


if __name__ == "__main__":
    DESCRIPTION = """
    Simple utility tool to convert automatically some weights on the hub to `safetensors` 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(
        "--force",
        action="store_true",
        help="Create the PR even if it already exists of if the model was already converted.",
    )
    args = parser.parse_args()
    model_id = args.model_id
    api = HfApi()
    convert(api, model_id, force=args.force)