from __future__ import annotations

import copy
import torch

from comfy.patcher_extension import PatcherInjection
from comfy.weight_adapter.bypass import BypassForwardHook
from comfy_api.latest import io

from .helper_functions import parse_expr
from .Parser.UnifiedMathVisitor import UnifiedMathVisitor
from .Stack import MrmthStack
from .ParseTree import MrmthParseTree

@staticmethod
def _parse_strengths(Expression, stack):
    tree = None
    variables = {}
    if isinstance(Expression,str):
        tree = parse_expr(Expression)
    else:
        tree = Expression
    visitor = UnifiedMathVisitor(variables, [1],torch.device("cpu"),state_storage=stack)
    result = visitor.visit(tree)
    return (list(result), stack)

class BatchLoraHooksNode(io.ComfyNode):
    """Scale existing bypass LoRA injections on a model/clip pair. Works with models/loaders/Load lora (Bypass) (For Debugging) or more_math batch lora node."""

    @staticmethod
    def _get_bypass_hooks(injections):
        if injections is None:
            return []

        hooks = []
        for injection in injections:
            closure = getattr(getattr(injection, "inject", None), "__closure__", None)
            if not closure:
                continue
            for cell in closure:
                cell_value = cell.cell_contents
                if hasattr(cell_value, "hooks") and hasattr(cell_value, "create_injections"):
                    manager_hooks = getattr(cell_value, "hooks", None)
                    if manager_hooks:
                        hooks.extend(manager_hooks)
                    break
        return hooks

    @staticmethod
    def _build_bypass_injection(hooks, scale: float):
        scaled_hooks = []
        for hook in hooks:
            adapter = copy.deepcopy(hook.adapter)
            scaled_hooks.append(BypassForwardHook(hook.module, adapter, multiplier=hook.multiplier * scale))

        def inject_all(model_patcher):
            for hook in scaled_hooks:
                hook.inject()

        def eject_all(model_patcher):
            for hook in scaled_hooks:
                hook.eject()

        return [PatcherInjection(inject=inject_all, eject=eject_all)]

    @classmethod
    def _scale_bypass_injections(cls, patcher, scale: float):
        if patcher is None:
            return

        injections = getattr(patcher, "get_injections", lambda _key: None)("bypass_lora")
        if injections is None:
            return

        hooks = cls._get_bypass_hooks(injections)
        if not hooks:
            return

        was_injected = getattr(patcher, "is_injected", False)
        if was_injected:
            patcher.eject_model()

        patcher.set_injections("bypass_lora", cls._build_bypass_injection(hooks, scale))

        if was_injected:
            patcher.inject_model()

    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="mrmth_LoraBatchRescale",
            category="More math",
            display_name="Scale LoRA (Batch)",
            inputs=[
                io.Model.Input(id="model", tooltip="The model with bypass LoRA injections"),
                io.Clip.Input(id="clip", tooltip="The clip with bypass LoRA injections"),
                io.MultiType.Input(
                    io.String.Input("model_strength", default="[1.0]", multiline=False),
                    types=[io.String, MrmthParseTree],
                    tooltip="List of model-side LoRA strengths",
                ),                io.MultiType.Input(
                    io.String.Input("clip_strength", default="[1.0]", multiline=False),
                    types=[io.String, MrmthParseTree],
                    tooltip="List of clip-side LoRA strengths",
                ),
                MrmthStack.Input(id="stack", optional=True)
            ],
            outputs=[
                io.Model.Output(id="model", tooltip="The model with scaled bypass LoRA injections", is_output_list=True),
                io.Clip.Output(id="clip", tooltip="The clip with scaled bypass LoRA injections", is_output_list=True),
                MrmthStack.Output(id="stack", tooltip="The stack"),
            ],
        )

    @classmethod
    def execute(cls, model, clip, model_strength, clip_strength, stack=None):
        stack = copy.deepcopy(stack) if stack is not None else {}
        weights0, stack = _parse_strengths(model_strength, {})
        weights1, stack = _parse_strengths(clip_strength, {})
        scaled_models = []
        scaled_clips = []
        if len(weights0) != len(weights1):
            raise ValueError("model_strength and clip_strength must have the same number of elements. Model: {}, Clip: {}".format(len(weights0), len(weights1)))
        for w0,w1 in zip(weights0, weights1):
            scaled_model = model.clone()
            scaled_clip = clip.clone()
            cls._scale_bypass_injections(scaled_model, w0)
            clip_patcher = getattr(scaled_clip, "patcher", None)
            cls._scale_bypass_injections(clip_patcher, w1)
            scaled_models.append(scaled_model)
            scaled_clips.append(scaled_clip)

        return (scaled_models, scaled_clips, stack)