"""The LoRA loader.

One LoRA per node, chained for more. v4 carried N LoRAs in a single node behind a
JSON list widget and a `FlexibleOptionalInputType` -- a dict subclass that lied to
ComfyUI's validator so arbitrary keys would pass. Chaining needs neither: it is
what ComfyUI's own LoraLoader does, every strength is a real socket, and the graph
shows how many LoRAs are actually applied instead of hiding the count in a string.
"""

import comfy.utils
import folder_paths
from comfy_api.latest import io

from ..._core import log
from .keys import match


NONE = "None"


class FunPackLoraLoader(io.ComfyNode):
    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="FunPackLoraLoader",
            display_name="FunPack LoRA Loader",
            category="FunPack/Loaders",
            description="Apply one LoRA to a model, and optionally its text encoder. "
                        "Chain these for more than one.",
            inputs=[
                io.Model.Input("model"),
                io.Combo.Input("lora_name", options=[NONE] + folder_paths.get_filename_list("loras"),
                               default=NONE, tooltip="None leaves the model untouched."),
                io.Float.Input("strength_model", default=1.0, min=-100.0, max=100.0, step=0.01,
                               tooltip="How strongly the LoRA changes the model. 0 = not at all, negative = the opposite."),
                io.Clip.Input("clip", optional=True,
                              tooltip="Leave unwired for a model-only LoRA."),
                io.Float.Input("strength_clip", default=1.0, min=-100.0, max=100.0, step=0.01,
                               optional=True, tooltip="How strongly the LoRA changes the text encoder. Only used when a clip is wired."),
            ],
            outputs=[
                io.Model.Output(display_name="model"),
                io.Clip.Output(display_name="clip"),
                io.String.Output(display_name="status"),
            ],
        )

    @classmethod
    def execute(cls, model, lora_name: str, strength_model: float,
                clip=None, strength_clip: float = 1.0) -> io.NodeOutput:
        if lora_name == NONE:       # kept in the pipeline so a LoRA is one pick away; nothing picked = no effect
            return io.NodeOutput(model, clip, "no LoRA picked")
        # Doing nothing is a valid request, but it must return the SAME objects
        # rather than clones: a clone here would silently drop patches a later
        # loader in the chain applied to the original.
        if strength_model == 0 and (clip is None or strength_clip == 0):
            # Loud, because "applied a LoRA at strength 0" is indistinguishable
            # from "forgot to set the strength" by looking at the picture.
            log.alert("FunPack LoRA Loader",
                      f"{lora_name} has both strengths at 0 and was not applied")
            return io.NodeOutput(model, clip, f"{lora_name}: both strengths are 0, not applied")

        path = folder_paths.get_full_path_or_raise("loras", lora_name)
        lora = comfy.utils.load_torch_file(path, safe_load=True)
        patches, how, dropped = match(model, clip, lora)
        if not patches:
            log.alert("FunPack LoRA Loader", f"{lora_name} matched NOTHING in this model, so it does nothing: "
                      "it is for another model, or names its layers in a way not known here")
            return io.NodeOutput(model, clip, f"{lora_name}: matched nothing, not applied")
        patched_model = model.clone()
        patched_model.add_patches(patches, strength_model)
        patched_clip = None
        if clip is not None:
            patched_clip = clip.clone()
            patched_clip.add_patches(patches, strength_clip)
        if dropped:
            log.alert("FunPack LoRA Loader", f"{lora_name}: {dropped} of {len(patches) + dropped} weights do not fit this "
                      "model's shapes and were left out (trained on another variant of it)")

        applied = [f"model={strength_model}"]
        applied.append(f"clip={strength_clip}" if clip is not None else "clip=not wired")
        found = f"{len(patches)} weights" + (f", keys {how}" if how != "as-is" else "") + (f", {dropped} dropped" if dropped else "")
        log.info("FunPack LoRA Loader", f"{lora_name} applied ({', '.join(applied)}; {found})")
        return io.NodeOutput(patched_model, patched_clip,
                             f"FunPack LoRA Loader | {lora_name} ({', '.join(applied)}; {found})")