#!/usr/bin/env python3
"""
export_text_enc.py — Export Qwen3-Embedding text encoder to ONNX.

The text encoder is a standard Qwen3Model (28 layers, H=1024, causal attention)
that takes BPE token IDs and produces hidden states for the condition encoder.

Usage:
    python export_text_enc.py --model-dir <path-to-Qwen3-Embedding-0.6B> --output <output.onnx>

Exports:
    text_encoder.onnx — Full 28-layer transformer
        Input:  input_ids [B, S] int64
        Output: hidden_states [B, S, 1024] fp16

    embed_lookup.bin — Raw embedding table (vocab_size * hidden_size * 2 bytes, BF16)
        Used for lyric token embedding lookup on CPU (no ONNX needed).
"""

import argparse
import os
import sys
import time
import struct
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn


class TextEncoderWrapper(nn.Module):
    """Wrapper around Qwen3Model that returns hidden_states as a flat tensor.
    
    ONNX inputs:
        input_ids:     [B, S] int64 — BPE token IDs
    
    ONNX output:
        hidden_states: [B, S, 1024] fp16 — last hidden state
    """
    
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    def forward(self, input_ids):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=None,  # causal mask generated internally
            output_hidden_states=False,
            return_dict=True,
        )
        return outputs.last_hidden_state


def load_model(model_dir: str, device: str = "cuda", dtype=torch.float32):
    """Load Qwen3-Embedding model from safetensors."""
    model_dir = Path(model_dir)
    
    # Fix Windows encoding issues
    if sys.platform == "win32":
        import io
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
        sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
    
    print(f"[export_text_enc] Loading model from {model_dir}...")
    t0 = time.time()
    
    from transformers import AutoModel, AutoConfig
    
    config = AutoConfig.from_pretrained(str(model_dir))
    # Force SDPA for ONNX export (no flash attention)
    config._attn_implementation = "sdpa"
    
    model = AutoModel.from_pretrained(
        str(model_dir),
        config=config,
        torch_dtype=dtype,
        trust_remote_code=True,
    )
    model = model.to(device)
    model.eval()
    
    t1 = time.time()
    n_params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"[export_text_enc] Model loaded in {t1-t0:.1f}s ({n_params:.0f}M params)")
    print(f"[export_text_enc] Config: {config.num_hidden_layers}L, H={config.hidden_size}, "
          f"heads={config.num_attention_heads}/{config.num_key_value_heads}")
    
    return model, config


def export_onnx(model, config, output_path: str, opset: int = 18):
    """Export the text encoder to ONNX."""
    device = next(model.parameters()).device
    dtype = next(model.parameters()).dtype
    
    wrapper = TextEncoderWrapper(model)
    wrapper.eval()
    
    # Dummy inputs for tracing
    B = 1
    S = 128  # typical sequence length
    
    dummy_input_ids = torch.randint(0, config.vocab_size, (B, S), device=device, dtype=torch.long)
    
    print(f"[export_text_enc] Tracing with shapes: input_ids={list(dummy_input_ids.shape)}")
    
    # Test forward pass
    print("[export_text_enc] Testing forward pass...")
    with torch.no_grad():
        test_out = wrapper(dummy_input_ids)
    print(f"[export_text_enc] Output shape: {list(test_out.shape)} "
          f"(expected [{B}, {S}, {config.hidden_size}])")
    
    # Export to ONNX
    print(f"[export_text_enc] Exporting to ONNX (opset {opset})...")
    t0 = time.time()
    
    torch.onnx.export(
        wrapper,
        (dummy_input_ids,),
        output_path,
        opset_version=opset,
        input_names=["input_ids"],
        output_names=["hidden_states"],
        dynamic_axes={
            "input_ids":     {0: "batch", 1: "seq_len"},
            "hidden_states": {0: "batch", 1: "seq_len"},
        },
        do_constant_folding=True,
        export_params=True,
    )
    
    t1 = time.time()
    file_size = os.path.getsize(output_path)
    print(f"[export_text_enc] Exported to {output_path}")
    print(f"[export_text_enc] File size: {file_size/1e6:.1f} MB")
    print(f"[export_text_enc] Export time: {t1-t0:.1f}s")
    
    return output_path


def export_embed_table(model, config, output_path: str):
    """Export the embedding table as a raw binary file for lyric lookup.
    
    The lyric path uses embed_tokens lookup only (no transformer layers).
    We export the table as float32 for direct CPU indexing.
    
    Format: raw float32 array [vocab_size, hidden_size]
    """
    embed_weight = model.embed_tokens.weight.detach().cpu().float().numpy()
    V, H = embed_weight.shape
    
    with open(output_path, "wb") as f:
        # Header: vocab_size (int32), hidden_size (int32)
        f.write(struct.pack("<II", V, H))
        # Raw float32 weights
        f.write(embed_weight.tobytes())
    
    file_size = os.path.getsize(output_path)
    print(f"[export_text_enc] Embedding table: [{V}, {H}] -> {output_path} ({file_size/1e6:.1f} MB)")


def export_null_cond(model_dir: str, output_path: str):
    """Export null_condition_emb from the DiT model as raw float32.
    
    This is a [2048] float32 vector used for classifier-free guidance padding.
    Read from the DiT safetensors since it lives there.
    """
    from safetensors.torch import load_file
    
    model_dir = Path(model_dir)
    st_path = model_dir / "model.safetensors"
    if not st_path.exists():
        # Try multi-shard
        for p in sorted(model_dir.glob("model-*.safetensors")):
            st = load_file(str(p))
            if "null_condition_emb" in st:
                vec = st["null_condition_emb"].detach().cpu().float().numpy()
                with open(output_path, "wb") as f:
                    f.write(struct.pack("<I", vec.shape[0]))
                    f.write(vec.tobytes())
                print(f"[export_text_enc] null_condition_emb: [{vec.shape[0]}] -> {output_path}")
                return
        print("[export_text_enc] WARNING: null_condition_emb not found")
        return
    
    st = load_file(str(st_path))
    if "null_condition_emb" not in st:
        print("[export_text_enc] WARNING: null_condition_emb not found in model.safetensors")
        return
    
    vec = st["null_condition_emb"].detach().cpu().float().numpy()
    with open(output_path, "wb") as f:
        f.write(struct.pack("<I", vec.shape[0]))
        f.write(vec.tobytes())
    print(f"[export_text_enc] null_condition_emb: [{vec.shape[0]}] -> {output_path}")


def verify_onnx(onnx_path: str, model, config):
    """Verify ONNX output matches PyTorch."""
    try:
        import onnxruntime as ort
    except ImportError:
        print("[export_text_enc] onnxruntime not installed, skipping verification")
        return
    
    device = next(model.parameters()).device
    wrapper = TextEncoderWrapper(model)
    wrapper.eval()
    
    # Test inputs
    B, S = 1, 64
    input_ids = torch.randint(0, config.vocab_size, (B, S), device=device, dtype=torch.long)
    
    # PyTorch reference
    with torch.no_grad():
        ref_out = wrapper(input_ids).cpu().float().numpy()
    
    # ONNX inference
    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
    sess = ort.InferenceSession(onnx_path, providers=providers)
    ort_out = sess.run(None, {
        "input_ids": input_ids.cpu().numpy(),
    })[0]
    
    # Compare
    max_diff = np.max(np.abs(ref_out - ort_out))
    mean_diff = np.mean(np.abs(ref_out - ort_out))
    print(f"[export_text_enc] Verification: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}")
    
    if max_diff < 0.05:
        print("[export_text_enc] PASS: ONNX output matches PyTorch (within FP16 tolerance)")
    else:
        print("[export_text_enc] WARNING: Large difference — may need investigation")


def main():
    parser = argparse.ArgumentParser(description="Export Qwen3-Embedding text encoder to ONNX")
    parser.add_argument("--model-dir", required=True,
                        help="Path to Qwen3-Embedding-0.6B directory")
    parser.add_argument("--output", default=None,
                        help="Output ONNX file (default: models/onnx/text_encoder.onnx)")
    parser.add_argument("--dit-dir", default=None,
                        help="Path to DiT model dir (for null_condition_emb export)")
    parser.add_argument("--opset", type=int, default=18,
                        help="ONNX opset version (default: 18)")
    parser.add_argument("--verify", action="store_true",
                        help="Verify ONNX output matches PyTorch")
    parser.add_argument("--device", default="cuda",
                        help="Device for model loading (default: cuda)")
    args = parser.parse_args()
    
    # Default output path
    if args.output is None:
        onnx_dir = Path(args.model_dir).parent / "onnx"
        onnx_dir.mkdir(parents=True, exist_ok=True)
        args.output = str(onnx_dir / "text_encoder.onnx")
    
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    output_dir = os.path.dirname(args.output)
    
    # Load model
    model, config = load_model(args.model_dir, device=args.device)
    
    # Export ONNX
    export_onnx(model, config, args.output, opset=args.opset)
    
    # Export embedding table for lyric lookup
    embed_path = os.path.join(output_dir, "embed_tokens.bin")
    export_embed_table(model, config, embed_path)
    
    # Export null_condition_emb if DiT dir provided
    if args.dit_dir:
        null_cond_path = os.path.join(output_dir, "null_condition_emb.bin")
        export_null_cond(args.dit_dir, null_cond_path)
    
    # Verify
    if args.verify:
        verify_onnx(args.output, model, config)
    
    print("[export_text_enc] Done!")


if __name__ == "__main__":
    main()