import os
import shutil
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import uuid

from fal.toolkit import clone_repository, FAL_REPOSITORY_DIR
from serverless_comfy.utils import get_impact_subpack_location


COMFYUI_IMPACT_PACK_URL = "https://github.com/ltdrdata/ComfyUI-Impact-Pack.git"

# URL of the custom node and the commit hash
# TRIMMED VERSION - Only essential nodes to speed up startup
CUSTOM_NODES = [
    # =====================================================================
    # FAL-SPECIFIC NODES (kept from original)
    # =====================================================================
    [
        "https://github.com/badayvedat/ComfyUI-fal-Connector.git",
        "ea3005c84be64ca0f7aab3e3eb0b62fda0ba24a1",
    ],
    [
        "https://{GITHUB_TOKEN}@github.com/fal-ai-community/Gourieff-ComfyUI-ReActor.git",
        "88ce71f7ecd42cba2ed3f5cb87f076b8b7a6728b",
    ],
    [
        "https://{GITHUB_TOKEN}@github.com/fal-ai-community/huchenlei-ComfyUI-layerdiffuse.git",
        "f747d4d7be6f6bb003688c9004abbd30f37f9945",
    ],
    [
        "https://{GITHUB_TOKEN}@github.com/fal-ai-community/42lux-ComfyUI-safety-checker.git",
        "03b73f436f5383e0c6eb0e26fb84d84cff05c5ea",
    ],
    [
        "https://github.com/mavibirdesmi/ComfyUI_LayerStyle_Advance.git",
        "1a3c371423f70eb28309ad3b5ae0f02bd7a8c006",
    ],
    # TODO: change this back to gokayfem when the PR is merged https://github.com/gokayfem/ComfyUI_VLM_nodes/pull/154
    [
        "https://github.com/mavibirdesmi/ComfyUI_VLM_nodes",
        "858ab8a13e2a3eaa328b3dc42f03af9b9ffbcf83",
    ],
    [
        "https://github.com/badayvedat/ComfyUI-Easy-Use.git",
        "b1ddd7b4d998c547c886be8d5bb4ab7f4e4a4579",
    ],
    # =====================================================================
    # ESSENTIAL NODES (1000+ stars or critical functionality)
    # =====================================================================
    # ComfyUI_IPAdapter_plus (5640 stars)
    [
        "https://github.com/cubiq/ComfyUI_IPAdapter_plus",
        "a0f451a5113cf9becb0847b92884cb10cbdec0ef",
    ],
    # ComfyUI-WanVideoWrapper (5469 stars)
    [
        "https://github.com/kijai/ComfyUI-WanVideoWrapper",
        "e867e642d43f2ff0a66125df2d59e3f40bdd4082",
    ],
    # comfyui_controlnet_aux (3620 stars)
    [
        "https://github.com/Fannovel16/comfyui_controlnet_aux",
        "12f35647f0d510e03b45a47fb420fe1245a575df",
    ],
    # ComfyUI-AnimateDiff-Evolved (3336 stars)
    [
        "https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved",
        "90fb1331201a4b29488089e4fbffc0d82cc6d0a9",
    ],
    # ComfyUI-GGUF (2920 stars)
    [
        "https://github.com/city96/ComfyUI-GGUF",
        "01f8845bf30d89fff293c7bd50187bc59d9d53ea",
    ],
    # ComfyUI-Impact-Pack (2848 stars)
    [
        COMFYUI_IMPACT_PACK_URL,
        "ba09fbc4c05688667e7bfda35823933c9fe09196",
    ],
    [
        "https://github.com/ltdrdata/ComfyUI-Impact-Subpack.git",
        "50c7b71a6a224734cc9b21963c6d1926816a97f1",
    ],
    # ComfyUI-Custom-Scripts (2831 stars)
    [
        "https://github.com/pythongosssss/ComfyUI-Custom-Scripts",
        "f2838ed5e59de4d73cde5c98354b87a8d3200190",
    ],
    # ComfyUI_LayerStyle (2760 stars)
    [
        "https://github.com/chflame163/ComfyUI_LayerStyle",
        "90f4bfb38aaf121292f429b20eff07c6c121697e",
    ],
    # ComfyUI-HunyuanVideoWrapper (2559 stars)
    [
        "https://github.com/kijai/ComfyUI-HunyuanVideoWrapper",
        "fcbd6729a9b0b8ff6037c598bbada4a6bdc6d967",
    ],
    # ComfyUI-LTXVideo (2446 stars) - using wzgrx fork with apply_rotary_emb fix
    [
        "https://github.com/wzgrx/ComfyUI-LTXVideo",
        "36cc0ac19e5b26f794037a520ecd98c7fa945bac",
    ],
    # ComfyUI-KJNodes (2009 stars)
    [
        "https://github.com/kijai/ComfyUI-KJNodes",
        "488e6607fe0261928c96e6ed6425cb82fa2a1ccf",
    ],
    # ComfyUI_InstantID (1757 stars)
    [
        "https://github.com/cubiq/ComfyUI_InstantID",
        "72495e806bc2ab9c41581e15ccaa1bcf83c477e8",
    ],
    # was-node-suite-comfyui (1693 stars)
    [
        "https://github.com/WASasquatch/was-node-suite-comfyui",
        "ea935d1044ae5a26efa54ebeb18fe9020af49a45",
    ],
    # x-flux-comfyui (1676 stars)
    [
        "https://github.com/XLabs-AI/x-flux-comfyui",
        "00328556efc9472410d903639dc9e68a8471f7ac",
    ],
    # ComfyUI-Crystools (1618 stars)
    [
        "https://github.com/crystian/ComfyUI-Crystools",
        "2f18256c5b5063937106f29a8e0a7db3ae3869b7",
    ],
    # ComfyUI-FramePackWrapper (1598 stars)
    [
        "https://github.com/kijai/ComfyUI-FramePackWrapper",
        "a7c4b704455aee0d016143f2fc232928cc0f1d83",
    ],
    # ComfyUI-RMBG (1544 stars)
    [
        "https://github.com/1038lab/ComfyUI-RMBG",
        "bd509b4750c81221b938684489c87187b0172208",
    ],
    # ComfyUI-CogVideoXWrapper (1534 stars)
    [
        "https://github.com/kijai/ComfyUI-CogVideoXWrapper",
        "fdb8abd2790b5459ddc7066c31861bb0b62e988b",
    ],
    # ComfyUI-Florence2 (1518 stars)
    [
        "https://github.com/kijai/ComfyUI-Florence2",
        "00b63382966a444a9fefacb65b8deb188d12a458",
    ],
    # ComfyUI-Fluxtapoz (1391 stars)
    [
        "https://github.com/logtd/ComfyUI-Fluxtapoz",
        "17c71bea20e932945e9c1be0586cfd4b7e51cbf6",
    ],
    # ComfyUI_UltimateSDUpscale (1369 stars)
    [
        "https://github.com/ssitu/ComfyUI_UltimateSDUpscale",
        "c48d60df1ee24b47c4e6acacaf1038aff2698cf3",
    ],
    # ComfyUI-VideoHelperSuite (1368 stars)
    [
        "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite",
        "537f3a02269e14ab4b0350d4189a52034e9a3b12",
    ],
    # ComfyUI-segment-anything-2 (1139 stars)
    [
        "https://github.com/kijai/ComfyUI-segment-anything-2",
        "0c35fff5f382803e2310103357b5e985f5437f32",
    ],
    # ComfyUI-IC-Light (1118 stars)
    [
        "https://github.com/kijai/ComfyUI-IC-Light",
        "22811d903af8a90100dccef1c4b579526649a6ab",
    ],
    # ComfyUI_Comfyroll_CustomNodes (1098 stars)
    [
        "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes",
        "d78b780ae43fcf8c6b7c6505e6ffb4584281ceca",
    ],
    # comfyui-inpaint-nodes (1082 stars)
    [
        "https://github.com/Acly/comfyui-inpaint-nodes",
        "1e2543e5619b3f6e88b7727243c8b9c5791a7b4a",
    ],
    # comfyui_segment_anything (1055 stars)
    [
        "https://github.com/storyicon/comfyui_segment_anything",
        "ab6395596399d5048639cdab7e44ec9fae857a93",
    ],
    # ComfyUI-WD14-Tagger (1038 stars)
    [
        "https://github.com/pythongosssss/ComfyUI-WD14-Tagger",
        "9e0a6e700299182fc05c58b62e7ad9f72182a78b",
    ],
    # ComfyUI_essentials (998 stars)
    [
        "https://github.com/cubiq/ComfyUI_essentials",
        "9d9f4bedfc9f0321c19faf71855e228c93bd0dc9",
    ],
    # ComfyUI-Advanced-ControlNet (915 stars)
    [
        "https://github.com/Kosinkadink/ComfyUI-Advanced-ControlNet",
        "2bde95a468ce5fd3959f646258606ae221fa1e17",
    ],
    # ComfyUI-Frame-Interpolation (905 stars)
    [
        "https://github.com/Fannovel16/ComfyUI-Frame-Interpolation",
        "a969c01dbccd9e5510641be04eb51fe93f6bfc3d",
    ],
    # PuLID_ComfyUI (903 stars)
    [
        "https://github.com/cubiq/PuLID_ComfyUI",
        "93e0c4c226b87b23c0009d671978bad0e77289ff",
    ],
    # sdxl_prompt_styler (895 stars)
    [
        "https://github.com/twri/sdxl_prompt_styler",
        "51068179927f79dce14f38c6b1984390ab242be2",
    ],
    # ComfyUI-MochiWrapper (799 stars)
    [
        "https://github.com/kijai/ComfyUI-MochiWrapper",
        "e1bd05240ac31b72166e9d952c75dd5735352311",
    ],
    # ComfyUI-Inspire-Pack (712 stars)
    [
        "https://github.com/ltdrdata/ComfyUI-Inspire-Pack",
        "d23db9aa544de9a6d4c609cb7005fa9e0d42031d",
    ],
    # ComfyUI-PuLID-Flux (704 stars)
    [
        "https://github.com/balazik/ComfyUI-PuLID-Flux",
        "a80912fc3435c358607bf4b43a58dbcbebdb09ff",
    ],
    # efficiency-nodes-comfyui (662 stars) - DISABLED: CompVisVDenoiser removed from comfy.samplers
    # [
    #     "https://github.com/LucianoCirino/efficiency-nodes-comfyui",
    #     "040e38864616d2359d279b706c3c2e8569f11db2",
    # ],
    # comfy_mtb (646 stars)
    [
        "https://github.com/melMass/comfy_mtb",
        "eabe43db79112d4e3674db343c5c184145e8396c",
    ],
    # ComfyUI-IC-Light-Native (637 stars)
    [
        "https://github.com/huchenlei/ComfyUI-IC-Light-Native",
        "40883a839d2a719ad3d3a454700609b94b1a9c65",
    ],
    # comfyui-tooling-nodes (593 stars)
    [
        "https://github.com/Acly/comfyui-tooling-nodes",
        "02202529129dd33d0622dcc72d0550d10bf40d62",
    ],
    # ComfyUI-Marigold (557 stars)
    [
        "https://github.com/kijai/ComfyUI-Marigold",
        "2adfec75c430b0751a98a82453dc2dd8077bca85",
    ],
    # ComfyUI-DepthAnythingV2 (378 stars)
    [
        "https://github.com/kijai/ComfyUI-DepthAnythingV2",
        "d505cbca99803fc63327b8305618a23e59a18b42",
    ],
    # comfyui-dynamicprompts (370 stars)
    [
        "https://github.com/adieyal/comfyui-dynamicprompts",
        "3f2fff32358cf39e21b8b440ca87eac9a8e2bade",
    ],
    # =====================================================================
    # ADDITIONAL ESSENTIAL NODES
    # =====================================================================
    [
        "https://github.com/mav-rik/facerestore_cf.git",
        "67f90bc6be976fb58169866155346b0da13bebee",
    ],
    [
        "https://github.com/BlenderNeko/ComfyUI_ADV_CLIP_emb",
        "63984deefb005da1ba90a1175e21d91040da38ab",
    ],
    [
        "https://github.com/comfyanonymous/ComfyUI_experiments.git",
        "934dba9d206e4738e0dac26a09b51f1dffcb4e44",
    ],
    [
        "https://github.com/BlenderNeko/ComfyUI_TiledKSampler",
        "2fd9b05d97ecffc604c642ffbb40220b182966b2",
    ],
    # Original fal nodes
    [
        "https://github.com/a1lazydog/ComfyUI-AudioScheduler.git",
        "fdf8ddd0f6934c031603b145a158e2060c715ca6",
    ],
    [
        "https://github.com/JettHu/ComfyUI-TCD.git",
        "888a8403b791333ce98c910a3b76f1842e1adee8",
    ],
    [
        "https://github.com/fofr/ComfyUI-HyperSDXL1StepUnetScheduler.git",
        "32471614dca844ac7be9a536d3cc6597e12e080a",
    ],
    [
        "https://github.com/edenartlab/eden_comfy_pipelines.git",
        "69d7f4f204a38626113686ec6c56281105419f1c",
    ],
    [
        "https://github.com/gokayfem/ComfyUI-PuLID-Flux-Enhanced.git",
        "d5214dccd7cf518cd7d38b83b7be3dca54304288",
    ],
    [
        "https://github.com/asagi4/ComfyUI-Adaptive-Guidance.git",
        "181641ca04ac524992f40e8717f5da01eb82ea79",
    ],
    [
        "https://github.com/jamesWalker55/comfyui-various.git",
        "5bd85aaf7616878471469c4ec7e11bbd0cef3bf2",
    ],
    [
        "https://github.com/ZenAI-Vietnam/ComfyUI_InfiniteYou.git",
        "8eb69d39170444b2370676a2f42358d9ef1572c0",
    ],
    [
        "https://github.com/theUpsider/ComfyUI-Logic.git",
        "42d4f3df45fb7f0dd6e2201a14c07d4dd09f235d",
    ],
    [
        "https://github.com/ostris/ComfyUI-Advanced-Vision.git",
        "a0d6e805fb3623a321c210c0dd1501fe5923eeb1",
    ],
    [
        "https://github.com/romakot321/romakot321-comfy.git",
        "319689ac08290255053e20d2d65bbba825548b0e",
    ],
    [
        "https://github.com/yichengup/Comfyui_Flux_Style_Adjust.git",
        "83de37b8c22b8ac219b7fe5cbb3723ad8a8ab024",
    ],
    [
        "https://github.com/digitaljohn/comfyui-propost.git",
        "df6a6d122498f57ad7195d58e07701a501c9dcb6",
    ],
    [
        "https://github.com/kijai/ComfyUI-Lotus.git",
        "07719d3195b9a31839e3c6c4875a9eb9d6b1734c",
    ],
    [
        "https://github.com/hike2008/ComfyUI_ADV_CLIP_emb",
        "6e51c704eceba2015f7dc89dfd8d77d95222be12",
    ],
    # =====================================================================
    # MISSING FROM MAIN BRANCH - ADDING BACK
    # =====================================================================
    # ComfyUI-BRIA_AI-RMBG
    [
        "https://github.com/ZHO-ZHO-ZHO/ComfyUI-BRIA_AI-RMBG.git",
        "827fcd63ff0cfa7fbc544b8d2f4c1e3f3012742d",
    ],
    # ComfyUI_FizzNodes
    [
        "https://github.com/FizzleDorf/ComfyUI_FizzNodes.git",
        "7d6ea60c55ebd1268bd76fa462da052852bff192",
    ],
    # Derfuu_ComfyUI_ModdedNodes
    [
        "https://github.com/Derfuu/Derfuu_ComfyUI_ModdedNodes.git",
        "d0905bed31249f2bd0814c67585cf4fe3c77c015",
    ],
    # masquerade-nodes-comfyui
    [
        "https://github.com/BadCafeCode/masquerade-nodes-comfyui.git",
        "432cb4d146a391b387a0cd25ace824328b5b61cf",
    ],
    # ComfyUI-AutomaticCFG
    [
        "https://github.com/Extraltodeus/ComfyUI-AutomaticCFG",
        "2e395317b65c05a97a0ef566c4a8c7969305dafa",
    ],
    # ComfyUI-BrushNet
    [
        "https://github.com/nullquant/ComfyUI-BrushNet.git",
        "505d8ef917ddf3896afd1926770ecc9b099704e2",
    ],
    # ComfyUI-SUPIR
    [
        "https://github.com/kijai/ComfyUI-SUPIR.git",
        "c5cb1ce099c56c2bdea57e19beb5105bdaa7c429",
    ],
    # ComfyUI_AdvancedRefluxControl
    [
        "https://github.com/kaibioinfo/ComfyUI_AdvancedRefluxControl",
        "2b95c2c866399ca1914b4da486fe52808f7a9c60",
    ],
    # ComfyUI-Inpaint-CropAndStitch
    [
        "https://github.com/lquesada/ComfyUI-Inpaint-CropAndStitch.git",
        "a5fbe06f766ec237f304fb65bb94c0f255060109",
    ],
    # ComfyUI-Detail-Daemon
    [
        "https://github.com/Jonseed/ComfyUI-Detail-Daemon.git",
        "f62d51009363083178e0ed8eb20d843ca130b58b",
    ],
    # ComfyUI-TeaCache
    [
        "https://github.com/welltop-cn/ComfyUI-TeaCache.git",
        "91dff8e31684ca70a5fda309611484402d8fa192",
    ],
    # ComfyUI-SeedVR2_VideoUpscaler
    [
        "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler.git",
        "58bc9e8bc946499352e0cb3a9fe0d0a61fd86791",
    ],
    # ComfyUI_Fill-Nodes - DISABLED: requires google-genai which conflicts with websockets<13.0 (gradio-client)
    # [
    #     "https://github.com/filliptm/ComfyUI_Fill-Nodes.git",
    #     "main",
    # ],
]


def _get_git_repo_name_from_url(url: str) -> str:
    return url.rstrip("/").rsplit("/", maxsplit=1)[-1].removesuffix(".git")


def clone_and_link_custom_node(
    url: str,
    commit_hash: str,
    global_dir: Path,
    local_dir: Path,
    global_repo_name: str,
    local_repo_name: str,
) -> Path:
    global_repo_path = clone_repository(
        url, commit_hash=commit_hash, target_dir=global_dir, repo_name=global_repo_name
    )
    # note: local dir uses a different name since some of the custom nodes
    # expect the repo to have the same name as the git repo
    local_repo_path = local_dir / local_repo_name

    try:
        local_repo_path.symlink_to(global_repo_path, target_is_directory=True)
    except FileExistsError:
        # check if it points to the same file
        if local_repo_path.resolve() == global_repo_path.resolve():
            return local_repo_path
        else:
            try:
                os.unlink(local_repo_path)
            except FileNotFoundError:
                pass

            try:
                local_repo_path.symlink_to(global_repo_path, target_is_directory=True)
            except FileExistsError:
                pass

    return local_repo_path


def clone_custom_nodes(repo_path: str | Path, custom_nodes: list[list[str]]):
    comfy_custom_node_dir = Path(repo_path) / "custom_nodes"
    global_comfy_custom_node_dir = Path(FAL_REPOSITORY_DIR) / "global_custom_nodes"
    cloned_nodes = {}

    with ThreadPoolExecutor(max_workers=12) as executor:
        for node in custom_nodes:
            url, commit_hash = node
            if "{GITHUB_TOKEN}" in url:
                if "GITHUB_TOKEN" not in os.environ:
                    raise ValueError(
                        f"GITHUB_TOKEN environment variable is not set, but reqired for {url}"
                    )
                url = url.format(GITHUB_TOKEN=os.getenv("GITHUB_TOKEN"))

            local_repo_name = _get_git_repo_name_from_url(url)
            global_repo_name = f"{local_repo_name}-{commit_hash[:8]}"

            future = executor.submit(
                clone_and_link_custom_node,
                url,
                commit_hash=commit_hash,
                local_dir=comfy_custom_node_dir,
                global_dir=global_comfy_custom_node_dir,
                local_repo_name=local_repo_name,
                global_repo_name=global_repo_name,
            )
            cloned_nodes[url] = future

        for url, future in cloned_nodes.items():
            try:
                cloned_nodes[url] = future.result()
            except Exception as e:
                print(f"Task failed: {e}")
                raise e

        impact_pack_dir = cloned_nodes.get(COMFYUI_IMPACT_PACK_URL)
        if (
            impact_pack_dir
            and get_impact_subpack_location(str(impact_pack_dir)) == "impact"
        ):
            # Impact Subpack requires a different target dir
            impact_subpack_url = (
                "https://github.com/ltdrdata/ComfyUI-Impact-Subpack.git"
            )
            impact_subpack_commit_hash = "bd4dedc4ea514dda9cfa333a0f42e296e32b8c10"
            impact_subpack_global_repo_name = (
                f"{_get_git_repo_name_from_url(impact_subpack_url)}-"
                f"{impact_subpack_commit_hash[:8]}"
            )

            clone_and_link_custom_node(
                impact_subpack_url,
                commit_hash=impact_subpack_commit_hash,
                local_dir=impact_pack_dir,
                global_dir=global_comfy_custom_node_dir,
                local_repo_name="impact_subpack",
                global_repo_name=impact_subpack_global_repo_name,
            )

    return cloned_nodes


def cleanup_unused_custom_nodes(repo_path: str | Path, cloned_nodes: dict[str, Path]):
    comfy_custom_node_dir = Path(repo_path) / "custom_nodes"
    custom_node_names = list(map(lambda x: x.name, cloned_nodes.values()))

    for node in comfy_custom_node_dir.iterdir():
        is_py_cache = node.name.endswith("__pycache__")
        is_for_deletion = ".for_deletion" in node.name
        node_exists = node.name in custom_node_names

        if (node.is_dir() or node.is_symlink()) and not is_py_cache and not node_exists:
            try:
                print(f"Removing {node}")

                # Handle symlinks directly with unlink
                if node.is_symlink():
                    os.unlink(node)
                else:
                    # atomic rename to random name and delete for real directories
                    temp_name = f"{node}.{uuid.uuid4().hex}.for_deletion"
                    os.rename(node, temp_name)
                    shutil.rmtree(temp_name)

            except Exception as e:
                print(f"Failed to remove {node}: {e}")

        # Also clean up any leftover .for_deletion directories
        elif is_for_deletion:
            try:
                print(f"Cleaning up leftover: {node}")
                if node.is_symlink():
                    os.unlink(node)
                else:
                    shutil.rmtree(node)
            except Exception as e:
                print(f"Failed to cleanup {node}: {e}")


def filter_custom_node_models(
    custom_nodes: list[list[str]], custom_node_models: dict[str, list[list[str]]]
) -> dict[str, list[list[str]]]:
    """Filter the custom node models to only include the ones that are needed"""
    custom_node_names = set(
        _get_git_repo_name_from_url(node[0]) for node in custom_nodes
    )

    filtered_custom_node_models = {}
    for custom_node_dir, models in custom_node_models.items():
        custom_node_name = custom_node_dir.split("/")[0]
        if custom_node_name in custom_node_names:
            filtered_custom_node_models[custom_node_dir] = models

    return filtered_custom_node_models
