import json
import logging
import os
from functools import partial
from typing import Union, Optional

import torch
from torch.hub import (
    load_state_dict_from_url,
    download_url_to_file,
    urlparse,
    HASH_REGEX,
)

try:
    from torch.hub import get_dir
except ImportError:
    from torch.hub import _get_torch_home as get_dir

from timm import __version__

try:
    from huggingface_hub import hf_hub_url
    from huggingface_hub import cached_download

    cached_download = partial(
        cached_download, library_name="timm", library_version=__version__
    )
except ImportError:
    hf_hub_url = None
    cached_download = None

_logger = logging.getLogger(__name__)


def get_cache_dir(child_dir=""):
    """
    Returns the location of the directory where models are cached (and creates it if necessary).
    """
    # Issue warning to move data if old env is set
    if os.getenv("TORCH_MODEL_ZOO"):
        _logger.warning(
            "TORCH_MODEL_ZOO is deprecated, please use env TORCH_HOME instead"
        )

    hub_dir = get_dir()
    child_dir = () if not child_dir else (child_dir,)
    model_dir = os.path.join(hub_dir, "checkpoints", *child_dir)
    os.makedirs(model_dir, exist_ok=True)
    return model_dir


def download_cached_file(url, check_hash=True, progress=False):
    parts = urlparse(url)
    filename = os.path.basename(parts.path)
    cached_file = os.path.join(get_cache_dir(), filename)
    if not os.path.exists(cached_file):
        _logger.info('Downloading: "{}" to {}\n'.format(url, cached_file))
        hash_prefix = None
        if check_hash:
            r = HASH_REGEX.search(filename)  # r is Optional[Match[str]]
            hash_prefix = r.group(1) if r else None
        download_url_to_file(url, cached_file, hash_prefix, progress=progress)
    return cached_file


def has_hf_hub(necessary=False):
    if hf_hub_url is None and necessary:
        # if no HF Hub module installed and it is necessary to continue, raise error
        raise RuntimeError(
            "Hugging Face hub model specified but package not installed. Run `pip install huggingface_hub`."
        )
    return hf_hub_url is not None


def hf_split(hf_id):
    rev_split = hf_id.split("@")
    assert (
        0 < len(rev_split) <= 2
    ), "hf_hub id should only contain one @ character to identify revision."
    hf_model_id = rev_split[0]
    hf_revision = rev_split[-1] if len(rev_split) > 1 else None
    return hf_model_id, hf_revision


def load_cfg_from_json(json_file: Union[str, os.PathLike]):
    with open(json_file, "r", encoding="utf-8") as reader:
        text = reader.read()
    return json.loads(text)


def _download_from_hf(model_id: str, filename: str):
    hf_model_id, hf_revision = hf_split(model_id)
    url = hf_hub_url(hf_model_id, filename, revision=hf_revision)
    return cached_download(url, cache_dir=get_cache_dir("hf"))


def load_model_config_from_hf(model_id: str):
    assert has_hf_hub(True)
    cached_file = _download_from_hf(model_id, "config.json")
    default_cfg = load_cfg_from_json(cached_file)
    default_cfg["hf_hub"] = (
        model_id  # insert hf_hub id for pretrained weight load during model creation
    )
    model_name = default_cfg.get("architecture")
    return default_cfg, model_name


def load_state_dict_from_hf(model_id: str):
    assert has_hf_hub(True)
    cached_file = _download_from_hf(model_id, "pytorch_model.bin")
    state_dict = torch.load(cached_file, map_location="cpu")
    return state_dict
