"""
FLUXNATION — Fused CUDA kernel for FLUX.1 SingleStreamBlock.

Architecture: PyTorch C++ Extension via torch.ops (no ctypes)
  - torch.ops.flux_ext.ssb_prepare: modulation + L1 GEMM + split QKV + RoPE
  - torch.ops.flux_ext.ssb_finish:  L2 GEMM + fused gate+residual+bias

Build: run build_ext.bat (or: py -3.10 setup_ext.py install)
"""

class FLUXNATION:
    """Main FLUXNATION kernel node. Pass-through MODEL — kernel is already active."""
    CATEGORY = "FLUXNATION"
    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "apply"

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"model": ("MODEL",)}}

    def apply(self, model):
        from .FLUXPATCH import _spike_controller
        spike = "SPIKE ON" if _spike_controller.enabled else "SPIKE OFF"
        print(f"[FLUXNATION] FP8 fused CUDA kernel ACTIVE | {spike}")
        return (model,)


class FLUXATTENHUT_Node:
    """Configure spike attention parameters."""
    CATEGORY = "FLUXNATION"
    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "configure"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model":       ("MODEL",),
                "switch_pct":  ("FLOAT", {"default": 0.50, "min": 0.1, "max": 0.9, "step": 0.05,
                                          "tooltip": "Fraction of steps before spike activates (0.50 = step 10/20)"}),
                "cap_pct":     ("FLOAT", {"default": 0.45, "min": 0.1, "max": 0.9, "step": 0.05,
                                          "tooltip": "Fraction of attention blocks to keep (0.45 = 45%)"}),
                "tau":         ("FLOAT", {"default": 0.05, "min": 0.01, "max": 0.5, "step": 0.01,
                                          "tooltip": "Block scoring threshold — lower = more blocks kept"}),
            }
        }

    def configure(self, model, switch_pct, cap_pct, tau):
        from .FLUXPATCH import _spike_controller
        _spike_controller.switch_pct = switch_pct
        _spike_controller.event_cap_pct = cap_pct
        _spike_controller.tau = tau
        print(f"[FLUXATTENHUT] switch={switch_pct*100:.0f}% cap={cap_pct*100:.0f}% tau={tau}")
        return (model,)


class FLUXSPIKEY_Node:
    """Enable or disable spike attention + step caching."""
    CATEGORY = "FLUXNATION"
    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "configure"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model":       ("MODEL",),
                "spike":       ("BOOLEAN", {"default": True,  "tooltip": "Enable block-sparse spike attention"}),
                "step_cache":  ("BOOLEAN", {"default": True,  "tooltip": "Cache attention output on alternating spike steps"}),
            }
        }

    def configure(self, model, spike, step_cache):
        from .FLUXPATCH import _spike_controller
        _spike_controller.enabled = spike
        _spike_controller.cache_enabled = step_cache
        print(f"[FLUXSPIKEY] spike={'ON' if spike else 'OFF'} cache={'ON' if step_cache else 'OFF'}")
        return (model,)


class ULT7RAFLUX:
    """Master status node — shows kernel state, spike config, and timing info."""
    CATEGORY = "FLUXNATION"
    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("model",)
    FUNCTION = "status"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model":      ("MODEL",),
                "spike":      ("BOOLEAN", {"default": True}),
                "step_cache": ("BOOLEAN", {"default": True}),
                "switch_pct": ("FLOAT",   {"default": 0.50, "min": 0.1, "max": 0.9, "step": 0.05}),
                "cap_pct":    ("FLOAT",   {"default": 0.45, "min": 0.1, "max": 0.9, "step": 0.05}),
                "tau":        ("FLOAT",   {"default": 0.05, "min": 0.01, "max": 0.5, "step": 0.01}),
            }
        }

    def status(self, model, spike, step_cache, switch_pct, cap_pct, tau):
        from .FLUXPATCH import _spike_controller
        _spike_controller.enabled       = spike
        _spike_controller.cache_enabled = step_cache
        _spike_controller.switch_pct    = switch_pct
        _spike_controller.event_cap_pct = cap_pct
        _spike_controller.tau           = tau
        print(
            f"[ULT7RAFLUX] FP8 fused kernel | "
            f"spike={'ON' if spike else 'OFF'} | "
            f"cache={'ON' if step_cache else 'OFF'} | "
            f"switch={switch_pct*100:.0f}% cap={cap_pct*100:.0f}% tau={tau}"
        )
        return (model,)


NODE_CLASS_MAPPINGS = {
    "FLUXNATION":   FLUXNATION,
    "FLUXATTENHUT": FLUXATTENHUT_Node,
    "FLUXSPIKEY":   FLUXSPIKEY_Node,
    "ULT7RAFLUX":   ULT7RAFLUX,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "FLUXNATION":   "FLUXNATION",
    "FLUXATTENHUT": "FLUXATTENHUT",
    "FLUXSPIKEY":   "FLUXSPIKEY",
    "ULT7RAFLUX":   "ULT7RAFLUX",
}

_EXT_LOADED = False


def _patch_safetensor_loading():
    """Fix Windows mmap access violation in safetensors loading.

    PyTorch 2.5.x + safetensors safe_open() on Windows triggers access violations
    when memory-mapping large files (>10GB). This patches ComfyUI's load_torch_file
    to use safetensors.torch.load_file() which loads the full file without mmap.
    """
    import sys
    if sys.platform != "win32":
        return

    try:
        import comfy.utils as comfy_utils
        import safetensors.torch
        import torch

        _original_load = comfy_utils.load_torch_file

        def _read_safetensor_streaming(path, target_device='cpu'):
            """Stream-load a safetensor file without mmap or full-file buffering.

            Reads the header to get tensor metadata, then reads each tensor
            individually from disk into a pre-allocated torch tensor.
            """
            import struct, json, numpy as np

            dtype_map = {
                'F16': (torch.float16, 2),
                'BF16': (torch.bfloat16, 2),
                'F32': (torch.float32, 4),
                'F64': (torch.float64, 8),
                'I8': (torch.int8, 1),
                'I16': (torch.int16, 2),
                'I32': (torch.int32, 4),
                'I64': (torch.int64, 8),
                'U8': (torch.uint8, 1),
                'BOOL': (torch.bool, 1),
                'F8_E4M3': (getattr(torch, 'float8_e4m3fn', torch.uint8), 1),
                'F8_E5M2': (getattr(torch, 'float8_e5m2', torch.uint8), 1),
            }
            np_dtype_map = {
                'F16': np.float16, 'BF16': np.void, 'F32': np.float32,
                'F64': np.float64, 'I8': np.int8, 'I16': np.int16,
                'I32': np.int32, 'I64': np.int64, 'U8': np.uint8,
                'BOOL': np.bool_,
            }

            with open(path, 'rb') as f:
                header_size = struct.unpack('<Q', f.read(8))[0]
                header_bytes = f.read(header_size)
                data_offset = 8 + header_size

                header = json.loads(header_bytes)
                metadata = header.pop('__metadata__', {})

                sd = {}
                for name, info in header.items():
                    dtype_str = info['dtype']
                    shape = info['shape']
                    offsets = info['data_offsets']
                    start, end = offsets[0], offsets[1]
                    nbytes = end - start

                    torch_dtype, elem_size = dtype_map.get(dtype_str, (torch.float32, 4))

                    f.seek(data_offset + start)
                    raw = f.read(nbytes)

                    if dtype_str == 'BF16':
                        # numpy doesn't support bfloat16 — read as uint16 then view
                        arr = np.frombuffer(raw, dtype=np.uint16).reshape(shape)
                        t = torch.from_numpy(arr.copy()).view(torch.bfloat16)
                    elif dtype_str in ('F8_E4M3', 'F8_E5M2'):
                        arr = np.frombuffer(raw, dtype=np.uint8).reshape(shape)
                        t = torch.from_numpy(arr.copy()).view(torch_dtype)
                    else:
                        np_dt = np_dtype_map.get(dtype_str, np.float32)
                        arr = np.frombuffer(raw, dtype=np_dt).reshape(shape)
                        t = torch.from_numpy(arr.copy())

                    if target_device != 'cpu':
                        t = t.to(device=target_device)
                    sd[name] = t

            return sd, metadata

        def _safe_load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):
            if device is None:
                device = torch.device("cpu")

            if not ckpt.lower().endswith((".safetensors", ".sft")):
                return _original_load(ckpt, safe_load=safe_load, device=device,
                                      return_metadata=return_metadata)

            import os
            size_gb = os.path.getsize(ckpt) / (1024**3)

            # Only use streaming loader for large files (>1GB) to avoid mmap crash
            if size_gb > 1.0:
                print(f"[FLUXNATION] Loading {os.path.basename(ckpt)} ({size_gb:.1f} GB) via streaming loader")
                try:
                    import time
                    t0 = time.time()
                    sd, metadata = _read_safetensor_streaming(ckpt, device.type)
                    elapsed = time.time() - t0
                    print(f"[FLUXNATION] Loaded {len(sd)} tensors in {elapsed:.1f}s (streaming, no-mmap)")
                    if return_metadata:
                        return sd, metadata
                    return sd
                except Exception as e:
                    print(f"[FLUXNATION] Streaming load failed ({e}), falling back to original")
                    return _original_load(ckpt, safe_load=safe_load, device=device,
                                          return_metadata=return_metadata)
            else:
                return _original_load(ckpt, safe_load=safe_load, device=device,
                                      return_metadata=return_metadata)

        comfy_utils.load_torch_file = _safe_load_torch_file
        print("[FLUXNATION] Patched safetensor loading (non-mmap, Windows crash fix)")
    except Exception as e:
        print(f"[FLUXNATION] Could not patch safetensor loading: {e}")


def _try_load_extension():
    """Try to import the compiled flux_ext extension."""
    global _EXT_LOADED
    try:
        import flux_ext  # noqa: F401 — triggers TORCH_LIBRARY registration
        _EXT_LOADED = True
        print("[FLUXNATION] Loaded PyTorch C++ extension successfully")
    except ImportError:
        import torch
        import os
        root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
        for candidate in [
            os.path.join(root, "flux_ext.pyd"),
            os.path.join(root, "flux_ext.so"),
        ]:
            if os.path.exists(candidate):
                try:
                    torch.ops.load_library(candidate)
                    _EXT_LOADED = True
                    print(f"[FLUXNATION] Loaded from {candidate}")
                    break
                except Exception as e:
                    print(f"[FLUXNATION] Failed to load {candidate}: {e}")

    if not _EXT_LOADED:
        print("[FLUXNATION] Extension not found. Run build_ext.bat to compile.")
        return False

    _register_fake_kernels()
    return True


def _register_fake_kernels():
    """Register Meta/FakeTensor implementations for torch.compile."""
    import torch

    if not hasattr(torch, 'library') or not hasattr(torch.library, 'register_fake'):
        print("[FLUXNATION] torch.library.register_fake not available, skipping")
        return

    try:
        @torch.library.register_fake("flux_ext::ssb_prepare")
        def _fake_prepare(x, vec, pe, mod_w, mod_w_scale, mod_b,
                          l1w, l1w_scale, l1b, qns, kns):
            B, S = x.size(0), x.size(1)
            q    = torch.empty([B, 24, S, 128], dtype=torch.float16, device=x.device)
            k    = torch.empty([B, 24, S, 128], dtype=torch.float16, device=x.device)
            v    = torch.empty([B, 24, S, 128], dtype=torch.float16, device=x.device)
            mlp  = torch.empty([B, S, 12288],   dtype=torch.float16, device=x.device)
            gate = torch.empty([B, 3072],        dtype=torch.float16, device=x.device)
            return (q, k, v, mlp, gate)

        @torch.library.register_fake("flux_ext::ssb_finish")
        def _fake_finish(x, attn, mlp, gate, l2w, l2w_scale, l2b):
            return torch.empty_like(x, dtype=torch.bfloat16)

        @torch.library.register_fake("flux_ext::ssb_cleanup")
        def _fake_cleanup():
            return

    except Exception as e:
        print(f"[FLUXNATION] Warning: could not register fake kernels: {e}")


# ── Initialization: patch safetensor loading + load extension ──
import os as _os
_patch_safetensor_loading()

if _os.environ.get("FLUX_EXT_DISABLE", "0") == "1":
    print("[FLUXNATION] DISABLED by FLUX_EXT_DISABLE env var (stock mode)")
elif _try_load_extension():
    from .FLUXPATCH import apply_patch
    apply_patch()
else:
    print("[FLUXNATION] DISABLED — extension not found. Run build_ext.bat.")