from typing import Optional

from .wan_video import WanVideoPipeline, ModelConfig
from ..models.comp_attn_model import (
    CompAttnConfig,
    CompAttnMergeUnit,
    CompAttnUnit,
    patch_cross_attention,
    wrap_model_fn,
)
from .wan_video import WanVideoUnit_PromptEmbedder, WanVideoUnit_CfgMerger


def attach_comp_attn(pipe: WanVideoPipeline) -> WanVideoPipeline:
    if getattr(pipe, "_comp_attn_attached", False):
        return pipe
    prompt_idx = None
    cfg_idx = None
    for idx, unit in enumerate(pipe.units):
        if prompt_idx is None and isinstance(unit, WanVideoUnit_PromptEmbedder):
            prompt_idx = idx
        if cfg_idx is None and isinstance(unit, WanVideoUnit_CfgMerger):
            cfg_idx = idx
    if prompt_idx is not None:
        pipe.units.insert(prompt_idx + 1, CompAttnUnit())
    else:
        pipe.units.append(CompAttnUnit())
    if cfg_idx is not None:
        pipe.units.insert(cfg_idx + 1, CompAttnMergeUnit())
    else:
        pipe.units.append(CompAttnMergeUnit())
    patch_cross_attention(pipe)
    wrap_model_fn(pipe)
    pipe._comp_attn_attached = True
    return pipe


class WanVideoCompAttnPipeline:
    """Comp-Attn 增强的视频生成 Pipeline
    
    支持两种标注模式：
    
    1. 显式标记模式（推荐）：
       在 prompt 中使用 <N>subject</N> 标记，索引与 bboxes 对应
       
       ```python
       prompt = "A <0>red car</0> drives left, a <1>blue bicycle</1> rides right"
       comp_attn = CompAttnConfig(
           bboxes=[car_bboxes, bike_bboxes],  # 按标记索引 <0>, <1> 对应
       )
       ```
       
    2. 隐式搜索模式（兼容旧版）：
       提供 subjects 列表，自动在 prompt 中搜索匹配
       
       ```python
       prompt = "A red car drives left, a blue bicycle rides right"
       comp_attn = CompAttnConfig(
           subjects=["red car", "blue bicycle"],
           bboxes=[car_bboxes, bike_bboxes],
       )
       ```
    """
    
    def __init__(self, pipe: WanVideoPipeline):
        self.pipe = attach_comp_attn(pipe)

    def __getattr__(self, name):
        return getattr(self.pipe, name)

    @staticmethod
    def from_pretrained(
        torch_dtype=None,
        device="cuda",
        model_configs: list[ModelConfig] = None,
        tokenizer_config: Optional[ModelConfig] = None,
        audio_processor_config: Optional[ModelConfig] = None,
        redirect_common_files: bool = True,
        use_usp: bool = False,
        vram_limit: float = None,
    ):
        pipe = WanVideoPipeline.from_pretrained(
            torch_dtype=torch_dtype,
            device=device,
            model_configs=model_configs or [],
            tokenizer_config=tokenizer_config,
            audio_processor_config=audio_processor_config,
            redirect_common_files=redirect_common_files,
            use_usp=use_usp,
            vram_limit=vram_limit,
        )
        return WanVideoCompAttnPipeline(pipe)

    def __call__(
        self,
        prompt: str,
        negative_prompt: str = "",
        comp_attn: Optional[CompAttnConfig] = None,
        **kwargs,
    ):
        num_frames = kwargs.get("num_frames")
        if num_frames is not None:
            self.pipe._comp_attn_num_frames = num_frames
        
        self.pipe._comp_attn_config = comp_attn
        self.pipe._comp_attn_last_prompt = prompt
        self.pipe._comp_attn_last_negative_prompt = negative_prompt
        return self.pipe(prompt=prompt, negative_prompt=negative_prompt, **kwargs)