#!/usr/bin/env python3
"""Inject MTP head from 27B donor GGUF into 40B Deckard GGUF.

Raw binary copy — no KV re-serialization. Inserts nextn_predict_layers
and MTP tensor info entries, copies everything else byte-for-byte.
"""

import argparse
import struct
import sys
import os
import io
from pathlib import Path
from gguf import GGUFReader


GGUF_MAGIC = 0x46554747
ALIGNMENT = 32

def align(x, a=ALIGNMENT):
    return (x + a - 1) & ~(a - 1)

def write_string(f, s):
    encoded = s.encode('utf-8')
    f.write(struct.pack('<Q', len(encoded)))
    f.write(encoded)

def write_tensor_info(f, name, n_dims, dims, ttype, offset):
    write_string(f, name)
    f.write(struct.pack('<I', n_dims))
    for d in dims:
        f.write(struct.pack('<Q', d))
    f.write(struct.pack('<I', ttype))
    f.write(struct.pack('<Q', offset))


def main():
    parser = argparse.ArgumentParser(description="Inject MTP head into GGUF")
    parser.add_argument('--target', required=True)
    parser.add_argument('--donor', required=True)
    parser.add_argument('--output', required=True)
    parser.add_argument('--source-layer', type=int, default=64)
    parser.add_argument('--dest-layer', type=int, default=96)
    parser.add_argument('--dry-run', action='store_true')
    args = parser.parse_args()

    target_path = Path(args.target).expanduser()
    donor_path = Path(args.donor).expanduser()
    output_path = Path(args.output).expanduser()

    if not target_path.exists():
        print(f"ERROR: Target not found: {target_path}"); sys.exit(1)
    if not donor_path.exists():
        print(f"ERROR: Donor not found: {donor_path}"); sys.exit(1)
    if output_path.exists() and not args.dry_run:
        r = input(f"Output exists: {output_path}\nOverwrite? [y/N] ").strip().lower()
        if r != 'y': sys.exit(0)

    target_size = os.path.getsize(target_path)

    print("=" * 60)
    print("MTP Head Injection (Raw Copy v2)")
    print("=" * 60)
    print(f"  Target: {target_path} ({target_size / (1024**3):.2f} GB)")
    print(f"  Donor:  {donor_path}")
    print(f"  Output: {output_path}")

    # ── Read donor ─────────────────────────────────────────────────────────
    print(f"\n[1/4] Reading donor...")
    donor_reader = GGUFReader(str(donor_path))
    prefix = f"blk.{args.source_layer}."
    mtp_tensors = []
    for tensor in donor_reader.tensors:
        if tensor.name.startswith(prefix):
            new_name = tensor.name.replace(prefix, f"blk.{args.dest_layer}.")
            mtp_tensors.append({
                'new_name': new_name,
                'shape': list(tensor.shape),
                'ttype': tensor.tensor_type.value,
                'data': bytes(tensor.data),
            })
            print(f"  {tensor.name} -> {new_name} ({tensor.data.nbytes:,} B)")
    print(f"  Total: {len(mtp_tensors)}")
    if not mtp_tensors:
        print("ERROR: No MTP tensors!"); sys.exit(1)

    # ── Analyze target structure ───────────────────────────────────────────
    print(f"\n[2/4] Analyzing target...")
    target_reader = GGUFReader(str(target_path))
    n_orig_tensors = len(target_reader.tensors)

    # Find architecture
    arch = "qwen35"
    for field in target_reader.fields.values():
        if field.name == "general.architecture":
            arch = str(bytes(field.parts[field.data[0]]), 'utf-8')
            break
    nextn_key = f"{arch}.nextn_predict_layers"
    has_nextn = any('nextn_predict_layers' in f.name for f in target_reader.fields.values())
    print(f"  Architecture: {arch}")
    print(f"  Tensors: {n_orig_tensors}")

    # Compute section boundaries
    # KV end = last non-GGUF field offset + its total parts size
    kv_end = 0
    for field in target_reader.fields.values():
        if field.name.startswith('GGUF.'):
            continue
        end = field.offset + sum(p.nbytes for p in field.parts)
        if end > kv_end:
            kv_end = end

    # TI size = sum of each tensor info entry
    ti_size = 0
    for t in target_reader.tensors:
        ti_size += 8 + len(t.name.encode('utf-8')) + 4 + len(t.shape) * 8 + 4 + 8

    ti_end = kv_end + ti_size
    data_start = align(ti_end)
    orig_tensor_data_size = target_size - data_start

    print(f"  KV end: {kv_end:,}")
    print(f"  TI: {kv_end:,} to {ti_end:,} ({ti_size:,} B)")
    print(f"  Data start: {data_start:,}")
    print(f"  Data size: {orig_tensor_data_size:,} B")

    # Read original header counts
    # GGUF v3 header: magic(4) + version(4) + n_tensors(8) + n_kv(8)
    with open(target_path, 'rb') as f:
        f.read(8)  # magic + version
        orig_n_tensors = struct.unpack('<Q', f.read(8))[0]
        orig_n_kv = struct.unpack('<Q', f.read(8))[0]
    print(f"  Original n_tensors={orig_n_tensors}, n_kv={orig_n_kv}")

    new_n_kv = orig_n_kv + (0 if has_nextn else 1)
    new_n_tensors = orig_n_tensors + len(mtp_tensors)

    if args.dry_run:
        print(f"\n[DRY RUN] Would produce n_kv={new_n_kv}, n_tensors={new_n_tensors}")
        sys.exit(0)

    # ── Build insert bytes ─────────────────────────────────────────────────
    print(f"\n[3/4] Preparing inserts...")

    # nextn KV entry
    nextn_bytes = b''
    if not has_nextn:
        nb = io.BytesIO()
        write_string(nb, nextn_key)
        nb.write(struct.pack('<I', 4))  # UINT32
        nb.write(struct.pack('<I', 1))  # value
        nextn_bytes = nb.getvalue()
        print(f"  nextn KV: {len(nextn_bytes)} B")

    # MTP tensor info entries
    # MTP tensor data offsets are relative to data section start
    # They go after the original tensor data
    mtp_ti = io.BytesIO()
    mtp_data_base = align(orig_tensor_data_size)
    cur_off = mtp_data_base
    for mt in mtp_tensors:
        aoff = align(cur_off)
        write_tensor_info(mtp_ti, mt['new_name'], len(mt['shape']),
                          mt['shape'], mt['ttype'], aoff)
        mt['rel_offset'] = aoff  # relative to data section start
        cur_off = aoff + len(mt['data'])
    mtp_ti_bytes = mtp_ti.getvalue()
    print(f"  MTP TI: {len(mtp_ti_bytes)} B")

    # ── Write output ───────────────────────────────────────────────────────
    print(f"\n[4/4] Writing output...")

    with open(output_path, 'wb') as out:
        # Patched header
        with open(target_path, 'rb') as src:
            hdr = bytearray(src.read(24))
        struct.pack_into('<Q', hdr, 8, new_n_tensors)
        struct.pack_into('<Q', hdr, 16, new_n_kv)
        out.write(hdr)
        print(f"  Header: 24 B (n_kv={new_n_kv}, n_tensors={new_n_tensors})")

        # Original KV section [24 .. kv_end)
        with open(target_path, 'rb') as src:
            src.seek(24)
            kv_bytes = src.read(kv_end - 24)
            out.write(kv_bytes)
        print(f"  Original KV: {len(kv_bytes):,} B")

        # Insert nextn KV
        if nextn_bytes:
            out.write(nextn_bytes)
            print(f"  Inserted nextn KV")

        # Original TI section [kv_end .. kv_end + ti_size)
        with open(target_path, 'rb') as src:
            src.seek(kv_end)
            ti_bytes = src.read(ti_size)
            out.write(ti_bytes)
        print(f"  Original TI: {len(ti_bytes):,} B")

        # MTP TI entries
        out.write(mtp_ti_bytes)
        print(f"  MTP TI: {len(mtp_ti_bytes)} B")

        # Pad to alignment
        pos = out.tell()
        new_data_start = align(pos)
        if new_data_start > pos:
            out.write(b'\x00' * (new_data_start - pos))
        print(f"  New data section at: {new_data_start:,}")

        # The original tensor offsets are relative to the original data section.
        # Our new data section starts at a different absolute position.
        # BUT the offsets stored in tensor info are relative to data section start.
        # Since we copied the TI verbatim, those relative offsets are still correct
        # as long as the tensor data is at the same relative positions.
        # So we just need to copy all tensor data starting at new_data_start.

        # Copy original tensor data
        print(f"  Copying tensor data ({orig_tensor_data_size / (1024**3):.2f} GB)...")
        CHUNK = 64 * 1024 * 1024
        with open(target_path, 'rb') as src:
            src.seek(data_start)
            remaining = orig_tensor_data_size
            copied = 0
            while remaining > 0:
                chunk = src.read(min(CHUNK, remaining))
                if not chunk:
                    break
                out.write(chunk)
                remaining -= len(chunk)
                copied += len(chunk)
                gb = copied / (1024**3)
                if int(gb * 2) > int((copied - len(chunk)) / (1024**3) * 2):
                    print(f"    {gb:.1f} GB...")
        print(f"    Copied: {copied:,} B")

        # MTP tensor data
        print(f"  Writing MTP data...")
        for mt in mtp_tensors:
            target_pos = new_data_start + mt['rel_offset']
            cur = out.tell()
            if target_pos > cur:
                out.write(b'\x00' * (target_pos - cur))
            out.write(mt['data'])
            print(f"    {mt['new_name']} ({len(mt['data']):,} B)")

    final_size = os.path.getsize(output_path)
    print(f"\n  Output: {final_size:,} B ({final_size / (1024**3):.2f} GB)")

    # Verify
    print(f"\nVerifying...")
    try:
        vr = GGUFReader(str(output_path))
        found = any('nextn_predict_layers' in f.name for f in vr.fields.values())
        mc = sum(1 for t in vr.tensors if f'blk.{args.dest_layer}' in t.name)
        print(f"  nextn_predict_layers: {'found' if found else 'MISSING'}")
        print(f"  MTP tensors: {mc}")
        print(f"  Total tensors: {len(vr.tensors)}")
        print(f"  {'PASSED' if found and mc > 0 else 'FAILED'}")
    except Exception as e:
        print(f"  Python verify failed: {e}")
        print(f"  Try llama-server to test.")


if __name__ == '__main__':
    main()