import fal
from fastapi.responses import StreamingResponse

from comfy_server.comfy import (
    COMFY_COMMIT_HASH,
    CUSTOM_NODE_MODELS,
    CUSTOM_NODES,
    MODELS,
    REPO_NAME,
)
from comfy_server.env import get_requirements
from comfy_server.models import Input, Output
from comfy_server.utils import build_fal_workflow

WAN_VACE_14B_URL = (
    "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/"
    "split_files/diffusion_models/wan2.1_vace_14B_fp16.safetensors"
)
WAN_VACE_14B_REPAIR_VERSION = "wan_vace_14b_redownload_v1"


def _patch_fal_exceptions():
    """Ensure both fal.exceptions.gpu and fal.exceptions._cuda exist."""
    import sys
    import types

    try:
        from fal.exceptions.gpu import CUDAOutOfMemoryException, _is_cuda_oom_exception
    except ImportError:
        try:
            from fal.exceptions._cuda import CUDAOutOfMemoryException, _is_cuda_oom_exception
        except ImportError:
            CUDAOutOfMemoryException = RuntimeError
            _is_cuda_oom_exception = lambda e: False

    for mod_name in ("fal.exceptions.gpu", "fal.exceptions._cuda"):
        if mod_name not in sys.modules:
            m = types.ModuleType(mod_name)
            m.CUDAOutOfMemoryException = CUDAOutOfMemoryException
            m._is_cuda_oom_exception = _is_cuda_oom_exception
            sys.modules[mod_name] = m


def _repair_wan_vace_14b_cache_once():
    import glob
    import shutil
    from pathlib import Path

    from fal.toolkit import FAL_PERSISTENT_DIR
    from fal.toolkit.utils.download_utils import FAL_MODEL_WEIGHTS_DIR, _hash_url

    repairs_dir = Path(FAL_PERSISTENT_DIR) / ".deployment_repairs"
    marker_path = repairs_dir / f"{WAN_VACE_14B_REPAIR_VERSION}.done"
    if marker_path.exists():
        return

    cached_weights_dir = Path(FAL_MODEL_WEIGHTS_DIR) / _hash_url(WAN_VACE_14B_URL)
    if cached_weights_dir.exists():
        print(f"Removing cached WAN VACE 14B weights: {cached_weights_dir}")
        shutil.rmtree(cached_weights_dir, ignore_errors=True)

    for repo_glob in (
        "/data/.fal/repos/*/models/diffusion_models/wan2.1_vace_14B_fp16.safetensors",
        "/data/.fal/repos/*/models/unet/wan2.1_vace_14B_fp16.safetensors",
    ):
        for model_symlink in glob.glob(repo_glob):
            model_path = Path(model_symlink)
            try:
                if model_path.is_symlink() or model_path.is_file():
                    print(f"Removing stale WAN VACE 14B link: {model_path}")
                    model_path.unlink(missing_ok=True)
            except Exception as exc:
                print(f"Failed to remove stale WAN VACE 14B link {model_path}: {exc}")

    repairs_dir.mkdir(parents=True, exist_ok=True)
    marker_path.write_text("done\n")


class Comfy(
    fal.App,
    min_concurrency=0,
    max_concurrency=2,
):  # type: ignore
    machine_type = "GPU-H100"
    requirements = get_requirements()
    host_kwargs = {
        "_scheduler": "nomad",
        "_scheduler_options": {
            "storage_region": "us-east",
        },
        "resolver": "uv",
        "keep_alive": 3600,
        "startup_timeout": 2700,  # 45 minutes
    }

    async def setup(self):
        _patch_fal_exceptions()
        _repair_wan_vace_14b_cache_once()

        from serverless_comfy import ComfyRuntime

        self.runtime = ComfyRuntime()

        self.runtime.setup_comfy(
            repo_name=REPO_NAME,
            commit_hash=COMFY_COMMIT_HASH,
            custom_nodes=CUSTOM_NODES,
            models=MODELS,
            custom_node_models=CUSTOM_NODE_MODELS,
        )

        # Fix TrentNodes MiniCPM: restore original file if corrupted by
        # old patch, then apply correct device mismatch fix
        import glob
        from urllib.request import urlopen as _urlopen
        TRENT_COMMIT = "05a4dc0d4a19224623256b0c9825b5fc975c187a"
        TRENT_RAW = f"https://raw.githubusercontent.com/TrentHunter82/TrentNodes/{TRENT_COMMIT}/utils/minicpm_wrapper.py"
        for wrapper in glob.glob("/data/.fal/repos/*/custom_nodes/TrentNodes/utils/minicpm_wrapper.py"):
            with open(wrapper, "r") as f:
                src = f.read()
            # If file has syntax error from old patch, restore original
            if "expected" not in src:
                try:
                    compile(src, wrapper, "exec")
                except SyntaxError:
                    print(f"Restoring corrupted {wrapper} from GitHub...")
                    src = _urlopen(TRENT_RAW).read().decode()
                    with open(wrapper, "w") as f:
                        f.write(src)
            # Apply device fix if not already patched
            if "_device_patched" not in src:
                lines = src.splitlines(True)
                for i, line in enumerate(lines):
                    if i > 400 and "_touch_model()" in line:
                        indent = line[:len(line) - len(line.lstrip())]
                        patch = (
                            f"\n"
                            f"{indent}# Fix device mismatch with device_map='auto' (int4)\n"
                            f"{indent}_model_cls = type(model)\n"
                            f"{indent}if not getattr(_model_cls, '_device_patched', False):\n"
                            f"{indent}    _model_cls.device = property(lambda self: self.llm.model.embed_tokens.weight.device)\n"
                            f"{indent}    _model_cls._device_patched = True\n"
                        )
                        lines[i] = line + patch
                        with open(wrapper, "w") as f:
                            f.writelines(lines)
                        print(f"Patched MiniCPM device fix: {wrapper}")
                        break

        self.runtime.start_server(startup_retries=200)

    @fal.endpoint("/")
    async def execute(self, input: Input) -> Output:
        comfy_input = self.build_comfy_input(input)

        output: Output = await self.runtime.execute_prompt(comfy_input)
        return output

    def build_comfy_input(self, input: Input):
        from serverless_comfy.models import ComfyInput

        fal_workflow = build_fal_workflow(input)
        comfy_workflow = ComfyInput.parse_obj(fal_workflow)

        return comfy_workflow

    @fal.endpoint("/stream")
    async def execute_stream(self, input: Input) -> StreamingResponse:
        comfy_input = self.build_comfy_input(input)
        return await self.runtime.execute_prompt_sse(comfy_input)

    @fal.endpoint("/stream/cancel")
    async def inference_stream_cancel(self):
        return await self.runtime.interrupt()

    @fal.endpoint("/object_info")
    async def node_class_mapping(self):
        import nodes as nodes_module
        out = {}
        for node_class_name, obj_class in nodes_module.NODE_CLASS_MAPPINGS.items():
            try:
                info = {}
                info['input'] = obj_class.INPUT_TYPES()
                info['input_order'] = {key: list(value.keys()) for (key, value) in obj_class.INPUT_TYPES().items()}
                info['output'] = obj_class.RETURN_TYPES
                info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES)
                info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output']
                info['name'] = node_class_name
                info['display_name'] = nodes_module.NODE_DISPLAY_NAME_MAPPINGS.get(node_class_name, node_class_name)
                info['description'] = getattr(obj_class, 'DESCRIPTION', '')
                info['python_module'] = getattr(obj_class, 'RELATIVE_PYTHON_MODULE', None) or 'nodes'
                info['category'] = getattr(obj_class, 'CATEGORY', None) or 'sd'
                info['output_node'] = getattr(obj_class, 'OUTPUT_NODE', False)
                if hasattr(obj_class, 'OUTPUT_TOOLTIPS'):
                    info['output_tooltips'] = obj_class.OUTPUT_TOOLTIPS
                if getattr(obj_class, 'DEPRECATED', False):
                    info['deprecated'] = True
                if getattr(obj_class, 'EXPERIMENTAL', False):
                    info['experimental'] = True
                if hasattr(obj_class, 'API_NODE'):
                    info['api_node'] = obj_class.API_NODE
                out[node_class_name] = info
            except Exception:
                pass
        return out

    @fal.endpoint("/extensions")
    async def get_extensions(self):
        """Return list of extension JS paths from ComfyUI."""
        import httpx
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.get("http://127.0.0.1:8188/extensions")
            return resp.json()

    @fal.endpoint("/extension_files")
    async def get_extension_files(self):
        """Return all extension JS file contents read directly from disk."""
        import os
        import glob
        result = {}

        # ComfyUI web extensions are served from custom_nodes/*/web/ or custom_nodes/*/js/
        repo_path = str(self.runtime.repo_path)
        web_root = os.path.join(repo_path, "web", "extensions")

        # Also check custom_nodes for their web directories
        custom_nodes_dir = os.path.join(repo_path, "custom_nodes")

        # Read core extensions
        if os.path.isdir(web_root):
            for root, dirs, files in os.walk(web_root):
                for f in files:
                    if f.endswith(".js"):
                        full_path = os.path.join(root, f)
                        rel_path = os.path.relpath(full_path, os.path.join(repo_path, "web"))
                        try:
                            with open(full_path, "r", encoding="utf-8", errors="ignore") as fh:
                                result["/" + rel_path] = fh.read()
                        except Exception:
                            pass

        # Read custom node extensions
        if os.path.isdir(custom_nodes_dir):
            for node_dir in os.listdir(custom_nodes_dir):
                node_path = os.path.join(custom_nodes_dir, node_dir)
                if not os.path.isdir(node_path):
                    continue
                # Check standard dirs + any dir containing .js files
                web_dirs = ["web", "js", "web_version/v1", "web_version/v2", "web_version/v3"]
                # Also scan __init__.py for WEB_DIRECTORY string literals
                try:
                    init_path = os.path.join(node_path, "__init__.py")
                    if os.path.isfile(init_path):
                        with open(init_path, "r", errors="ignore") as fh:
                            for line in fh:
                                if "WEB_DIRECTORY" in line and "=" in line:
                                    val = line.split("=", 1)[1].strip().strip("'\"").strip("./")
                                    if val and val not in web_dirs and "/" not in val.replace("web_version/", ""):
                                        web_dirs.append(val)
                except Exception:
                    pass
                for sub in web_dirs:
                    ext_dir = os.path.join(node_path, sub)
                    if not os.path.isdir(ext_dir):
                        continue
                    for root, dirs, files in os.walk(ext_dir):
                        for f in files:
                            if f.endswith(".js"):
                                full_path = os.path.join(root, f)
                                rel_path = os.path.relpath(full_path, ext_dir)
                                key = f"/extensions/{node_dir}/{rel_path}"
                                try:
                                    with open(full_path, "r", encoding="utf-8", errors="ignore") as fh:
                                        result[key] = fh.read()
                                except Exception:
                                    pass

        return result

    @fal.endpoint("/cancel")
    async def inference_cancel(self):
        return await self.runtime.interrupt()
