from __future__ import annotations

import re

_BLOCK_RE = re.compile(r"^(blocks\.\d+)\.(?!block\.)")


def _insert_block_wrapper(key: str) -> str:
    if key.startswith("blocks.") and ".block." not in key:
        return _BLOCK_RE.sub(r"\1.block.", key, count=1)
    return key


def _remove_block_wrapper(key: str) -> str:
    if key.startswith("blocks."):
        return key.replace(".block.", ".", 1)
    return key


def diffusers_to_official_key(key: str) -> str:
    """Map WanAnimate2 Diffusers/refactor keyspace to Wan-Video/Wan-Animate-2 keyspace.

    OrbitQuant Project-A was repaired into the current Diffusers WanAnimate2 keyspace:
      blocks.0.self_attn.to_q.weight
    The official Wan-Animate-2 repo uses:
      blocks.0.block.self_attn.q.weight
    """
    k = key
    if k.startswith("model.diffusion_model."):
        k = k[len("model.diffusion_model.") :]

    # Only attention projection names differ. FFN names stay ffn.0 / ffn.2.
    replacements = (
        (".self_attn.to_out.0.", ".self_attn.o."),
        (".cross_attn.to_out.0.", ".cross_attn.o."),
        (".self_attn.to_q.", ".self_attn.q."),
        (".self_attn.to_k.", ".self_attn.k."),
        (".self_attn.to_v.", ".self_attn.v."),
        (".cross_attn.to_q.", ".cross_attn.q."),
        (".cross_attn.to_k.", ".cross_attn.k."),
        (".cross_attn.to_v.", ".cross_attn.v."),
        (".cross_attn.add_k_proj.", ".cross_attn.k_img."),
        (".cross_attn.add_v_proj.", ".cross_attn.v_img."),
        (".cross_attn.norm_added_k.", ".cross_attn.norm_k_img."),
    )
    for old, new in replacements:
        k = k.replace(old, new)
    return _insert_block_wrapper(k)


def official_to_diffusers_key(key: str) -> str:
    """Inverse of diffusers_to_official_key for Wan-Animate-2 transformer keys."""
    k = _remove_block_wrapper(key)
    replacements = (
        (".self_attn.o.", ".self_attn.to_out.0."),
        (".cross_attn.o.", ".cross_attn.to_out.0."),
        (".self_attn.q.", ".self_attn.to_q."),
        (".self_attn.k.", ".self_attn.to_k."),
        (".self_attn.v.", ".self_attn.to_v."),
        (".cross_attn.q.", ".cross_attn.to_q."),
        (".cross_attn.k.", ".cross_attn.to_k."),
        (".cross_attn.v.", ".cross_attn.to_v."),
        (".cross_attn.k_img.", ".cross_attn.add_k_proj."),
        (".cross_attn.v_img.", ".cross_attn.add_v_proj."),
        (".cross_attn.norm_k_img.", ".cross_attn.norm_added_k."),
    )
    for old, new in replacements:
        k = k.replace(old, new)
    return k


def target_key_for_module(module_name: str, keyspace: str = "official") -> str:
    key = f"{module_name}.weight" if module_name else "weight"
    if keyspace == "official":
        return key
    if keyspace == "diffusers":
        return official_to_diffusers_key(key)
    raise ValueError(f"Unsupported keyspace: {keyspace}")