"""Shared loading options.

Not a module: it announces nothing and ships no node. `registry.scan` only walks
two levels of DIRECTORIES, so a plain file at the domain level is a helper the
loaders import and nothing else.

These options are on the loader node rather than in a settings payload, and the
line is worth stating: a loader widget describes the FILE and how it is read,
which the loader cannot be handed by anyone else. Feature knobs -- strengths,
sigma windows, anything a modifier owns -- are settings, and never appear here.
"""

import json

import torch

from .._core import log

# Named exactly as core's UNETLoader names them where they overlap, so a value
# copied from a stock workflow means the same thing here.
WEIGHT_DTYPES = ["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2", "fp16", "bf16", "fp32"]
COMPUTE_DTYPES = ["default", "fp16", "bf16", "fp32"]
VAE_DTYPES = ["bfloat16", "float16", "float32"]

_BY_NAME = {
    "fp16": torch.float16,
    "bf16": torch.bfloat16,
    "fp32": torch.float32,
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float32": torch.float32,
    "fp8_e4m3fn": getattr(torch, "float8_e4m3fn", None),
    "fp8_e5m2": getattr(torch, "float8_e5m2", None),
    # fp8_e4m3fn plus comfy's fp8 matmul path; see weight_model_options.
    "fp8_e4m3fn_fast": getattr(torch, "float8_e4m3fn", None),
}


def dtype_of(name):
    """The torch dtype a choice names, or None for 'default'/an unsupported build."""
    return _BY_NAME.get(str(name or "default"))


def unsupported_dtypes(model, weight_dtype, compute_dtype):
    """A forced dtype the model does not run in, one line each. A model says which it runs in (H3: bf16/fp32;
    v4 measured fp16 breaking it). Allowed, but said."""
    supported = getattr(getattr(getattr(model, "model", None), "model_config", None), "supported_inference_dtypes", None) or []
    runs_in = ", ".join(str(d).replace("torch.", "") for d in supported)
    return [f"{what} dtype {picked} is not one this model runs in ({runs_in})"
            for what, picked in (("weight", weight_dtype), ("compute", compute_dtype))
            if picked in ("fp16", "bf16", "fp32") and supported and dtype_of(picked) not in supported]


def fp32_over_bf16(weight_dtype, compute_dtype, stored=None):
    """Why fp32 compute over bf16 weights is noisy, or None. It runs, but the activations go fp32 while the norm layers keep
    their bf16 weights, so torch warns "Mismatch dtype between input and weight". Allowed, but said."""
    if compute_dtype == "fp32" and (weight_dtype == "bf16" or (weight_dtype == "default" and stored == "bf16")):
        return "bf16 weights dtype with fp32 precision is a bad combo for H3. Don't expect miracles: it runs, but slowly."


_SHORT = {"bfloat16": "bf16", "float16": "fp16", "float32": "fp32"}


def stored_weights(state_dict):
    """What a file's weights are stored as, by element count: "bf16", or "int8+bf16" for a
    mixed file. Read from the tensors, not the weight_dtype widget: a pre-quantized file
    loaded with weight_dtype bf16 still runs int8 layers, and the log must say so."""
    sizes = {}
    for t in state_dict.values():
        if torch.is_tensor(t) and t.dim() >= 2:
            name = str(t.dtype).replace("torch.", "")
            sizes[_SHORT.get(name, name)] = sizes.get(_SHORT.get(name, name), 0) + t.numel()
    total = sum(sizes.values())
    return "+".join(n for n in sorted(sizes, key=sizes.get, reverse=True) if sizes[n] > total * 0.1) or None


def slow_int8_build():
    """Why int8 models sample slowly on this torch build, or None.

    Measured on an RTX PRO 6000 Blackwell, same 4-step H3 run: torch 2.10+cu128 sampled the
    int8 model ~2x SLOWER than bf16; the same torch from the cu130 index made int8 ~25% faster.
    """
    cuda = getattr(torch.version, "cuda", None)
    try:
        if not cuda or not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10:
            return None
        if int(cuda.split(".")[0]) >= 13:
            return None
    except (RuntimeError, ValueError):
        return None
    # --force-reinstall is required: pip counts 2.10.0+cu128 as already satisfying torch==2.10.0.
    return (f"torch is built for CUDA {cuda}: on this Blackwell GPU int8 models sample about 2x slower "
            f"than on a CUDA 13 build. Stop ComfyUI and reinstall torch, torchvision and torchaudio "
            f"(same versions) with --force-reinstall --index-url https://download.pytorch.org/whl/cu130")


def weight_model_options(weight_dtype):
    """comfy `model_options` for a weight dtype. The mapping core's UNETLoader uses."""
    options = {}
    dtype = dtype_of(weight_dtype)
    if dtype is not None:
        options["dtype"] = dtype
    if weight_dtype == "fp8_e4m3fn_fast":
        options["fp8_optimizations"] = True
    return options


def attention_choices():
    """Backends this ComfyUI can actually run.

    Read from ComfyUI's own registry rather than a private import, so the list
    offers exactly what the machine has -- sage where SageAttention is installed,
    flash where flash-attn is -- and never lies about the rest. 'default' leaves
    whatever ComfyUI itself selected from its launch flags.
    """
    names = []
    try:
        from comfy.ldm.modules.attention import REGISTERED_ATTENTION_FUNCTIONS
        names = sorted(REGISTERED_ATTENTION_FUNCTIONS.keys())
    except Exception:                            # noqa: BLE001 -- older ComfyUI
        pass
    return ["default"] + names


def attention_override(name):
    """A transformer_options override routing every attention call to `name`.

    ComfyUI wraps each implementation with `wrap_attn`, which hands the override
    the original function plus the call's arguments -- so calling the WRAPPED
    replacement would re-enter that machinery. The unwrapped one is what gets
    called.
    """
    if not name or name == "default":
        return None
    try:
        from comfy.ldm.modules.attention import get_attention_function
    except ImportError:
        log.warning("attention", "this ComfyUI has no attention registry, so the "
                                 "chosen backend cannot be applied")
        return None
    chosen = get_attention_function(name, None)
    if chosen is None:
        log.warning("attention", f"backend {name!r} is not available on this machine, "
                                 f"so it was not applied")
        return None
    inner = getattr(chosen, "__wrapped__", chosen)

    def override(_func, *args, **kwargs):
        return inner(*args, **kwargs)

    # Some backends carry a zero-copy path that takes the attention tensors still
    # in their containers, and `wrap_attn` only offers it when the OVERRIDE has
    # the attribute -- otherwise it unpacks each container and calls the slow way.
    # Forgetting to forward it is invisible: the backend is selected, the maths is
    # right, and the fast path it exists for is silently never used. ComfyUI's own
    # `ModelPatcher.set_model_optimized_attention` forwards it for this reason.
    container_function = getattr(chosen, "container_function", None)
    if container_function is not None:
        override.container_function = container_function

    return override


def set_fp16_accumulation(enabled):
    """torch's fp16 accumulation switch, or None where the build has no such knob.

    PROCESS-WIDE, and there is no scoped version of it: torch has one flag for the
    whole interpreter. So this cannot mean "for this model" however it is spelled
    on a node -- a second loader setting it differently changes the first model's
    maths, and because ComfyUI caches loader outputs the last value set outlives
    the run that set it, reaching models loaded by core's own nodes.

    It is therefore never changed silently. Every transition says so once, which
    is the difference between a documented global and an invisible one. The real
    fix is to move this to a machine-level setting, where its scope is honest.
    """
    matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
    if matmul is None or not hasattr(matmul, "allow_fp16_accumulation"):
        return None
    wanted = bool(enabled)
    before = bool(getattr(matmul, "allow_fp16_accumulation", False))
    matmul.allow_fp16_accumulation = wanted
    if before != wanted:
        log.alert("fp16 accumulation",
                  f"now {wanted} for EVERY model in this process, not just this one -- "
                  f"torch has a single process-wide flag")
    return wanted


def head_bank_metadata(state_dict, metadata):
    """File metadata that declares how many output heads a few-step H3 file stacks, or the input untouched.

    A file with several stacked heads per stream normally says so in its safetensors
    config (`transformer.pdd_steps`). A re-saved copy can lose that line while keeping
    explicit head sizes, and ComfyUI then builds a single head and fails on the biases.
    The count is in the weights anyway: the patch projections hold the true per-head
    width, and the head bias holds that width times the count.

    Reads shapes only, so it works on any device, dtype or quantised storage.
    """
    try:
        bias = state_dict["final_layer.video_out.bias"]
        abias = state_dict["final_layer.audio_out.bias"]
        video = state_dict["video_patch_proj.weight"].shape[1]
        audio = state_dict["audio_patch_proj.weight"].shape[1]
    except (KeyError, AttributeError, IndexError):
        return metadata, None
    if not video or not audio or bias.shape[0] % video or abias.shape[0] % audio:
        return metadata, None
    steps = bias.shape[0] // video
    if steps < 2 or abias.shape[0] // audio != steps:
        return metadata, None
    try:
        config = json.loads((metadata or {}).get("config") or "{}")
    except ValueError:
        config = {}
    transformer = config.setdefault("transformer", {})
    if transformer.get("pdd_steps", 1) == steps:
        return metadata, None
    declared = transformer.get("pdd_steps")
    transformer["pdd_steps"] = steps
    fixed = dict(metadata or {})
    fixed["config"] = json.dumps(config, separators=(",", ":"))
    return fixed, (f"output heads: the file stacks {steps} per stream but its metadata "
                   f"{'says ' + str(declared) if declared else 'does not say so'}; using {steps}")


_loaded_as = {}        # model file -> [run id, the settings seen in that run, the settings seen in the run before]


def free_stale_copy_after_run(name, settings) -> bool:
    """Loading a file again with other settings (attention, SLA method, dtype...) leaves the earlier copy in ComfyUI's node cache,
    until RAM runs short and ComfyUI evicts it half way, which is when "memory leak with model X" shows up. So when the same file
    comes back with settings the run before did not have, ask ComfyUI to unload its models and reset its cache once this run ends
    (its own /free flag): the next run loads the file once more, and only one copy was ever around. Two loaders of one file in one
    graph are no change: each run sees the same set of settings, so nothing is flagged and nothing reloads every run."""
    try:
        from server import PromptServer
        server = PromptServer.instance
        running = list((getattr(server.prompt_queue, "currently_running", None) or {}).values())      # ComfyUI's running item: (number, prompt id, ...)
        run = running[0][1] if running else object()                        # no run id (a plain script, a test): every call is its own run
    except Exception:           # noqa: BLE001 -- no server (tests, a plain script): nothing to ask
        return False
    key = json.dumps(settings, sort_keys=True, default=str)
    seen = _loaded_as.get(name)
    if seen is None or seen[0] != run:
        seen = _loaded_as[name] = [run, set(), seen[1] if seen else None]
    seen[1].add(key)
    previous = seen[2]
    if previous is None or key in previous:
        return False
    try:
        server.prompt_queue.set_flag("free_memory", True)
    except Exception as exc:    # noqa: BLE001 -- a ComfyUI without the flag: said, so the leak warning is not a mystery
        log.warning("FunUI Diffusion Model Loader", f"{name} came back with other settings, but this ComfyUI cannot be asked to free the earlier copy ({type(exc).__name__}); "
                    "a 'memory leak with model' warning may follow")
        return False
    log.info("FunUI Diffusion Model Loader",
             f"{name} was loaded before with other settings: the earlier copy is freed when this run ends, so the next run loads it once more")
    return True