import os
import sys
from pathlib import Path
from fal.toolkit import FAL_PERSISTENT_DIR, clone_repository

from serverless_comfy.utils import patch_onnx_runtime
from serverless_comfy.comfy_model_registry import (
    CUSTOM_NODE_MODELS,
    MODELS,
    cleanup_unused_models,
    download_models,
)
from serverless_comfy.custom_comfy_nodes import (
    CUSTOM_NODES,
    cleanup_unused_custom_nodes,
    clone_custom_nodes,
    filter_custom_node_models,
)

__CURR_DIR = Path(__file__).resolve().parent


def get_requirements():
    with open(__CURR_DIR / "requirements.txt") as fp:
        requirements = fp.read().splitlines()
        requirements = [req for req in requirements if not req.startswith("#")]

    return requirements


def setup_insightface(comfy_repo_path: Path):
    try:
        import insightface
    except ImportError:
        print("Insightface is not installed")
        return

    INSIGHTFACE_HOME = str(FAL_PERSISTENT_DIR / ".insightface")

    # we make sure that the directory exists
    Path(INSIGHTFACE_HOME).mkdir(parents=True, exist_ok=True)

    os.environ["INSIGHTFACE_HOME"] = INSIGHTFACE_HOME

    insightface.utils.storage.BASE_REPO_URL = (
        "https://storage.googleapis.com/fal-models/insightface"
    )
    insightface.utils.storage.INSIGHTFACE_HOME = INSIGHTFACE_HOME

    models_dir = comfy_repo_path / "models" / "insightface"

    try:
        models_dir.symlink_to(INSIGHTFACE_HOME)
    except FileExistsError:
        pass


def setup_comfy(
    repo_name: str,
    custom_nodes: list[list[str]] | None = None,
    models: dict[str, list[list[str]]] | None = None,
    custom_node_models: dict[str, list[list[str]]] | None = None,
    force: bool = False,
    commit_hash: str = "c5de4955bb91a2b136027a698aaecb8d19e3d892",
    should_cleanup_unused_models: bool = True,
):
    # WARNING: After updating the commit hash, make sure to look for the
    # changes in `comfy_overrides.py` and update the patched functions
    # accordingly.

    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
    os.environ["HF_HOME"] = f"{FAL_PERSISTENT_DIR}/.cache/huggingface"
    os.environ["TRANSFORMERS_CACHE"] = (
        f"{FAL_PERSISTENT_DIR}/.cache/huggingface/transformers"
    )
    os.environ["TORCH_HOME"] = f"{FAL_PERSISTENT_DIR}/.cache/torch"

    repo_path = clone_repository(
        "https://github.com/comfyanonymous/ComfyUI",
        commit_hash=commit_hash,
        repo_name=repo_name,
        include_to_path=True,
        force=force,
    )

    sys.argv = [
        sys.executable,
        "--input-directory",
        "/tmp/comfy/input",
        "--output-directory",
        "/tmp/comfy/output",
        "--temp-directory",
        "/tmp/comfy",
        # NOTE: this gives us a slight speedup, but gets us into CUDA OOMs
        # "--gpu-only",
        "--force-fp16",
    ]

    # Used by Fannovel16/comfyui_controlnet_aux
    # https://github.com/Fannovel16/comfyui_controlnet_aux/blob/4cd233c5d7afe2e51bf57ee7a5ba7e6fcb9cbb22/src/custom_controlnet_aux/util.py#L299-L301
    os.makedirs("/tmp/ckpts", exist_ok=True)

    models = models if models is not None else MODELS

    custom_nodes = custom_nodes if custom_nodes is not None else CUSTOM_NODES
    custom_node_models = (
        custom_node_models if custom_node_models is not None else CUSTOM_NODE_MODELS
    )

    custom_node_models = filter_custom_node_models(custom_nodes, custom_node_models)

    cleanup_input_directory()

    setup_insightface(repo_path)
    cloned_nodes = clone_custom_nodes(repo_path, custom_nodes=custom_nodes)
    cleanup_unused_custom_nodes(repo_path, cloned_nodes=cloned_nodes)

    download_models(repo_path, models=models, custom_node_models=custom_node_models)
    if should_cleanup_unused_models:
        cleanup_unused_models(
            repo_path, models=models, custom_node_models=custom_node_models
        )
    patch_onnx_runtime()
    patch_facexlib()
    patch_rembg()

    os.environ["COMFYUI_EASY_USE_DISABLE_STABILITY_API"] = "1"

    return repo_path


def cleanup_input_directory():
    input_dir = Path("/tmp/comfy/input")
    if input_dir.exists():
        for file in input_dir.iterdir():
            file.unlink()


def patch_facexlib():
    try:
        import facexlib.utils.misc
    except ImportError:
        print("Facexlib is not installed")
        return

    FACEXLIB_HOME = str(FAL_PERSISTENT_DIR / ".facexlib")

    os.environ["FACEXLIB_HOME"] = FACEXLIB_HOME
    facexlib.utils.misc.ROOT_DIR = FACEXLIB_HOME


def patch_rembg():
    os.environ["U2NET_HOME"] = str(Path(FAL_PERSISTENT_DIR) / "remove_background_model")
    os.environ["NUMBA_CACHE_DIR"] = str(Path(FAL_PERSISTENT_DIR) / ".numba_cache")
