"""Classes and functions for managing AI models"""

import logging
import os
import re
import sys
import urllib.parse
from functools import lru_cache, wraps

import requests
import torch
from huggingface_hub import (
    HfFolder,
    hf_hub_download as _hf_hub_download,
    try_to_load_from_cache,
)
from omegaconf import OmegaConf
from refiners.foundationals.latent_diffusion import SD1UNet
from refiners.foundationals.latent_diffusion.model import LatentDiffusionModel
from safetensors.torch import load_file

from imaginairy import config as iconfig
from imaginairy.config import IMAGE_WEIGHTS_SHORT_NAMES, ModelArchitecture
from imaginairy.modules import attention
from imaginairy.utils import get_device, instantiate_from_config
from imaginairy.utils.model_cache import memory_managed_model
from imaginairy.utils.named_resolutions import normalize_image_size
from imaginairy.utils.paths import PKG_ROOT

logger = logging.getLogger(__name__)

MOST_RECENTLY_LOADED_MODEL = None


class HuggingFaceAuthorizationError(RuntimeError):
    pass


def load_tensors(tensorfile, map_location=None):
    if tensorfile == "empty":
        # used for testing
        return {}
    if tensorfile.endswith((".ckpt", ".pth", ".bin")):
        return torch.load(tensorfile, map_location=map_location)
    if tensorfile.endswith(".safetensors"):
        return load_file(tensorfile, device=map_location)

    return load_file(tensorfile, device=map_location)

    # raise ValueError(f"Unknown tensorfile type: {tensorfile}")


def load_state_dict(weights_location, half_mode=False, device=None):
    if device is None:
        device = get_device()

    if weights_location.startswith("http"):
        ckpt_path = get_cached_url_path(weights_location, category="weights")
    else:
        ckpt_path = weights_location
    logger.info(f"Loading model {ckpt_path} onto {get_device()} backend...")
    state_dict = None
    # weights_cache_key = (ckpt_path, half_mode)
    # if weights_cache_key in GLOBAL_WEIGHTS_CACHE:
    #     return GLOBAL_WEIGHTS_CACHE.get(weights_cache_key)

    try:
        state_dict = load_tensors(ckpt_path, map_location="cpu")
    except FileNotFoundError as e:
        if e.errno == 2:
            logger.error(
                f'Error: "{ckpt_path}" not a valid path to model weights.\nPreconfigured models you can use: {IMAGE_WEIGHTS_SHORT_NAMES}.'
            )
            sys.exit(1)
        raise
    except RuntimeError as e:
        err_str = str(e)
        if (
            "PytorchStreamReader failed reading zip archive" in err_str
            and weights_location.startswith("http")
        ):
            logger.warning("Corrupt checkpoint. deleting and re-downloading...")
            os.remove(ckpt_path)
            ckpt_path = get_cached_url_path(weights_location, category="weights")
            state_dict = load_tensors(ckpt_path, map_location="cpu")
        if state_dict is None:
            raise

    state_dict = state_dict.get("state_dict", state_dict)

    if half_mode:
        state_dict = {k: v.half() for k, v in state_dict.items()}

    # change device
    state_dict = {k: v.to(device) for k, v in state_dict.items()}

    # GLOBAL_WEIGHTS_CACHE.set(weights_cache_key, state_dict)

    return state_dict


def load_model_from_config(config, weights_location, half_mode=False):
    model = instantiate_from_config(config.model)
    base_model_dict = load_state_dict(weights_location, half_mode=half_mode)
    model.init_from_state_dict(base_model_dict)
    if half_mode:
        model = model.half()
    model.to(get_device())
    model.eval()
    return model


def load_model_from_config_old(
    config, weights_location, control_weights_locations=None, half_mode=False
):
    model = instantiate_from_config(config.model)
    base_model_dict = load_state_dict(weights_location, half_mode=half_mode)
    model.init_from_state_dict(base_model_dict)

    control_weights_locations = control_weights_locations or []
    controlnets = []
    for control_weights_location in control_weights_locations:
        controlnet_state_dict = load_state_dict(
            control_weights_location, half_mode=half_mode
        )
        controlnet_state_dict = {
            k.replace("control_model.", ""): v for k, v in controlnet_state_dict.items()
        }
        controlnet = instantiate_from_config(model.control_stage_config)
        controlnet.load_state_dict(controlnet_state_dict)
        controlnet.to(get_device())
        controlnets.append(controlnet)
        model.set_control_models(controlnets)

    if half_mode:
        model = model.half()
    print("halved")

    model.to(get_device())
    print("moved to device")
    model.eval()
    print("set to eval mode")
    return model


def add_controlnet(base_state_dict, controlnet_state_dict):
    """Merges a base sd15 model with a controlnet model."""
    for key in controlnet_state_dict:
        base_state_dict[key] = controlnet_state_dict[key]
    return base_state_dict


def get_diffusion_model(
    weights_location=iconfig.DEFAULT_MODEL_WEIGHTS,
    config_path="configs/stable-diffusion-v1.yaml",
    control_weights_locations=None,
    half_mode=None,
    for_inpainting=False,
):
    """
    Load a diffusion model.

    Weights location may also be shortcut name, e.g. "SD-1.5"
    """
    try:
        return _get_diffusion_model(
            weights_location,
            config_path,
            half_mode,
            for_inpainting,
            control_weights_locations=control_weights_locations,
        )
    except HuggingFaceAuthorizationError as e:
        if for_inpainting:
            logger.warning(
                f"Failed to load inpainting model. Attempting to fall-back to standard model.   {e!s}"
            )
            return _get_diffusion_model(
                iconfig.DEFAULT_MODEL_WEIGHTS,
                config_path,
                half_mode,
                for_inpainting=False,
                control_weights_locations=control_weights_locations,
            )
        raise


def _get_diffusion_model(
    weights_location=iconfig.DEFAULT_MODEL_WEIGHTS,
    model_architecture="configs/stable-diffusion-v1.yaml",
    half_mode=None,
    for_inpainting=False,
    control_weights_locations=None,
):
    """
    Load a diffusion model.

    Weights location may also be shortcut name, e.g. "SD-1.5"
    """
    global MOST_RECENTLY_LOADED_MODEL

    model_weights_config = resolve_model_weights_config(
        model_weights=weights_location,
        default_model_architecture=model_architecture,
        for_inpainting=for_inpainting,
    )
    # some models need the attention calculated in float32
    if model_weights_config is not None:
        attention.ATTENTION_PRECISION_OVERRIDE = (
            model_weights_config.forced_attn_precision
        )
    else:
        attention.ATTENTION_PRECISION_OVERRIDE = "default"
    diffusion_model = _load_diffusion_model(
        config_path=model_weights_config.architecture.config_path,
        weights_location=weights_location,
        half_mode=half_mode,
    )
    MOST_RECENTLY_LOADED_MODEL = diffusion_model
    if control_weights_locations:
        controlnets = []
        for control_weights_location in control_weights_locations:
            controlnets.append(load_controlnet(control_weights_location, half_mode))
        diffusion_model.set_control_models(controlnets)

    return diffusion_model


def get_diffusion_model_refiners(
    weights_config: iconfig.ModelWeightsConfig,
    for_inpainting=False,
    dtype=None,
) -> LatentDiffusionModel:
    """Load a diffusion model."""
    return _get_diffusion_model_refiners(
        weights_location=weights_config.weights_location,
        for_inpainting=for_inpainting,
        dtype=dtype,
    )


@lru_cache(maxsize=1)
def _get_diffusion_model_refiners(
    weights_location: str,
    for_inpainting: bool = False,
    device=None,
    dtype=torch.float16,
) -> LatentDiffusionModel:
    """
    Load a diffusion model.

    Weights location may also be shortcut name, e.g. "SD-1.5"
    """
    from imaginairy.modules.refiners_sd import (
        SD1AutoencoderSliced,
        StableDiffusion_1,
        StableDiffusion_1_Inpainting,
    )

    global MOST_RECENTLY_LOADED_MODEL

    device = device or get_device()

    (
        vae_weights,
        unet_weights,
        text_encoder_weights,
    ) = load_stable_diffusion_compvis_weights(weights_location)

    StableDiffusionCls: type[LatentDiffusionModel]
    if for_inpainting:
        unet = SD1UNet(in_channels=9)
        StableDiffusionCls = StableDiffusion_1_Inpainting
    else:
        unet = SD1UNet(in_channels=4)
        StableDiffusionCls = StableDiffusion_1
    logger.debug(f"Using class {StableDiffusionCls.__name__}")

    sd = StableDiffusionCls(
        device=device, dtype=dtype, lda=SD1AutoencoderSliced(), unet=unet
    )
    logger.debug("Loading VAE")
    sd.lda.load_state_dict(vae_weights)

    logger.debug("Loading text encoder")
    sd.clip_text_encoder.load_state_dict(text_encoder_weights)

    logger.debug("Loading UNet")
    sd.unet.load_state_dict(unet_weights, strict=False)

    logger.debug(f"'{weights_location}' Loaded")

    MOST_RECENTLY_LOADED_MODEL = sd

    sd.set_self_attention_guidance(enable=True)

    return sd


@memory_managed_model("stable-diffusion", memory_usage_mb=1951)
def _load_diffusion_model(config_path, weights_location, half_mode):
    model_config = OmegaConf.load(f"{PKG_ROOT}/{config_path}")

    # only run half-mode on cuda. run it by default
    half_mode = half_mode is None and get_device() == "cuda"

    model = load_model_from_config(
        config=model_config,
        weights_location=weights_location,
        half_mode=half_mode,
    )
    return model


@memory_managed_model("controlnet")
def load_controlnet(control_weights_location, half_mode):
    controlnet_state_dict = load_state_dict(
        control_weights_location, half_mode=half_mode
    )
    controlnet_state_dict = {
        k.replace("control_model.", ""): v for k, v in controlnet_state_dict.items()
    }
    control_stage_config = OmegaConf.load(f"{PKG_ROOT}/configs/control-net-v15.yaml")[
        "model"
    ]["params"]["control_stage_config"]
    controlnet = instantiate_from_config(control_stage_config)
    controlnet.load_state_dict(controlnet_state_dict)
    controlnet.to(get_device())
    return controlnet


def resolve_model_weights_config(
    model_weights: str | iconfig.ModelWeightsConfig,
    default_model_architecture: str | None = None,
    for_inpainting: bool = False,
) -> iconfig.ModelWeightsConfig:
    """Resolve weight and config path if they happen to be shortcuts."""
    if isinstance(model_weights, iconfig.ModelWeightsConfig):
        return model_weights

    if not isinstance(model_weights, str):
        msg = f"Invalid model weights: {model_weights}"
        raise ValueError(msg)  # noqa

    if default_model_architecture is not None and not isinstance(
        default_model_architecture, str
    ):
        msg = f"Invalid model architecture: {default_model_architecture}"
        raise ValueError(msg)

    if for_inpainting:
        model_weights_config = iconfig.MODEL_WEIGHT_CONFIG_LOOKUP.get(
            f"{model_weights.lower()}-inpaint", None
        )
        if model_weights_config:
            return model_weights_config

    model_weights_config = iconfig.MODEL_WEIGHT_CONFIG_LOOKUP.get(
        model_weights.lower(), None
    )
    if model_weights_config:
        return model_weights_config

    if not default_model_architecture:
        msg = "You must specify the model architecture when loading custom weights."
        raise ValueError(msg)

    default_model_architecture = default_model_architecture.lower()
    model_architecture_config = None
    if for_inpainting:
        model_architecture_config = iconfig.MODEL_ARCHITECTURE_LOOKUP.get(
            f"{default_model_architecture}-inpaint", None
        )

    if not model_architecture_config:
        model_architecture_config = iconfig.MODEL_ARCHITECTURE_LOOKUP.get(
            default_model_architecture, None
        )

    if model_architecture_config is None:
        msg = f"Invalid model architecture: {default_model_architecture}"
        raise ValueError(msg)

    model_weights_config = iconfig.ModelWeightsConfig(
        name="Custom Loaded",
        aliases=[],
        architecture=model_architecture_config,
        weights_location=model_weights,
        defaults={},
    )

    return model_weights_config


def get_model_default_image_size(model_architecture: str | ModelArchitecture | None):
    if isinstance(model_architecture, str):
        model_architecture = iconfig.MODEL_ARCHITECTURE_LOOKUP.get(
            model_architecture, None
        )
    default_size = None
    if model_architecture:
        default_size = model_architecture.defaults.get("size")

    if default_size is None:
        default_size = 512
    default_size = normalize_image_size(default_size)
    return default_size


def get_current_diffusion_model():
    return MOST_RECENTLY_LOADED_MODEL


def get_cache_dir():
    xdg_cache_home = os.getenv("XDG_CACHE_HOME", None)
    if xdg_cache_home is None:
        user_home = os.getenv("HOME", None)
        if user_home:
            xdg_cache_home = os.path.join(user_home, ".cache")

    if xdg_cache_home is not None:
        return os.path.join(xdg_cache_home, "imaginairy")

    return os.path.join(os.path.dirname(__file__), ".cached-aimg")


def get_cached_url_path(url, category=None):
    """
    Gets the contents of a url, but caches the response indefinitely.

    While we attempt to use the cached_path from huggingface transformers, we fall back
    to our own implementation if the url does not provide an etag header, which `cached_path`
    requires.  We also skip the `head` call that `cached_path` makes on every call if the file
    is already cached.
    """

    try:
        if url.startswith("https://huggingface.co"):
            return huggingface_cached_path(url)
    except (OSError, ValueError):
        pass
    filename = url.split("/")[-1]
    dest = get_cache_dir()
    if category:
        dest = os.path.join(dest, category)
    os.makedirs(dest, exist_ok=True)

    # Replace possibly illegal destination path characters
    safe_filename = re.sub('[*<>:"|?]', "_", filename)
    dest_path = os.path.join(dest, safe_filename)
    if os.path.exists(dest_path):
        return dest_path

    # check if it's saved at previous path and rename it
    old_dest_path = os.path.join(dest, filename)
    if os.path.exists(old_dest_path):
        os.rename(old_dest_path, dest_path)
        return dest_path

    r = requests.get(url)

    with open(dest_path, "wb") as f:
        f.write(r.content)
    return dest_path


def check_huggingface_url_authorized(url):
    if not url.startswith("https://huggingface.co/"):
        return None
    token = HfFolder.get_token()
    headers = {}
    if token is not None:
        headers["authorization"] = f"Bearer {token}"
    response = requests.head(url, allow_redirects=True, headers=headers, timeout=5)
    if response.status_code == 401:
        msg = "Unauthorized access to HuggingFace model. This model requires a huggingface token.  Please login to HuggingFace or set HUGGING_FACE_HUB_TOKEN to your User Access Token. See https://huggingface.co/docs/huggingface_hub/quick-start#login for more information"
        raise HuggingFaceAuthorizationError(msg)
    return None


@wraps(_hf_hub_download)
def hf_hub_download(*args, **kwargs):
    """
    backwards compatible wrapper for huggingface's hf_hub_download.

    they changed the argument name from `use_auth_token` to `token`
    """

    try:
        return _hf_hub_download(*args, **kwargs)
    except TypeError as e:
        if "unexpected keyword argument 'token'" in str(e):
            kwargs["use_auth_token"] = kwargs.pop("token")
            return _hf_hub_download(*args, **kwargs)
        raise


def huggingface_cached_path(url):
    # bypass all the HEAD calls done by the default `cached_path`
    repo, commit_hash, filepath = extract_huggingface_repo_commit_file_from_url(url)
    dest_path = try_to_load_from_cache(
        repo_id=repo, revision=commit_hash, filename=filepath
    )
    if not dest_path:
        check_huggingface_url_authorized(url)
        token = HfFolder.get_token()
        logger.info(f"Downloading {url} from huggingface")
        dest_path = hf_hub_download(
            repo_id=repo, revision=commit_hash, filename=filepath, token=token
        )
        # make a refs folder so caching works
        # work-around for
        # https://github.com/huggingface/huggingface_hub/pull/1306
        # https://github.com/brycedrennan/imaginAIry/issues/171
        refs_url = dest_path[: dest_path.index("/snapshots/")] + "/refs/"
        os.makedirs(refs_url, exist_ok=True)
    return dest_path


def extract_huggingface_repo_commit_file_from_url(url):
    parsed_url = urllib.parse.urlparse(url)
    path_components = parsed_url.path.strip("/").split("/")

    repo = "/".join(path_components[0:2])
    assert path_components[2] == "resolve"
    commit_hash = path_components[3]
    filepath = "/".join(path_components[4:])

    return repo, commit_hash, filepath


def download_diffusers_weights(repo, sub, filename):
    from imaginairy.utils.model_manager import get_cached_url_path

    url = f"https://huggingface.co/{repo}/resolve/main/{sub}/{filename}"
    return get_cached_url_path(url, category="weights")


@lru_cache
def load_stable_diffusion_diffusers_weights(diffusers_repo, device=None):
    from imaginairy.utils import get_device
    from imaginairy.weight_management.conversion import cast_weights
    from imaginairy.weight_management.utils import (
        COMPONENT_NAMES,
        FORMAT_NAMES,
        MODEL_NAMES,
    )

    if device is None:
        device = get_device()
    vae_weights_path = download_diffusers_weights(
        repo=diffusers_repo, sub="vae", filename="diffusion_pytorch_model.safetensors"
    )
    vae_weights = open_weights(vae_weights_path, device=device)
    vae_weights = cast_weights(
        source_weights=vae_weights,
        source_model_name=MODEL_NAMES.SD15,
        source_component_name=COMPONENT_NAMES.VAE,
        source_format=FORMAT_NAMES.DIFFUSERS,
        dest_format=FORMAT_NAMES.REFINERS,
    )

    unet_weights_path = download_diffusers_weights(
        repo=diffusers_repo, sub="unet", filename="diffusion_pytorch_model.safetensors"
    )
    unet_weights = open_weights(unet_weights_path, device=device)
    unet_weights = cast_weights(
        source_weights=unet_weights,
        source_model_name=MODEL_NAMES.SD15,
        source_component_name=COMPONENT_NAMES.UNET,
        source_format=FORMAT_NAMES.DIFFUSERS,
        dest_format=FORMAT_NAMES.REFINERS,
    )

    text_encoder_weights_path = download_diffusers_weights(
        repo=diffusers_repo, sub="text_encoder", filename="model.safetensors"
    )
    text_encoder_weights = open_weights(text_encoder_weights_path, device=device)
    text_encoder_weights = cast_weights(
        source_weights=text_encoder_weights,
        source_model_name=MODEL_NAMES.SD15,
        source_component_name=COMPONENT_NAMES.TEXT_ENCODER,
        source_format=FORMAT_NAMES.DIFFUSERS,
        dest_format=FORMAT_NAMES.REFINERS,
    )

    return vae_weights, unet_weights, text_encoder_weights


def open_weights(filepath, device=None):
    from imaginairy.utils import get_device

    if device is None:
        device = get_device()

    if "safetensor" in filepath.lower():
        from refiners.fluxion.utils import safe_open

        with safe_open(path=filepath, framework="pytorch", device=device) as tensors:
            state_dict = {
                key: tensors.get_tensor(key) for key in tensors.keys()  # noqa
            }
    else:
        import torch

        state_dict = torch.load(filepath, map_location=device)

    while "state_dict" in state_dict:
        state_dict = state_dict["state_dict"]

    return state_dict


def load_stable_diffusion_compvis_weights(weights_url):
    from imaginairy.utils import get_device
    from imaginairy.utils.model_manager import get_cached_url_path
    from imaginairy.weight_management.conversion import cast_weights
    from imaginairy.weight_management.utils import (
        COMPONENT_NAMES,
        FORMAT_NAMES,
        MODEL_NAMES,
    )

    weights_path = get_cached_url_path(weights_url, category="weights")
    logger.info(f"Loading weights from {weights_path}")
    state_dict = open_weights(weights_path, device=get_device())

    text_encoder_prefix = "cond_stage_model."
    cut_start = len(text_encoder_prefix)
    text_encoder_state_dict = {
        k[cut_start:]: v
        for k, v in state_dict.items()
        if k.startswith(text_encoder_prefix)
    }
    text_encoder_state_dict = cast_weights(
        source_weights=text_encoder_state_dict,
        source_model_name=MODEL_NAMES.SD15,
        source_component_name=COMPONENT_NAMES.TEXT_ENCODER,
        source_format=FORMAT_NAMES.COMPVIS,
        dest_format=FORMAT_NAMES.DIFFUSERS,
    )
    text_encoder_state_dict = cast_weights(
        source_weights=text_encoder_state_dict,
        source_model_name=MODEL_NAMES.SD15,
        source_component_name=COMPONENT_NAMES.TEXT_ENCODER,
        source_format=FORMAT_NAMES.DIFFUSERS,
        dest_format=FORMAT_NAMES.REFINERS,
    )

    vae_prefix = "first_stage_model."
    cut_start = len(vae_prefix)
    vae_state_dict = {
        k[cut_start:]: v for k, v in state_dict.items() if k.startswith(vae_prefix)
    }
    vae_state_dict = cast_weights(
        source_weights=vae_state_dict,
        source_model_name=MODEL_NAMES.SD15,
        source_component_name=COMPONENT_NAMES.VAE,
        source_format=FORMAT_NAMES.COMPVIS,
        dest_format=FORMAT_NAMES.DIFFUSERS,
    )
    vae_state_dict = cast_weights(
        source_weights=vae_state_dict,
        source_model_name=MODEL_NAMES.SD15,
        source_component_name=COMPONENT_NAMES.VAE,
        source_format=FORMAT_NAMES.DIFFUSERS,
        dest_format=FORMAT_NAMES.REFINERS,
    )

    unet_prefix = "model."
    cut_start = len(unet_prefix)
    unet_state_dict = {
        k[cut_start:]: v for k, v in state_dict.items() if k.startswith(unet_prefix)
    }
    unet_state_dict = cast_weights(
        source_weights=unet_state_dict,
        source_model_name=MODEL_NAMES.SD15,
        source_component_name=COMPONENT_NAMES.UNET,
        source_format=FORMAT_NAMES.COMPVIS,
        dest_format=FORMAT_NAMES.DIFFUSERS,
    )

    unet_state_dict = cast_weights(
        source_weights=unet_state_dict,
        source_model_name=MODEL_NAMES.SD15,
        source_component_name=COMPONENT_NAMES.UNET,
        source_format=FORMAT_NAMES.DIFFUSERS,
        dest_format=FORMAT_NAMES.REFINERS,
    )

    return vae_state_dict, unet_state_dict, text_encoder_state_dict