"""Convert a ComfyUI-native Wan-style Qwen-Image VAE state dict (comfy.ldm.wan.vae.WanVAE,
keys like `encoder.downsamples.0.residual.0.gamma`) into the key layout expected by
diffusers' AutoencoderKLQwenImage (keys like `encoder.down_blocks.0.norm1.gamma`).

Both classes are literal re-implementations of the same architecture (diffusers'
own docstring: "QwenImageVAE is further fine-tuned from the Wan Video VAE"), so the
mapping below mirrors each class's construction loop exactly rather than hardcoding
per-key tables. Every call is validated against the real diffusers module's own
state_dict() keys/shapes before use, so any topology mismatch (e.g. an unexpected
config) raises loudly instead of silently loading misaligned weights.
"""

RESIDUAL_SUBKEYS = {
    "residual.0": "norm1",
    "residual.2": "conv1",
    "residual.3": "norm2",
    "residual.6": "conv2",
}


def _flat_encoder_layout(dim_mult, num_res_blocks, attn_scales):
    """Reproduce Encoder3d/QwenImageEncoder3d's flat down-path construction.

    Returns a list of ('res'|'attn'|'down', flat_index) — identical ordering in both
    comfy's `downsamples` Sequential and diffusers' `down_blocks` ModuleList, since
    both are built by the same loop over the same config.
    """
    layout = []
    scale = 1.0
    n_stages = len(dim_mult)
    for i in range(n_stages):
        for _ in range(num_res_blocks):
            layout.append("res")
            if scale in attn_scales:
                layout.append("attn")
        if i != n_stages - 1:
            layout.append("down")
            scale /= 2.0
    return layout


def _flat_decoder_stage_layout(dim_mult, num_res_blocks, attn_scales):
    """Reproduce Decoder3d/QwenImageDecoder3d's up-path construction.

    Returns a list of stages; each stage is a list of ('res'|'attn') entries plus a
    trailing 'up' marker if that stage upsamples. comfy's `upsamples` Sequential is the
    flat concatenation of all stages; diffusers' `up_blocks` is one QwenImageUpBlock per
    stage (its own `.resnets` / `.upsamplers`).
    """
    stages = []
    scale = 1.0 / 2 ** (len(dim_mult) - 2)
    n_stages = len(dim_mult)
    for i in range(n_stages):
        stage = []
        for _ in range(num_res_blocks + 1):
            stage.append("res")
            if scale in attn_scales:
                stage.append("attn")
        if i != n_stages - 1:
            stage.append("up")
            scale *= 2.0
        stages.append(stage)
    return stages


def _map_residual_block(comfy_prefix, diffusers_prefix, mapping):
    for comfy_sub, diffusers_sub in RESIDUAL_SUBKEYS.items():
        mapping[f"{comfy_prefix}.{comfy_sub}"] = f"{diffusers_prefix}.{diffusers_sub}"
    mapping[f"{comfy_prefix}.shortcut"] = f"{diffusers_prefix}.conv_shortcut"


def _map_attention_block(comfy_prefix, diffusers_prefix, mapping):
    mapping[f"{comfy_prefix}.norm"] = f"{diffusers_prefix}.norm"
    mapping[f"{comfy_prefix}.to_qkv"] = f"{diffusers_prefix}.to_qkv"
    mapping[f"{comfy_prefix}.proj"] = f"{diffusers_prefix}.proj"


def _map_resample(comfy_prefix, diffusers_prefix, mapping):
    mapping[f"{comfy_prefix}.resample"] = f"{diffusers_prefix}.resample"
    mapping[f"{comfy_prefix}.time_conv"] = f"{diffusers_prefix}.time_conv"


def _map_middle(comfy_prefix, diffusers_prefix, mapping):
    # comfy: Sequential[ResidualBlock, AttentionBlock, ResidualBlock]
    # diffusers: QwenImageMidBlock(resnets=[ResidualBlock, ResidualBlock], attentions=[AttentionBlock])
    _map_residual_block(f"{comfy_prefix}.0", f"{diffusers_prefix}.resnets.0", mapping)
    _map_attention_block(f"{comfy_prefix}.1", f"{diffusers_prefix}.attentions.0", mapping)
    _map_residual_block(f"{comfy_prefix}.2", f"{diffusers_prefix}.resnets.1", mapping)


def build_prefix_map(dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=()):
    """Build a dict of module-path prefixes: comfy WanVAE -> diffusers AutoencoderKLQwenImage."""
    mapping = {}

    # top level
    mapping["conv1"] = "quant_conv"
    mapping["conv2"] = "post_quant_conv"

    # --- encoder ---
    mapping["encoder.conv1"] = "encoder.conv_in"
    enc_layout = _flat_encoder_layout(dim_mult, num_res_blocks, attn_scales)
    res_i = attn_i = 0
    for flat_idx, kind in enumerate(enc_layout):
        comfy_prefix = f"encoder.downsamples.{flat_idx}"
        diffusers_prefix = f"encoder.down_blocks.{flat_idx}"
        if kind == "res":
            _map_residual_block(comfy_prefix, diffusers_prefix, mapping)
            res_i += 1
        elif kind == "attn":
            _map_attention_block(comfy_prefix, diffusers_prefix, mapping)
            attn_i += 1
        elif kind == "down":
            _map_resample(comfy_prefix, diffusers_prefix, mapping)
    _map_middle("encoder.middle", "encoder.mid_block", mapping)
    mapping["encoder.head.0"] = "encoder.norm_out"
    mapping["encoder.head.2"] = "encoder.conv_out"

    # --- decoder ---
    mapping["decoder.conv1"] = "decoder.conv_in"
    _map_middle("decoder.middle", "decoder.mid_block", mapping)
    dec_stages = _flat_decoder_stage_layout(dim_mult, num_res_blocks, attn_scales)
    flat_idx = 0
    for stage_idx, stage in enumerate(dec_stages):
        res_j = 0
        for kind in stage:
            comfy_prefix = f"decoder.upsamples.{flat_idx}"
            if kind == "res":
                diffusers_prefix = f"decoder.up_blocks.{stage_idx}.resnets.{res_j}"
                _map_residual_block(comfy_prefix, diffusers_prefix, mapping)
                res_j += 1
            elif kind == "attn":
                # diffusers QwenImageUpBlock has no attention support; only reachable if
                # attn_scales is non-empty, which real Qwen-Image-Edit VAE configs don't use.
                raise NotImplementedError(
                    "VAE config has attn_scales set in the decoder up-path; this converter "
                    "was only validated against the real Qwen-Image-Edit config (attn_scales=[])."
                )
            elif kind == "up":
                diffusers_prefix = f"decoder.up_blocks.{stage_idx}.upsamplers.0"
                _map_resample(comfy_prefix, diffusers_prefix, mapping)
            flat_idx += 1
    mapping["decoder.head.0"] = "decoder.norm_out"
    mapping["decoder.head.2"] = "decoder.conv_out"

    return mapping


def convert_state_dict(comfy_sd, dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=()):
    """Convert a comfy WanVAE state dict into diffusers AutoencoderKLQwenImage key layout."""
    prefix_map = build_prefix_map(dim_mult, num_res_blocks, attn_scales)
    # longest-prefix-first so e.g. "encoder.downsamples.10.residual.0" doesn't get
    # accidentally matched against a shorter unrelated prefix
    sorted_prefixes = sorted(prefix_map.keys(), key=len, reverse=True)

    out = {}
    unmatched = []
    for key, tensor in comfy_sd.items():
        matched = False
        for comfy_prefix in sorted_prefixes:
            if key == comfy_prefix or key.startswith(comfy_prefix + "."):
                suffix = key[len(comfy_prefix):]
                out[prefix_map[comfy_prefix] + suffix] = tensor
                matched = True
                break
        if not matched:
            unmatched.append(key)

    if unmatched:
        raise RuntimeError(
            f"vae_convert: {len(unmatched)} comfy VAE keys had no mapping, e.g. {unmatched[:5]}"
        )
    return out


def convert_and_validate(comfy_sd, target_model, dim_mult=(1, 2, 4, 4), num_res_blocks=2, attn_scales=()):
    """Convert and check the result exactly matches target_model's own state_dict keys/shapes."""
    converted = convert_state_dict(comfy_sd, dim_mult, num_res_blocks, attn_scales)
    target_sd = target_model.state_dict()

    converted_keys = set(converted.keys())
    target_keys = set(target_sd.keys())
    missing = target_keys - converted_keys
    extra = converted_keys - target_keys
    if missing or extra:
        raise RuntimeError(
            "vae_convert: converted VAE state dict doesn't match diffusers model exactly. "
            f"missing={sorted(missing)[:5]} extra={sorted(extra)[:5]}"
        )

    bad_shapes = [
        (k, tuple(converted[k].shape), tuple(target_sd[k].shape))
        for k in converted_keys
        if tuple(converted[k].shape) != tuple(target_sd[k].shape)
    ]
    if bad_shapes:
        raise RuntimeError(f"vae_convert: shape mismatches after conversion, e.g. {bad_shapes[:5]}")

    return converted