import torch
import torch.nn as nn
import torch.nn.functional as F
import sys
import os
import gdown
import inspect
import textwrap
import math

# --- 0. FORCE CPU PATCH ---
# Patches T5 to avoid "No NVIDIA driver" error
t5_path = "/app/Wan2.1/wan/modules/t5.py"
if os.path.exists(t5_path):
    print("--- PATCHING T5.PY ---")
    with open(t5_path, 'r') as f:
        content = f.read()
    new_content = content.replace("device=torch.cuda.current_device(),", "device='cpu',")
    with open(t5_path, 'w') as f:
        f.write(new_content)

sys.path.append("/app/Wan2.1")

# --- 1. DOWNLOAD WEIGHTS ---
print("--- DOWNLOADING MODEL ---")
file_id = '1pf-fhkWayK3_sv9bujvT2Od04ki2Wh4y'
url = f'https://drive.google.com/uc?id={file_id}'
ckpt_path = "/app/wan13bNSFWFull_v20I2v.safetensors"

if not os.path.exists(ckpt_path):
    gdown.download(url, ckpt_path, quiet=False)

# --- 2. CRITICAL PATCHES ---

# A. RECURSION-PROOF LINEAR FIX
_orig_linear = F.linear  # <--- SAVES THE ORIGINAL FUNCTION

def safe_linear(input, weight, bias=None):
    # Force mixed-precision to match (FP32 input + FP16 weight = FP16 input)
    if input.dtype == torch.float32 and weight.dtype == torch.float16:
        input = input.half()
    elif input.dtype == torch.float16 and weight.dtype == torch.float32:
        weight = weight.half()
        if bias is not None: bias = bias.half()
    return _orig_linear(input, weight, bias) # <--- CALLS SAVED FUNCTION

# Apply the patch
torch.nn.functional.linear = safe_linear
torch.nn.modules.linear.F.linear = safe_linear

# B. ROPE FIX (CPU Compatible)
def manual_rope_apply_real(x, grid_sizes, freqs):
    n, l, h, d = x.shape
    if freqs.dim() == 2:
        l_freqs, d_half = freqs.shape
        if l_freqs != l:
            freqs_in = freqs.transpose(0, 1).unsqueeze(0).float()
            freqs_out = F.interpolate(freqs_in, size=l, mode='linear', align_corners=False)
            freqs = freqs_out.squeeze(0).transpose(0, 1)
        cos = torch.cos(freqs)
        sin = torch.sin(freqs)
    elif freqs.dim() == 3:
        l_freqs = freqs.shape[0]
        if l_freqs != l:
             f_flat = freqs.flatten(1)
             f_in = f_flat.transpose(0, 1).unsqueeze(0).float()
             f_out = F.interpolate(f_in, size=l, mode='linear', align_corners=False)
             freqs = f_out.squeeze(0).transpose(0, 1).view(l, -1, 2)
        cos = freqs[..., 0]
        sin = freqs[..., 1]
    else:
        cos = freqs[..., 0]
        sin = freqs[..., 1]
    x_pairs = x.float().view(n, l, h, d // 2, 2)
    r = x_pairs[..., 0]
    i = x_pairs[..., 1]
    cos = cos.unsqueeze(0).unsqueeze(2)
    sin = sin.unsqueeze(0).unsqueeze(2)
    out_r = r * cos - i * sin
    out_i = r * sin + i * cos
    return torch.stack([out_r, out_i], dim=-1).flatten(3).type_as(x)

import wan.modules.model
wan.modules.model.rope_apply = manual_rope_apply_real

# C. FLASH ATTENTION FIX
def manual_attention(q, k, v, dropout_p=0.0, **kwargs):
    if q.dtype != torch.float16: q = q.half()
    if k.dtype != torch.float16: k = k.half()
    if v.dtype != torch.float16: v = v.half()
    q = q.transpose(1, 2)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)
    out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
    return out.transpose(1, 2)

import wan.modules.attention
wan.modules.attention.flash_attention = manual_attention
wan.modules.attention.FLASH_ATTN_2_AVAILABLE = False
wan.modules.model.flash_attention = manual_attention

# D. LAYERNORM FIX
class ManualWanLayerNorm(nn.LayerNorm):
    def forward(self, x):
        orig_type = x.dtype
        return F.layer_norm(x.float(), self.normalized_shape, self.weight.float(), self.bias.float(), self.eps).to(orig_type)

wan.modules.model.WanLayerNorm = ManualWanLayerNorm

# --- 3. DYNAMIC CLASS PATCHER ---
# Removes assert statements that crash the export
from wan.modules.model import WanModel, sinusoidal_embedding_1d

def patch_all_wan_classes():
    print("=== APPLYING DYNAMIC PATCHES ===")
    target_module = wan.modules.model
    classes = [obj for name, obj in inspect.getmembers(target_module) if inspect.isclass(obj) and obj.__module__ == target_module.__name__]
    
    shared_globals = {
        'torch': torch, 'math': math, 'F': F,
        'amp': torch.cuda.amp if hasattr(torch.cuda, 'amp') else None,
        'sinusoidal_embedding_1d': sinusoidal_embedding_1d,
        'rope_apply': manual_rope_apply_real,
        'nn': nn, 'flash_attention': manual_attention,
    }
    
    for cls in classes:
        shared_globals[cls.__name__] = cls

    for cls in classes:
        if not hasattr(cls, 'forward') or cls.__name__ == "ManualWanLayerNorm": continue
        try:
            src = inspect.getsource(cls.forward)
        except OSError: continue
            
        src = textwrap.dedent(src)
        lines = src.splitlines()
        new_lines = []
        modified = False
        
        for line in lines:
            indent = len(line) - len(line.lstrip())
            # FIX 1: Remove asserts that check for float32
            if "assert" in line and ("float32" in line or "dtype" in line):
                new_lines.append(" " * indent + "pass")
                modified = True
            # FIX 2: Force time embedding to Half
            elif cls.__name__ == "WanModel" and "e = self.time_embedding(t, shift)" in line:
                new_lines.append(line.replace("e = self.time_embedding(t, shift)", "e = self.time_embedding(t, shift).half()"))
                modified = True
            # FIX 3: Force QKV to Half
            elif "q, k, v = qkv_fn(x)" in line:
                new_lines.append(" " * indent + "x = x.half()")
                new_lines.append(line)
                modified = True
            else:
                new_lines.append(line)
        
        if modified:
            new_src = chr(10).join(new_lines)
            exec_scope = shared_globals.copy()
            exec_scope[cls.__name__] = cls 
            try:
                exec(new_src, exec_scope)
                cls.forward = exec_scope['forward']
                print(f"   [Patcher] SUCCESS: {cls.__name__} patched.")
            except Exception as e:
                print(f"   [Patcher] FAILED {cls.__name__}: {e}")

# --- 4. EXPORT LOGIC (PART 1) ---
class HeadIdentity(nn.Module):
    def forward(self, x, e=None): return x

def convert():
    print("--- STARTING CONVERSION: PART 1 ---")
    patch_all_wan_classes()
    
    cfg = {
        "t5_model": "umt5_xxl", "t5_dtype": torch.float16, "text_len": 512,
        "param_dtype": torch.float16, "num_train_timesteps": 1000, "sample_fps": 16,
        "sample_neg_prompt": "", "t5_checkpoint": "", "t5_tokenizer": "",
        "vae_checkpoint": "", "vae_stride": (4, 8, 8), "patch_size": (1, 2, 2),
        "dim": 1536, "num_heads": 12, "num_layers": 30, "window_size": (-1, -1),
        "qk_norm": True, "cross_attn_norm": True, "eps": 1e-6, "rope_max_len": 4096,
        "base_dim": 256, "model_type": "i2v"
    }

    with torch.no_grad():
        model = WanModel.from_config(cfg)
        model.eval().to("cpu")
        model.half()
        
        from safetensors.torch import load_file
        print(f"Loading weights from {ckpt_path}...")
        sd = load_file(ckpt_path)
        model.load_state_dict(sd, strict=False)

        old = model.patch_embedding
        model.patch_embedding = torch.nn.Conv3d(72, old.out_channels, old.kernel_size, old.stride, old.padding, bias=(old.bias is not None)).half()
        
        # EXPORT FIRST 15 BLOCKS
        all_blocks_ref = list(model.blocks)
        model.blocks = nn.ModuleList()
        model.blocks.extend(all_blocks_ref[:15]) 
        
        model.head = HeadIdentity()
        model.unpatchify = lambda x, grid_sizes: x
        
        # Dummy inputs
        dtype = torch.float16
        args = (
            [torch.randn(36, 1, 32, 32, dtype=dtype)], 
            torch.tensor([10.0], dtype=torch.float32), 
            torch.randn(1, 512, 4096, dtype=dtype),
            torch.tensor([4096]), 
            torch.randn(1, 257, 1280, dtype=dtype),
            [torch.randn(36, 1, 32, 32, dtype=dtype)]
        )

        output_file = "/app/output/wan13b_part1.onnx"
        print(f"Exporting to {output_file}...")
        
        torch.onnx.export(
            model, args, output_file,
            input_names=['latents', 'timestep', 'context', 'seq_len', 'clip_fea', 'y'],
            output_names=['mid_block_states'],
            opset_version=14,
            do_constant_folding=False,
            export_params=True,
            keep_initializers_as_inputs=True
        )
        print("--- SUCCESS ---")

if __name__ == "__main__":
    convert()