import logging
import os
import threading

import torch
import torch.nn as nn

_LOG = logging.getLogger(__name__)

_DEFAULT_FORCE_PARAMETER_STATIC_SHAPES = getattr(
    torch._dynamo.config, "force_parameter_static_shapes", None
)
_DEFAULT_GRAPH_PARTITION = getattr(torch._inductor.config, "graph_partition", None)

try:
    from comfy_api.torch_helpers import set_torch_compile_wrapper as _set_torch_compile_wrapper
except Exception:
    _set_torch_compile_wrapper = None

try:
    from comfy_extras.nodes_torch_compile import skip_torch_compile_dict as _skip_torch_compile_dict
except Exception:
    _skip_torch_compile_dict = None


class _LazyCompiled(nn.Module):
    """
    Wraps any nn.Module (e.g. a UNet or a single transformer block)
    and replaces its forward pass with torch.compile **lazily** on the
    first call.  All other attributes / helpers are proxied so that the
    module behaves like the original one.
    """

    def __init__(self, module: nn.Module, **compile_kw):
        super().__init__()
        self._orig = module
        self._compiled = None
        self._compile_kw = compile_kw
        self._compile_lock = threading.Lock()

    def __getattr__(self, name):
        if name in {"_orig", "_compiled", "_compile_kw", "_compile_lock"}:
            return super().__getattr__(name)
        return getattr(self._orig, name)

    def modules(self):
        return self._orig.modules()

    def children(self):
        return self._orig.children()

    def named_modules(self, *a, **k):
        return self._orig.named_modules(*a, **k)

    def state_dict(self, *a, **k):
        return self._orig.state_dict(*a, **k)

    @property
    def dtype(self):
        return self._orig.dtype

    @property
    def device(self):
        return self._orig.device

    def train(self, mode: bool = True):
        self._orig.train(mode)
        return super().train(mode)

    def eval(self):
        return self.train(False)

    def to(self, *args, **kwargs):
        self._orig.to(*args, **kwargs)
        return super().to(*args, **kwargs)

    def forward(self, *args, **kwargs):
        if self._compiled is None:
            with self._compile_lock:
                if self._compiled is None:
                    self._compiled = torch.compile(self._orig, **self._compile_kw)
        return self._compiled(*args, **kwargs)


class TorchCompileModel_LoRASafe:
    """LoRA-safe torch.compile that also works with Flux block layouts."""

    _GENERIC_BLOCK_LAYER_TYPES = (
        "layers",
        "transformer_blocks",
        "blocks",
        "visual_transformer_blocks",
        "text_transformer_blocks",
    )

    _FLUX_DOUBLE_PREFERRED_SUFFIXES = (
        "img_attn.qkv",
        "img_attn.proj",
        "txt_attn.qkv",
        "txt_attn.proj",
    )
    _FLUX_DOUBLE_CUDAGRAPH_SAFE_SUFFIXES = (
        "img_attn.qkv",
        "txt_attn.qkv",
    )

    _FLUX_SINGLE_PREFERRED_SUFFIXES = ("linear1", "linear2")

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "backend": (["inductor", "cudagraphs", "nvfuser"],),
                "mode": (["default", "reduce-overhead", "max-autotune"],),
                "fullgraph": ("BOOLEAN", {"default": False}),
                "dynamic": ("BOOLEAN", {"default": False}),
                "disable_cudagraphs": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "True -> pass torch.compile options {'triton.cudagraphs': False} "
                        "for inductor backend. Useful to avoid cudaMallocAsync/cudagraph issues "
                        "in some environments.",
                    },
                ),
                "compile_transformer_only": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "True -> compile discovered transformer targets only. "
                        "Flux models compile leaf submodules inside double_blocks/single_blocks; "
                        "SD-style models often compile transformer_blocks/blocks.",
                    },
                ),
                "debug_torch_logs": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "Enable PyTorch compile/cudagraph debug logs via "
                        "torch._logging.set_logs (ignored if TORCH_LOGS env var is already set).",
                    },
                ),
                "skip_dynamic_cudagraphs": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "For cudagraph-capable runs, skip cudagraphing functions with "
                        "dynamic-shape inputs. Recommended for Flux.",
                    },
                ),
                "allow_flux_inferred_targets": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "False (recommended) -> Flux compile targets are strict and limited "
                        "to qkv/proj/linear1/linear2 leaves only. "
                        "True -> also allow inferred Flux target names for compatibility.",
                    },
                ),
                "fallback_to_full_model_if_no_targets": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": "If transformer-only target discovery finds nothing, "
                        "fall back to compiling diffusion_model. "
                        "False keeps execution uncompiled instead of widening target scope.",
                    },
                ),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "patch"
    CATEGORY = "model/optimisation 🛠️"
    EXPERIMENTAL = True

    @staticmethod
    def _check_backend(backend: str):
        if backend == "nvfuser" and not torch.cuda.is_available():
            raise ValueError("nvfuser backend requires a CUDA GPU. " "Select 'inductor' instead.")
        if backend == "cudagraphs":
            if not torch.cuda.is_available():
                raise ValueError("cudagraphs backend needs a CUDA GPU.")
            cap = torch.cuda.get_device_capability()
            if cap[0] < 7:
                raise ValueError(
                    "cudagraphs works reliably on GPUs with "
                    "compute capability 7.0 or higher "
                    f"(yours is {cap[0]}.{cap[1]})."
                )

    @staticmethod
    def _is_leaf_compile_module(module: nn.Module) -> bool:
        if not isinstance(module, nn.Module):
            return False
        if any(True for _ in module.children()):
            return False
        return any(True for _ in module.parameters(recurse=False))

    @classmethod
    def _discover_flux_all_leaf_keys(cls, diffusion_model: nn.Module):
        keys = []
        for name, module in diffusion_model.named_modules():
            if not name or not cls._is_leaf_compile_module(module):
                continue
            keys.append(f"diffusion_model.{name}")
        return list(dict.fromkeys(keys))

    @staticmethod
    def _is_flux_model(diffusion_model: nn.Module) -> bool:
        return hasattr(diffusion_model, "double_blocks") and hasattr(
            diffusion_model, "single_blocks"
        )

    @classmethod
    def _try_preferred_suffixes(cls, block: nn.Module, suffixes):
        targets = []
        for suffix in suffixes:
            obj = block
            ok = True
            for part in suffix.split("."):
                if not hasattr(obj, part):
                    ok = False
                    break
                obj = getattr(obj, part)
            if ok and cls._is_leaf_compile_module(obj):
                targets.append(suffix)
        return targets

    @classmethod
    def _discover_flux_double_block_suffixes(
        cls,
        block: nn.Module,
        allow_inferred_targets: bool = False,
        cudagraph_capable: bool = False,
    ):
        preferred_suffixes = (
            cls._FLUX_DOUBLE_CUDAGRAPH_SAFE_SUFFIXES
            if cudagraph_capable
            else cls._FLUX_DOUBLE_PREFERRED_SUFFIXES
        )
        preferred = cls._try_preferred_suffixes(block, preferred_suffixes)
        if preferred or not allow_inferred_targets:
            return preferred

        inferred = []
        for name, module in block.named_modules():
            if not name or not cls._is_leaf_compile_module(module):
                continue
            lname = name.lower()
            parts = lname.split(".")
            tail = parts[-1]
            if (
                (("img_attn" in lname or "txt_attn" in lname) and tail in {"qkv", "proj"})
                or ("img_mlp" in lname)
                or ("txt_mlp" in lname)
            ):
                inferred.append(name)
        return list(dict.fromkeys(inferred))

    @classmethod
    def _discover_flux_single_block_suffixes(
        cls,
        block: nn.Module,
        allow_inferred_targets: bool = False,
        cudagraph_capable: bool = False,
    ):
        if cudagraph_capable:
            return []

        preferred = cls._try_preferred_suffixes(block, cls._FLUX_SINGLE_PREFERRED_SUFFIXES)
        if preferred or not allow_inferred_targets:
            return preferred

        inferred = []
        for name, module in block.named_modules():
            if not name or not cls._is_leaf_compile_module(module):
                continue
            lname = name.lower()
            if lname.endswith("linear1") or lname.endswith("linear2"):
                inferred.append(name)
        return list(dict.fromkeys(inferred))

    @classmethod
    def _discover_compile_keys(
        cls,
        diffusion_model,
        allow_flux_inferred_targets: bool = False,
        cudagraph_capable: bool = False,
    ):
        keys = []

        # Flux special-case: compile only leaf tensor-heavy modules, not full block forward.
        if cls._is_flux_model(diffusion_model):
            double_blocks = getattr(diffusion_model, "double_blocks")
            single_blocks = getattr(diffusion_model, "single_blocks")

            if isinstance(double_blocks, (nn.ModuleList, list, tuple)):
                for i, block in enumerate(double_blocks):
                    for suffix in cls._discover_flux_double_block_suffixes(
                        block,
                        allow_inferred_targets=allow_flux_inferred_targets,
                        cudagraph_capable=cudagraph_capable,
                    ):
                        keys.append(f"diffusion_model.double_blocks.{i}.{suffix}")

            if isinstance(single_blocks, (nn.ModuleList, list, tuple)):
                for i, block in enumerate(single_blocks):
                    for suffix in cls._discover_flux_single_block_suffixes(
                        block,
                        allow_inferred_targets=allow_flux_inferred_targets,
                        cudagraph_capable=cudagraph_capable,
                    ):
                        keys.append(f"diffusion_model.single_blocks.{i}.{suffix}")

            if keys:
                return list(dict.fromkeys(keys))

        for layer_name in cls._GENERIC_BLOCK_LAYER_TYPES:
            if hasattr(diffusion_model, layer_name):
                blocks = getattr(diffusion_model, layer_name)
                if not isinstance(blocks, (nn.ModuleList, list, tuple)):
                    continue
                for i in range(len(blocks)):
                    keys.append(f"diffusion_model.{layer_name}.{i}")
        return list(dict.fromkeys(keys))

    @staticmethod
    def _build_compile_kwargs(backend, mode, fullgraph, dynamic, disable_cudagraphs):
        compile_kw = {
            "backend": backend,
            "fullgraph": fullgraph,
            "dynamic": dynamic,
        }

        options = {}
        if _skip_torch_compile_dict is not None:
            options["guard_filter_fn"] = _skip_torch_compile_dict
        if disable_cudagraphs and backend == "inductor":
            options["triton.cudagraphs"] = False

        # torch.compile rejects calls that specify both `mode` and `options`.
        # When we need extra backend options as well as a preset mode, expand
        # the mode preset into concrete backend options and merge them.
        if mode != "default":
            if options:
                list_mode_options = getattr(torch._inductor, "list_mode_options", None)
                if list_mode_options is None:
                    raise RuntimeError(
                        "Selected a non-default torch.compile mode together with backend options, "
                        "but this PyTorch build does not expose torch._inductor.list_mode_options() "
                        "to merge them safely."
                    )
                merged_options = dict(list_mode_options(mode, dynamic) or {})
                merged_options.update(options)
                compile_kw["options"] = merged_options
            else:
                compile_kw["mode"] = mode
        elif options:
            compile_kw["options"] = options

        return compile_kw

    @staticmethod
    def _configure_torch_logs(debug_torch_logs: bool):
        if not debug_torch_logs:
            return
        # TORCH_LOGS takes precedence over torch._logging.set_logs(...).
        if os.environ.get("TORCH_LOGS"):
            _LOG.warning(
                "TORCH_LOGS is already set; using environment logging instead of "
                "torch._logging.set_logs."
            )
            return
        try:
            torch._logging.set_logs(
                dynamo=logging.INFO,
                inductor=logging.INFO,
                cudagraphs=True,
                cudagraph_static_inputs=True,
                perf_hints=True,
                graph_breaks=True,
                guards=True,
                recompiles=True,
                recompiles_verbose=True,
            )
        except Exception as e:
            _LOG.warning(f"Failed to enable torch logging: {e}")

    def patch(
        self,
        model,
        backend,
        mode,
        fullgraph,
        dynamic,
        disable_cudagraphs,
        compile_transformer_only,
        debug_torch_logs=False,
        skip_dynamic_cudagraphs=True,
        allow_flux_inferred_targets=False,
        fallback_to_full_model_if_no_targets=False,
    ):
        self._check_backend(backend)

        self._configure_torch_logs(debug_torch_logs)

        cudagraph_capable = backend == "cudagraphs" or (
            backend == "inductor" and not disable_cudagraphs
        )

        if _DEFAULT_FORCE_PARAMETER_STATIC_SHAPES is not None:
            torch._dynamo.config.force_parameter_static_shapes = (
                False if cudagraph_capable else _DEFAULT_FORCE_PARAMETER_STATIC_SHAPES
            )
        if _DEFAULT_GRAPH_PARTITION is not None:
            torch._inductor.config.graph_partition = (
                True if cudagraph_capable else _DEFAULT_GRAPH_PARTITION
            )

        # This flag is process-global, so set it on every call to match the
        # current node input instead of leaking state from earlier runs.
        if backend in ("inductor", "cudagraphs"):
            torch._inductor.config.triton.cudagraph_skip_dynamic_graphs = bool(
                skip_dynamic_cudagraphs
            )

        try:
            m = model.clone(disable_dynamic=True)
        except TypeError:
            m = model.clone()
        dm = m.get_model_object("diffusion_model")

        compile_kw = self._build_compile_kwargs(
            backend, mode, fullgraph, dynamic, disable_cudagraphs
        )

        is_flux_model = self._is_flux_model(dm)

        if compile_transformer_only:
            compile_keys = self._discover_compile_keys(
                dm,
                allow_flux_inferred_targets=(
                    allow_flux_inferred_targets and not cudagraph_capable
                ),
                cudagraph_capable=cudagraph_capable,
            )
            if not compile_keys:
                if fallback_to_full_model_if_no_targets and not (
                    cudagraph_capable and is_flux_model
                ):
                    _LOG.warning(
                        "TorchCompileModel_LoRASafe: no known transformer compile targets found; "
                        "falling back to whole diffusion model compile."
                    )
                    compile_keys = ["diffusion_model"]
                else:
                    if cudagraph_capable and is_flux_model:
                        _LOG.warning(
                            "TorchCompileModel_LoRASafe: no cudagraph-safe Flux transformer "
                            "compile targets found; skipping compile instead of widening target scope."
                        )
                    else:
                        _LOG.warning(
                            "TorchCompileModel_LoRASafe: no known transformer compile targets found; "
                            "skipping compile instead of falling back to whole model."
                        )
                    return (m,)
        elif cudagraph_capable and is_flux_model:
            compile_keys = self._discover_flux_all_leaf_keys(dm)
            if not compile_keys:
                _LOG.warning(
                    "TorchCompileModel_LoRASafe: Flux full-leaf cudagraph compile requested, "
                    "but no leaf targets were found; skipping compile instead of compiling the "
                    "whole WrapperExecutor-driven diffusion_model path."
                )
                return (m,)
            _LOG.warning(
                "TorchCompileModel_LoRASafe: Flux whole-model cudagraph compile requested; "
                "using wide leaf-module compile instead to avoid WrapperExecutor / "
                "torch.compiler.disable replay failures."
            )
        else:
            compile_keys = ["diffusion_model"]

        if cudagraph_capable and is_flux_model:
            label = (
                "Flux cudagraph transformer targets"
                if compile_transformer_only
                else "Flux cudagraph wide-leaf targets"
            )
            _LOG.warning("%s: %s", label, compile_keys[:16])
            _LOG.warning("%s count: %d", label, len(compile_keys))

        if _set_torch_compile_wrapper is not None:
            try:
                _set_torch_compile_wrapper(model=m, keys=compile_keys, **compile_kw)
                return (m,)
            except Exception as e:
                raise RuntimeError(
                    f"Failed to apply torch.compile wrapper via comfy_api.torch_helpers: {e}"
                ) from e

        for key in compile_keys:
            if key == "diffusion_model":
                target = dm
            else:
                obj = dm
                for part in key.split(".")[1:]:
                    obj = obj[int(part)] if part.isdigit() else getattr(obj, part)
                target = obj
            m.add_object_patch(key, _LazyCompiled(target, **compile_kw))

        return (m,)


NODE_CLASS_MAPPINGS = {
    "TorchCompileModel_LoRASafe": TorchCompileModel_LoRASafe,
}