#!/usr/bin/env python3
"""Dump tensor info from qwen3_assets.gguf and extract embeddings to npy."""
import sys, os, struct
import numpy as np

GGUF_PATH = "/models/qwen3-tts/qwen3_assets.gguf"
OUT_DIR = "/tmp/embeddings"

# Try using the gguf library
try:
    from gguf import GGUFReader
    print("Using gguf library GGUFReader")
    reader = GGUFReader(GGUF_PATH)
    print(f"Tensors in {GGUF_PATH}:")
    for i, tensor in enumerate(reader.tensors):
        print(f"  [{i}] name={tensor.name}, shape={tensor.shape}, type={tensor.tensor_type}")
    
    # Name mapping: GGUF tensor name -> npy filename
    NAME_MAP = {
        "text_embd": "text_embedding_projected.npy",
        "proj.weight": "proj_weight.npy",
        "proj.bias": "proj_bias.npy",
    }
    for j in range(16):
        NAME_MAP[f"codec_embd.{j}"] = f"codec_embedding_{j}.npy"

    os.makedirs(OUT_DIR, exist_ok=True)
    
    for tensor in reader.tensors:
        name = tensor.name
        outname = NAME_MAP.get(name, name.replace("/", "_").replace(".", "_") + ".npy")
        
        # tensor.data is a numpy array (may be quantized view)
        data = tensor.data
        print(f"  Processing: {name} -> {outname}, raw shape={data.shape}, dtype={data.dtype}")
        
        # For Q8_0: block size 32, each block = 2 bytes scale + 32 bytes ints
        # The gguf library should dequantize automatically via .data
        # If dtype is already float, use as-is. Otherwise cast.
        if data.dtype in (np.float32, np.float64):
            arr = data.astype(np.float32)
        elif data.dtype == np.float16:
            arr = data.astype(np.float32)
        else:
            # Quantized — try dequantizing manually for Q8_0
            print(f"    WARNING: dtype={data.dtype}, attempting Q8_0 dequant for shape {tensor.shape}")
            target_shape = list(tensor.shape)
            # Q8_0: block_size=32, each block: 1 fp16 scale + 32 int8
            n_elements = 1
            for d in target_shape:
                n_elements *= d
            n_blocks = n_elements // 32
            raw = data.tobytes()
            # Each Q8_0 block: 2 bytes (fp16 scale) + 32 bytes (int8 quants) = 34 bytes
            block_size = 34
            if len(raw) == n_blocks * block_size:
                scales = np.zeros(n_blocks, dtype=np.float32)
                quants = np.zeros(n_elements, dtype=np.float32)
                for bi in range(n_blocks):
                    offset = bi * block_size
                    s = np.frombuffer(raw[offset:offset+2], dtype=np.float16)[0]
                    scales[bi] = float(s)
                    qs = np.frombuffer(raw[offset+2:offset+block_size], dtype=np.int8)
                    quants[bi*32:(bi+1)*32] = qs.astype(np.float32) * float(s)
                arr = quants.reshape(target_shape)
            else:
                print(f"    ERROR: Cannot dequantize, raw_bytes={len(raw)}, expected={n_blocks * block_size}")
                arr = data.astype(np.float32) if data.dtype.kind == 'f' else None
                if arr is None:
                    print(f"    SKIPPING tensor {name}")
                    continue
        
        outpath = os.path.join(OUT_DIR, outname)
        np.save(outpath, arr)
        print(f"  Saved: {outpath} shape={arr.shape} dtype={arr.dtype}")
    
    print(f"\nDone! Files in {OUT_DIR}:")
    for f in sorted(os.listdir(OUT_DIR)):
        sz = os.path.getsize(os.path.join(OUT_DIR, f))
        print(f"  {f} ({sz} bytes)")

except ImportError:
    print("gguf library not available, trying manual parse...")
    # Minimal GGUF tensor listing
    with open(GGUF_PATH, "rb") as f:
        magic = f.read(4)
        print(f"Magic: {magic}")
        version = struct.unpack("<I", f.read(4))[0]
        print(f"Version: {version}")
        n_tensors = struct.unpack("<Q", f.read(8))[0]
        n_kv = struct.unpack("<Q", f.read(8))[0]
        print(f"Tensors: {n_tensors}, KV pairs: {n_kv}")