from .helper_functions import generate_dim_variables, parse_expr, as_tensor, get_v_variable, get_f_variable, get_tensor_device, LazyVariableDict, checkLazyNew
from .Parser.UnifiedMathVisitor import MathDict, ChangeTrackingDict, UnifiedMathVisitor
import comfy
import comfy.model_management
import comfy.model_patcher
import comfy.lora
import torch


def _get_model_patcher(obj):
    # A MODEL input *is* a ModelPatcher already - not something wrapping one.
    if isinstance(obj, comfy.model_patcher.ModelPatcher):
        return obj
    patcher = getattr(obj, "patcher", None)
    if patcher is None and hasattr(obj, "model"):
        patcher = getattr(obj.model, "patcher", None)
    return patcher


def get_effective_state_dict(obj):
    if hasattr(obj, "patcher") and hasattr(obj.patcher, "model_state_dict"):
        return obj.patcher.model_state_dict()
    if hasattr(obj, "model_state_dict"):
        return obj.model_state_dict()
    if hasattr(obj, "model") and hasattr(obj.model, "patcher") and hasattr(obj.model.patcher, "model_state_dict"):
        return obj.model.patcher.model_state_dict()
    if hasattr(obj, "model") and hasattr(obj.model, "model_state_dict"):
        return obj.model.model_state_dict()
    if hasattr(obj, "model") and hasattr(obj.model, "state_dict"):
        return obj.model.state_dict()
    if hasattr(obj, "state_dict"):
        return obj.state_dict()
    return None


def get_base_weight(obj, key):
    """True pre-patch weight for `key`, regardless of whether the patcher
    currently has this key's patches baked into the live model (`backup`
    holds the pristine value whenever that happened; `use_ejected` does not
    undo weight patches, only injections, so it can't be relied on here)."""
    patcher = _get_model_patcher(obj)
    if patcher is not None:
        backup = getattr(patcher, "backup", None)
        if backup is not None:
            bk = backup.get(key, None)
            if bk is not None:
                return bk.weight
        model = getattr(patcher, "model", None)
        if model is not None:
            try:
                weight, _, _ = comfy.model_patcher.get_key_weight(model, key)
                return weight
            except AttributeError:
                pass

    sd = get_effective_state_dict(obj)
    if sd is not None:
        return sd.get(key, None)
    return None


def get_patched_weight(obj, key, base_weight):
    """Fully patched weight for `key`, computed fresh from `base_weight` so
    the result never depends on whether the model happens to be baked
    in-place already (reusing an already-baked weight as input would
    double-apply patches whose strength_model != 1.0)."""
    if base_weight is None:
        return None
    patcher = _get_model_patcher(obj)
    if patcher is not None:
        patches = getattr(patcher, "patches", None)
        if patches is not None and key in patches:
            return comfy.lora.calculate_weight(patches[key], base_weight.clone(), key)
    return base_weight


def coerce_like(weight, reference, device=None):
    if isinstance(weight, torch.Tensor) and isinstance(reference, torch.Tensor):
        target_device = device if device is not None else reference.device
        if weight.device != target_device or weight.dtype != reference.dtype:
            return weight.to(device=target_device, dtype=reference.dtype)
    return weight


def calculate_patches(Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0, use_compute_device=False):
    """Legacy calculate_patches for backward compatibility."""
    return calculate_patches_autogrow(Model, V={"V0": a, "V1": b, "V2": c, "V3": d}, F={"F0": w, "F1": x, "F2": y, "F3": z}, pbar=None, mapping={"a": "V0", "b": "V1", "c": "V2", "d": "V3", "w": "F0", "x": "F1", "y": "F2", "z": "F3"}, use_compute_device=use_compute_device)

def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=None, use_compute_device=False):
    """
    Calculate patches for model-like objects (Model, VAE, CLIP) using Autogrow inputs.
    Iterates over the UNION of keys from all input models to support merging disjoint architectures/patches.

    In autogrow mode:
    - The expression is evaluated once PER KEY (per-layer iteration)
    - V0, V1, V2, V3 contain only the weight tensor for the current key (not all keys)
    - Expression must return a TENSOR (the patch for this key)
    - Returning a dict will raise an error; use dict_mode=True instead

    Performance note:
    - If expression applies the same operation to all keys (e.g., V0 *= 2), it will iterate
      through each key in Python, making it O(num_keys) in Python overhead
    - For efficiency with bulk operations on multiple keys, use dict_mode=True instead
    - dict_mode evaluates the expression once with V0 as a complete dict, allowing batch operations

    Args:
        Expr: The math expression string or parse tree.
        V: Dictionary of input models (Autogrow Input).
        F: Dictionary of input floats (Autogrow Float).
        mapping: Optional dict mapping legacy alias names to V keys (e.g. {"a": "V0"}).
    """
    if V is None: V = {}
    if F is None: F = {}
    if mapping is None: mapping = {}

    needed_vars = checkLazyNew(Expr, V, F)
    needed_v_names = [v_name for v_name in V.keys() if v_name in needed_vars]
    if not needed_v_names:
        needed_v_names = [v_name for v_name in V.keys() if V.get(v_name) is not None]

    # Collect all unique keys from all models, preserving order from first model
    all_keys_list = []  # Preserves order
    seen_keys = set()   # Fast O(1) duplicate checking

    models = [v for v in V.values() if v is not None]
    if not models:
        return {}

    for m in models:
        sd = get_effective_state_dict(m)
        if sd is None:
            sd_keys = []
        else:
            sd_keys = sd.keys()

        for key in sd_keys:
            if key not in seen_keys:
                seen_keys.add(key)
                all_keys_list.append(key)

    tree = None
    if isinstance(Expr,str):
        tree = parse_expr(Expr)
    else:
        tree = Expr
    patches = {}
    original_device = None
    compute_device = comfy.model_management.get_torch_device() if use_compute_device else None
    needs_deltas = not (isinstance(Expr, str) and "_d" not in Expr)
    const_f_vars = {k: val if val is not None else 0.0 for k, val in F.items()}
    const_f_alias_vars = {
        alias: (F[target] if F[target] is not None else 0.0)
        for alias, target in mapping.items()
        if target in F
    }
    mapping_items = tuple(mapping.items())

    current_state_dicts = {}
    for v_name in needed_v_names:
        v_val = V.get(v_name)
        if v_val is None:
            continue
        current_sd = get_effective_state_dict(v_val)
        if current_sd is None:
            continue
        if original_device is None:
            original_device = get_tensor_device(current_sd)
        current_state_dicts[v_name] = current_sd

    # Progress bar if possible (comfy.utils.ProgressBar might assume unthreaded?)
    # Just skip for utility or use if substantial.
    layer_count = len(all_keys_list)

    def make_lazy(value_fn):
        def lazy_eval():
            return value_fn()
        lazy_eval.is_lazy_var = True
        return lazy_eval

    for layer_idx, key in enumerate(all_keys_list):

        variables = LazyVariableDict(const_f_vars)
        # Also populate mapped aliases for F (w, x, y, z)
        variables.update(const_f_alias_vars)

        variables["L"] = float(layer_idx)
        variables["layer"] = float(layer_idx)
        variables["LC"] = float(layer_count)
        variables["layer_count"] = float(layer_count)
        variables["K"] = key
        variables["key"] = key

        # Inject weights for this key from V models
        valid_key = False
        ref_tensor = None
        present_weights = {}
        present_bases = {}

        for v_name in V.keys():
            current_sd = current_state_dicts.get(v_name)
            if current_sd is None:
                continue
            raw_tensor = current_sd.get(key, None)
            if raw_tensor is not None:
                v_val = V.get(v_name)
                base_tensor = get_base_weight(v_val, key) if v_val is not None else None
                if base_tensor is None:
                    base_tensor = raw_tensor
                base_tensor = coerce_like(base_tensor, raw_tensor, device=compute_device)
                patched_tensor = get_patched_weight(v_val, key, base_tensor) if v_val is not None else None
                if patched_tensor is None:
                    patched_tensor = raw_tensor
                patched_tensor = coerce_like(patched_tensor, raw_tensor, device=compute_device)
                present_weights[v_name] = patched_tensor
                if needs_deltas:
                    present_bases[v_name] = base_tensor
                if ref_tensor is None:
                    ref_tensor = patched_tensor
                valid_key = True

        if not valid_key:
            continue

        # Find reference shape
        if ref_tensor is None:
            continue # Should not happen if valid_key is true

        # Fill missing models with lazy values so V[n] keeps working
        v_cache = [None]

        def lazy_v_stack():
            if v_cache[0] is None:
                v_cache[0] = get_v_variable(variables)
            return v_cache[0][0]

        lazy_v_stack.is_lazy_var = True

        def lazy_v_count():
            if v_cache[0] is None:
                v_cache[0] = get_v_variable(variables)
            return float(v_cache[0][1])

        lazy_v_count.is_lazy_var = True

        for v_name in V.keys():
            if v_name in present_weights:
                w_tensor = present_weights[v_name]
                variables[v_name] = make_lazy(lambda wt=w_tensor: wt)
                if needs_deltas:
                    base_tensor = present_bases.get(v_name, w_tensor)
                    variables[f"{v_name}_d"] = make_lazy(lambda wt=w_tensor, bt=base_tensor: wt - bt)
            else:
                variables[v_name] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))
                if needs_deltas:
                    variables[f"{v_name}_d"] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))

        # Populate aliases for V (a, b, c, d)
        for alias, target in mapping_items:
            if target in variables:
                variables[alias] = variables[target]
            elif target in V:
                # V key exists in input dict but this specific layer key is missing
                variables[alias] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))
            else:
                # Target doesn't exist at all (e.g., V1 not provided)
                # Check if it's a V-key pattern and zero-fill
                if target.startswith("V") and target[1:].isdigit():
                    variables[alias] = make_lazy(lambda rt=ref_tensor: torch.zeros_like(rt))

        variables["V"] = lazy_v_stack
        variables["Vcnt"] = lazy_v_count
        variables["V_count"] = lazy_v_count

        f_stacked, f_cnt = get_f_variable(F)
        if f_stacked is not None:
             variables["F"] = f_stacked
             variables["Fcnt"] = float(f_cnt)
             variables["F_count"] = float(f_cnt)

        variables = variables | generate_dim_variables(ref_tensor)

        # Read V0 before visiting, since the expression may rebind it
        original_v0 = variables.get("V0")

        visitor = UnifiedMathVisitor(variables, ref_tensor.shape,state_storage=stack)
        res = visitor.visit(tree)

        # In autogrow mode, expression must return a single tensor (patch for this key)
        # If a dict is returned, that's an error - use dict_mode instead
        if isinstance(res, dict) and not isinstance(res, torch.Tensor):
            raise ValueError(f"Autogrow mode expects tensor output, got dict. Use dict_mode=True for dict output.")

        res = as_tensor(res, ref_tensor.shape)

        # Extract original for this specific key
        original = None
        if original_v0 is not None:
            if isinstance(original_v0, torch.Tensor):
                original = original_v0
            elif isinstance(original_v0, dict):
                # Should not happen in autogrow mode - V0 is always a tensor
                original = original_v0.get(key)

        if original is None:
            original = torch.zeros_like(res)

        diff = res - original

        # Clean up: don't store zero patches
        if not torch.all(diff == 0):
            patches[key] = (diff,)

        if pbar is not None:
            pbar.update(1)

    if compute_device is not None and original_device is not None and original_device != compute_device:
        patches = {
            key: ((value[0].to(device=original_device),) if isinstance(value, tuple) and value and torch.is_tensor(value[0]) else value)
            for key, value in patches.items()
        }

    return patches


def calculate_patches_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=None, use_compute_device=False):
    """
    Calculate patches for model-like objects (Model, VAE, CLIP) using Dictionary mode.
    Evaluates the expression once for all keys simultaneously.

    In dict_mode:
    - The expression is evaluated ONCE (not per-key)
    - V0, V1, V2, V3 are MathDict objects containing all weights across all keys
    - Expression can modify individual keys: V0[key] *= 2
    - Expression can apply bulk operations: V0 *= 2
    - Expression must return a DICT/MATHDICT with exactly the same keys as the input models
    - The returned dict must preserve the original key structure (cannot add/remove keys)

    Performance:
    - More efficient than autogrow_mode for bulk operations since evaluation is done once
    - Suitable for complex transformations that need access to all keys simultaneously

    Args:
        Expr: The math expression string or parse tree.
        V: Dictionary of input models (Dict Mode Input).
        F: Dictionary of input floats (Dict Mode Float).
        mapping: Optional dict mapping legacy alias names to V keys.
    """
    if V is None:
        V = {}
    if F is None:
        F = {}
    if mapping is None:
        mapping = {}

    needed_vars = checkLazyNew(Expr, V, F)
    needed_v_names = [name for name in V if name in needed_vars]
    if "V0" in V and "V0" not in needed_v_names:
        needed_v_names.append("V0")
    if not needed_v_names:
        needed_v_names = [name for name, value in V.items() if value is not None]

    models = [value for value in V.values() if value is not None]
    if not models:
        return {}

    all_keys = []
    seen_keys = set()
    for model in models:
        state_dict = get_effective_state_dict(model)
        if state_dict is None:
            continue
        for key in state_dict:
            if key not in seen_keys:
                seen_keys.add(key)
                all_keys.append(key)

    tree = parse_expr(Expr) if isinstance(Expr, str) else Expr
    compute_device = comfy.model_management.get_torch_device() if use_compute_device else None
    original_device = None
    state_dicts = {}
    for name in needed_v_names:
        value = V.get(name)
        if value is None:
            continue
        state_dict = get_effective_state_dict(value)
        if state_dict is None:
            continue
        if original_device is None:
            original_device = get_tensor_device(state_dict)
        state_dicts[name] = state_dict

    if not state_dicts:
        return {}

    all_keys = [key for key in all_keys if any(key in state_dict for state_dict in state_dicts.values())]

    # Mirrors calculate_patches_autogrow's needs_deltas: a plain substring
    # check on the expression text, since checkLazyNew's returned variable
    # names never carry "_d" suffixes (only bare V{n}/alias names), so a
    # membership test against needed_vars can never match here.
    needs_deltas = not (isinstance(Expr, str) and "_d" not in Expr)
    delta_names = set(needed_v_names) if needs_deltas else set()
    weights = {name: {} for name in needed_v_names}
    bases = {name: {} for name in delta_names}
    references = {}

    for key in all_keys:
        for name in needed_v_names:
            state_dict = state_dicts.get(name)
            raw_weight = state_dict.get(key) if state_dict is not None else None
            if raw_weight is None:
                continue
            model = V[name]
            base_weight = get_base_weight(model, key)
            if base_weight is None:
                base_weight = raw_weight
            base_weight = coerce_like(base_weight, raw_weight, device=compute_device)
            patched_weight = get_patched_weight(model, key, base_weight)
            if patched_weight is None:
                patched_weight = raw_weight
            patched_weight = coerce_like(patched_weight, raw_weight, device=compute_device)
            weights[name][key] = patched_weight
            if name in delta_names:
                bases[name][key] = base_weight
            if key not in references:
                references[key] = patched_weight

    if not references:
        return {}

    def weight_dict(name):
        values = {key: weights[name].get(key, torch.zeros_like(references[key])) for key in all_keys}
        if name == "V0":
            return ChangeTrackingDict(values)
        return MathDict(values)

    variables = LazyVariableDict({name: value if value is not None else 0.0 for name, value in F.items()})
    for alias, target in mapping.items():
        if target in F:
            variables[alias] = F[target] if F[target] is not None else 0.0

    for name in needed_v_names:
        variables[name] = weight_dict(name)

    for alias, target in mapping.items():
        if target in variables:
            variables[alias] = variables[target]

    for name in delta_names:
        def make_delta_dict(name=name):
            return MathDict({
                key: weights[name][key] - bases[name][key] if key in weights[name] else torch.zeros_like(references[key])
                for key in all_keys
            })

        make_delta_dict.is_lazy_var = True
        variables[f"{name}_d"] = make_delta_dict

    visitor_device = compute_device if compute_device is not None else next(iter(references.values())).device
    # Snapshot V0's entries BEFORE visiting, since the expression may reassign them
    v0_dict = variables.get("V0", MathDict({key: torch.zeros_like(reference) for key, reference in references.items()}))
    originals = dict(v0_dict)
    result = UnifiedMathVisitor(variables, device=visitor_device, state_storage=stack).visit(tree)
    if not isinstance(result, (MathDict, dict)):
        if isinstance(result, (float, int)):
            result = MathDict({key: torch.full_like(reference, result) for key, reference in references.items()})
        else:
            raise ValueError("Dictionary mode expressions must return a dictionary")

    if tuple(result) != tuple(all_keys):
        raise ValueError("Dictionary mode expressions must preserve all model keys")

    patches = {}

    # Optimization: if result is a ChangeTrackingDict, only process changed keys
    changed_keys_set = None
    if isinstance(result, ChangeTrackingDict):
        changed_keys_set = result.get_changed_keys()
        keys_to_process = changed_keys_set
    else:
        # Fallback: process all keys
        keys_to_process = all_keys

    for key in keys_to_process:
        value = result.get(key)
        if value is None:
            # Key was removed or doesn't exist
            if pbar is not None:
                pbar.update(1)
            continue

        if not torch.is_tensor(value):
            value = torch.full_like(references[key], value)

        diff = value - originals[key]
        if not torch.all(diff == 0):
            patches[key] = (diff,)

        if pbar is not None:
            pbar.update(1)

    if compute_device is not None and original_device is not None and original_device != compute_device:
        patches = {key: (value[0].to(device=original_device),) for key, value in patches.items()}
    return patches