from pathlib import Path

from serverless_comfy.download_utils import download_model_weights

from .animatediff_models import ANIMATEDIFF_MODELS
from .animatediff_motion_lora import ANIMATEDIFF_MOTION_LORA_MODELS
from .bert_base_uncased import BERT_BASE_UNCASED
from .checkpoints import CHECKPOINTS
from .clip import CLIP
from .clip_vision import CLIP_VISION
from .controlnet import CONTROLNET
from .custom_nodes import BRIA_REMBG_MODELS
from .embeddings import EMBEDDINGS
from .facedetection import FACEDETECTION
from .facerestore_models import FACERESTORE_MODELS
from .gligen import GLIGEN
from .grounding_dino import GROUNDING_DINO
from .hypernetworks import HYPERNETWORKS
from .insightface import INSIGHTFACE
from .instantid import INSTANTID
from .ipadapter import IPADAPTER
from .llava_checkpoints import LLAVA_CHECKPOINTS
from .loras import LORAS
from .mmdets import MMDETS
from .photomaker import PHOTOMAKER
from .pulid import PULID
from .safety_checker import SAFETY_CHECKER
from .sams import SAMS
from .style_models import STYLE_MODELS
from .ultralytics import ULTRALYTICS
from .unet import UNET
from .upscale_models import UPSCALE_MODELS
from .vae import VAE
from .vae_approx import VAE_APPROX
from .xlabs import XLABS
from .inpaint import INPAINT
from .infiniteyou import INFINITE_YOU
from .seedvr2 import SEEDVR2

# New model categories
from .diffusion_models import DIFFUSION_MODELS
from .text_encoders import TEXT_ENCODERS
from .audio_encoders import AUDIO_ENCODERS
from .model_patches import MODEL_PATCHES
from .latent_upscale_models import LATENT_UPSCALE_MODELS

MODELS = {
    "animatediff_models": ANIMATEDIFF_MODELS,
    "animatediff_motion_lora": ANIMATEDIFF_MOTION_LORA_MODELS,
    "bert-base-uncased": BERT_BASE_UNCASED,
    "checkpoints": CHECKPOINTS,
    "clip": CLIP,
    "clip_vision": CLIP_VISION,
    "controlnet": CONTROLNET,
    "embeddings": EMBEDDINGS,
    "facedetection": FACEDETECTION,
    "facerestore_models": FACERESTORE_MODELS,
    "gligen": GLIGEN,
    "grounding-dino": GROUNDING_DINO,
    "hypernetworks": HYPERNETWORKS,
    "insightface": INSIGHTFACE,
    "instantid": INSTANTID,
    "ipadapter": IPADAPTER,
    "loras": LORAS,
    "mmdets": MMDETS,
    "photomaker": PHOTOMAKER,
    "sams": SAMS,
    "style_models": STYLE_MODELS,
    "ultralytics": ULTRALYTICS,
    "unet": UNET,
    "upscale_models": UPSCALE_MODELS,
    "vae": VAE,
    "pulid": PULID,
    "vae_approx": VAE_APPROX,
    "LLavacheckpoints": LLAVA_CHECKPOINTS,
    "safety_checker": SAFETY_CHECKER,
    "xlabs": XLABS,
    "inpaint": INPAINT,
    "InfiniteYou": INFINITE_YOU,
    "SEEDVR2": SEEDVR2,
    # New model categories
    "diffusion_models": DIFFUSION_MODELS,
    "text_encoders": TEXT_ENCODERS,
    "audio_encoders": AUDIO_ENCODERS,
    "model_patches": MODEL_PATCHES,
    "latent_upscale_models": LATENT_UPSCALE_MODELS,
}

CUSTOM_NODE_MODELS = {
    "ComfyUI-BRIA_AI-RMBG/RMBG-1.4": BRIA_REMBG_MODELS,
}


def download_models(
    comfy_repo_path: str | Path,
    models: dict[str, list[list[str]]],
    custom_node_models: dict[str, list[list[str]]],
):
    comfy_models_dir = Path(comfy_repo_path) / "models"
    custom_nodes_dir = Path(comfy_repo_path) / "custom_nodes"

    for subfolder, model_list in models.items():
        for model in model_list:
            try:
                model_name, model_url = model
            except ValueError:
                print(f"Error unpacking model {model}")
                raise
            # Issue: the consumers can see the logs of downloading models
            # including the URL of our private models
            # print("Downloading", model_name, "from", model_url)
            download_and_symlink(model_url, comfy_models_dir / subfolder / model_name)

    for subfolder, model_list in custom_node_models.items():
        for model in model_list:
            model_name, model_url = model
            download_and_symlink(model_url, custom_nodes_dir / subfolder / model_name)


def download_and_symlink(url: str, path: str | Path):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    downloaded_path = download_model_weights(url)

    # if url ends with tar then it needs to be extracted
    if not path.exists():
        try:
            if url.endswith(".tar"):
                # TODO: impl.
                return
            else:
                path.symlink_to(downloaded_path)
        except FileExistsError:
            # If the file already exists, remove it and create a new symlink
            path.unlink()
            path.symlink_to(downloaded_path)


def cleanup_unused_models(
    comfy_repo_path: str | Path,
    models: dict[str, list[list[str]]],
    custom_node_models: dict[str, list[list[str]]],
):
    comfy_models_dir = Path(comfy_repo_path) / "models"

    for subfolder, model_list in models.items():
        model_dir = comfy_models_dir / subfolder

        model_dirs_to_check = [
            comfy_models_dir / "checkpoints",
            comfy_models_dir / "unet",
            comfy_models_dir / "vae",
            comfy_models_dir / "loras",
        ]

        model_set = {model[0] for model in model_list}

        if model_dir not in model_dirs_to_check:
            continue

        print(f"Cleaning up {model_dir.name}")
        for model in model_dir.iterdir():
            if model.is_file() and model.name not in model_set:
                try:
                    print(f"Removing {model.name}")
                    model.unlink(missing_ok=True)
                except Exception as e:
                    print(e)
                    pass
