import argparse
import json
import os
import shutil
from collections import defaultdict
from tempfile import TemporaryDirectory
from typing import Dict, List, Optional, Set, Tuple

import torch
from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download
from huggingface_hub.file_download import repo_folder_name
from safetensors.torch import _find_shared_tensors, _is_complete, load_file, save_file


COMMIT_DESCRIPTION = """
This is an automated PR created with https://huggingface.co/spaces/safetensors/convert

This new file is equivalent to `pytorch_model.bin` but safe in the sense that
no arbitrary code can be put into it.

These files also happen to load much faster than their pytorch counterpart:
https://colab.research.google.com/github/huggingface/notebooks/blob/main/safetensors_doc/en/speed.ipynb

The widgets on your model page will run using this model even if this is not merged
making sure the file actually works.

If you find any issues: please report here: https://huggingface.co/spaces/safetensors/convert/discussions

Feel free to ignore this PR.
"""

ConversionResult = Tuple[List["CommitOperationAdd"], List[Tuple[str, "Exception"]]]


def _remove_duplicate_names(
    state_dict: Dict[str, torch.Tensor],
    *,
    preferred_names: Optional[List[str]] = None,
    discard_names: Optional[List[str]] = None,
) -> Dict[str, List[str]]:
    if preferred_names is None:
        preferred_names = []
    preferred_names = set(preferred_names)
    if discard_names is None:
        discard_names = []
    discard_names = set(discard_names)

    shareds = _find_shared_tensors(state_dict)
    to_remove = defaultdict(list)

    for shared in shareds:
        complete_names = set([name for name in shared if _is_complete(state_dict[name])])
        if not complete_names:
            if len(shared) == 1:
                # Force contiguous
                name = list(shared)[0]
                state_dict[name] = state_dict[name].clone()
                complete_names = {name}
            else:
                raise RuntimeError(
                    f"Error while trying to find names to remove to save state dict, but found no suitable name to keep "
                    f"for saving amongst: {shared}. None covers the entire storage. Refusing to save/load the model "
                    f"since you could be storing much more memory than needed. Please refer to "
                    f"https://huggingface.co/docs/safetensors/torch_shared_tensors"
                )

        keep_name = sorted(list(complete_names))[0]

        # Prefer keys not in discard list
        preferred = complete_names.difference(discard_names)
        if preferred:
            keep_name = sorted(list(preferred))[0]

        # Prefer explicitly preferred names (if any)
        if preferred_names:
            preferred2 = preferred_names.intersection(complete_names)
            if preferred2:
                keep_name = sorted(list(preferred2))[0]

        for name in sorted(shared):
            if name != keep_name:
                to_remove[keep_name].append(name)

    return to_remove


def get_discard_names(model_id: str, revision: Optional[str], folder: str, token: Optional[str]) -> List[str]:
    try:
        import transformers

        config_filename = hf_hub_download(
            repo_id=model_id, revision=revision, filename="config.json", token=token, cache_dir=folder
        )
        with open(config_filename, "r") as f:
            config = json.load(f)

        architecture = config.get("architectures", [None])[0]
        if not architecture:
            return []

        class_ = getattr(transformers, architecture, None)
        if class_ is None:
            return []

        # Name depends on transformers version
        discard_names = getattr(class_, "_tied_weights_keys", [])
        if discard_names is None:
            discard_names = []
        return list(discard_names)
    except Exception:
        return []


class AlreadyExists(Exception):
    pass


def check_file_size(sf_filename: str, pt_filename: str):
    sf_size = os.stat(sf_filename).st_size
    pt_size = os.stat(pt_filename).st_size

    if pt_size > 0 and (sf_size - pt_size) / pt_size > 0.01:
        raise RuntimeError(
            f"The file size difference is more than 1%:\n"
            f" - {sf_filename}: {sf_size}\n"
            f" - {pt_filename}: {pt_size}\n"
        )


def rename(pt_filename: str) -> str:
    filename, _ext = os.path.splitext(pt_filename)
    local = f"{filename}.safetensors"
    local = local.replace("pytorch_model", "model")
    return local


def convert_multi(
    model_id: str, *, revision: Optional[str], folder: str, token: Optional[str], discard_names: List[str]
) -> ConversionResult:
    filename = hf_hub_download(
        repo_id=model_id, revision=revision, filename="pytorch_model.bin.index.json", token=token, cache_dir=folder
    )
    with open(filename, "r") as f:
        data = json.load(f)

    filenames = set(data["weight_map"].values())
    local_filenames = []
    for fname in filenames:
        pt_filename = hf_hub_download(repo_id=model_id, revision=revision, filename=fname, token=token, cache_dir=folder)
        sf_filename = rename(pt_filename)
        sf_filename = os.path.join(folder, sf_filename)
        convert_file(pt_filename, sf_filename, discard_names=discard_names)
        local_filenames.append(sf_filename)

    index = os.path.join(folder, "model.safetensors.index.json")
    with open(index, "w") as f:
        newdata = {k: v for k, v in data.items()}
        newmap = {k: rename(v) for k, v in data["weight_map"].items()}
        newdata["weight_map"] = newmap
        json.dump(newdata, f, indent=4)
    local_filenames.append(index)

    operations = [
        CommitOperationAdd(path_in_repo=os.path.basename(local), path_or_fileobj=local)
        for local in local_filenames
    ]
    errors: List[Tuple[str, "Exception"]] = []
    return operations, errors


def convert_single(
    model_id: str, *, revision: Optional[str], folder: str, token: Optional[str], discard_names: List[str]
) -> ConversionResult:
    pt_filename = hf_hub_download(
        repo_id=model_id, revision=revision, filename="pytorch_model.bin", token=token, cache_dir=folder
    )

    sf_name = "model.safetensors"
    sf_filename = os.path.join(folder, sf_name)
    convert_file(pt_filename, sf_filename, discard_names)
    operations = [CommitOperationAdd(path_in_repo=sf_name, path_or_fileobj=sf_filename)]
    errors: List[Tuple[str, "Exception"]] = []
    return operations, errors


def convert_file(pt_filename: str, sf_filename: str, discard_names: List[str]):
    # Safer behavior across PyTorch versions: explicitly set weights_only=True.
    # PyTorch 2.6 changed defaults for security reasons. :contentReference[oaicite:4]{index=4}
    loaded = torch.load(pt_filename, map_location="cpu", weights_only=True)

    if isinstance(loaded, dict) and "state_dict" in loaded:
        loaded = loaded["state_dict"]

    if not isinstance(loaded, dict):
        raise RuntimeError(f"Expected a state_dict-like dict in {pt_filename}, got {type(loaded)}")

    to_removes = _remove_duplicate_names(loaded, discard_names=discard_names)

    metadata = {"format": "pt"}
    for kept_name, to_remove_group in to_removes.items():
        for to_remove in to_remove_group:
            if to_remove not in metadata:
                metadata[to_remove] = kept_name
            del loaded[to_remove]

    loaded = {k: v.contiguous() for k, v in loaded.items()}

    os.makedirs(os.path.dirname(sf_filename), exist_ok=True)
    save_file(loaded, sf_filename, metadata=metadata)

    check_file_size(sf_filename, pt_filename)

    reloaded = load_file(sf_filename)
    for k in loaded:
        if not torch.equal(loaded[k], reloaded[k]):
            raise RuntimeError(f"The output tensors do not match for key {k}")


def previous_pr(api: HfApi, model_id: str, pr_title: str, revision: Optional[str]) -> Optional[Discussion]:
    """
    Check if a PR with the same title already exists.
    Uses get_repo_discussions(discussion_type="pull_request"). :contentReference[oaicite:5]{index=5}
    """
    try:
        base_sha = api.model_info(model_id, revision=revision).sha
        discussions = api.get_repo_discussions(
            repo_id=model_id,
            discussion_type="pull_request",
            discussion_status="all",
        )
    except Exception:
        return None

    for d in discussions:
        if not (d.is_pull_request and d.title == pr_title and d.status in {"open", "closed"}):
            continue

        # Best-effort check to avoid duplicates:
        # - If git_reference exists, compare the base commit if possible.
        # - If that fails, return the matching PR anyway (safer than duplicating).
        try:
            if getattr(d, "git_reference", None):
                commits = api.list_repo_commits(model_id, revision=d.git_reference)
                # list_repo_commits returns newest-first; base is typically last.
                if commits and commits[-1].commit_id == base_sha:
                    return d
        except Exception:
            return d

        # If we can’t confirm base SHA, still consider it a duplicate by title
        return d

    return None


def convert_generic(
    model_id: str,
    *,
    revision: Optional[str],
    folder: str,
    filenames: Set[str],
    token: Optional[str],
    discard_names: List[str],
) -> ConversionResult:
    operations: List[CommitOperationAdd] = []
    errors: List[Tuple[str, "Exception"]] = []

    extensions = {".bin", ".ckpt"}
    for filename in filenames:
        prefix, ext = os.path.splitext(filename)
        if ext not in extensions:
            continue

        pt_filename = hf_hub_download(
            repo_id=model_id, revision=revision, filename=filename, token=token, cache_dir=folder
        )
        dirname, raw_filename = os.path.split(filename)

        if raw_filename == "pytorch_model.bin":
            sf_in_repo = os.path.join(dirname, "model.safetensors")
        else:
            sf_in_repo = f"{prefix}.safetensors"

        sf_filename = os.path.join(folder, sf_in_repo)
        try:
            convert_file(pt_filename, sf_filename, discard_names=discard_names)
            operations.append(CommitOperationAdd(path_in_repo=sf_in_repo, path_or_fileobj=sf_filename))
        except Exception as e:
            errors.append((pt_filename, e))

    return operations, errors


def convert(api: HfApi, model_id: str, revision: Optional[str] = None, force: bool = False) -> Tuple[CommitInfo, List[Tuple[str, "Exception"]]]:
    pr_title = "Adding `safetensors` variant of this model"
    info = api.model_info(model_id, revision=revision)
    filenames = set(s.rfilename for s in info.siblings)
    library_name = getattr(info, "library_name", None)

    with TemporaryDirectory() as d:
        folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
        os.makedirs(folder, exist_ok=True)

        new_pr: Optional[CommitInfo] = None
        errors: List[Tuple[str, "Exception"]] = []

        discard_names = get_discard_names(model_id, revision=revision, folder=folder, token=api.token)

        operations: Optional[List[CommitOperationAdd]] = None

        pr = previous_pr(api, model_id, pr_title, revision=revision)
        if any(fn.endswith(".safetensors") for fn in filenames) and not force:
            raise AlreadyExists(f"Model {model_id} is already converted, skipping.")
        if pr is not None and getattr(pr, "author", None) == "SFconvertbot" and not force:
            url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
            raise AlreadyExists(f"Model {model_id} already has an open PR: {url}")

        if library_name == "transformers":
            if "pytorch_model.bin" in filenames:
                operations, errors = convert_single(model_id, revision=revision, folder=folder, token=api.token, discard_names=discard_names)
            elif "pytorch_model.bin.index.json" in filenames:
                operations, errors = convert_multi(model_id, revision=revision, folder=folder, token=api.token, discard_names=discard_names)
            else:
                raise RuntimeError(f"Model {model_id} doesn't look like a standard PyTorch Transformers model. Cannot convert.")
        else:
            operations, errors = convert_generic(model_id, revision=revision, folder=folder, filenames=filenames, token=api.token, discard_names=discard_names)

        if not operations:
            raise RuntimeError("No files to convert.")

        # Re-check for duplicate PR right before commit (race safety)
        pr = previous_pr(api, model_id, pr_title, revision=revision)
        if pr is not None and not force:
            url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
            raise AlreadyExists(f"Model {model_id} already has an open PR: {url}")

        # IMPORTANT: create_pr=True cannot be used with revision != None/"main". :contentReference[oaicite:6]{index=6}
        # Workaround: always open PR against main, but anchor it to the requested revision commit if provided.
        commit_revision = revision
        parent_commit = None
        if revision is not None and revision != "main":
            parent_commit = api.model_info(model_id, revision=revision).sha
            commit_revision = "main"

        new_pr = api.create_commit(
            repo_id=model_id,
            revision=commit_revision,
            parent_commit=parent_commit,
            operations=operations,
            commit_message=pr_title,
            commit_description=COMMIT_DESCRIPTION,
            create_pr=True,
        )
        # create_commit returns CommitInfo with pr_url when create_pr=True. :contentReference[oaicite:7]{index=7}
        print(f"PR created at {new_pr.pr_url}")

        return new_pr, errors


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("--revision", type=str, help="The revision to convert")
    parser.add_argument("--force", action="store_true", help="Create PR even if it exists or model already converted.")
    parser.add_argument("-y", action="store_true", help="Ignore safety prompt")
    args = parser.parse_args()

    model_id = args.model_id
    api = HfApi()

    if args.y:
        txt = "y"
    else:
        txt = input(
            "This conversion script will unpickle a pickled file, which is inherently unsafe. If you do not trust this "
            "file, use https://huggingface.co/spaces/safetensors/convert or another hosted solution. Continue [Y/n] ? "
        )

    if txt.lower() in {"", "y"}:
        commit_info, errors = convert(api, model_id, revision=args.revision, force=args.force)
        string = f"""
### Success 🔥
Yay! This model was successfully converted and a PR was opened here:
[{commit_info.pr_url}]({commit_info.pr_url})
"""
        if errors:
            string += "\nErrors during conversion:\n"
            string += "\n".join(f"Error while converting {filename}: {e}, skipped" for filename, e in errors)
        print(string)
    else:
        print(f"Answer was `{txt}` aborting.")