import os
import struct
from typing import BinaryIO, Dict, NamedTuple

import numpy as np
from tqdm import tqdm
from huggingface_hub import hf_hub_download

WORK_DIR = "/content/gguf_patch"
CHUNK_SIZE = 8 * 1024 * 1024

MODELS = {
    "source": ("bartowski/Qwen_Qwen3.5-9B-GGUF", "Qwen_Qwen3.5-9B-Q4_K_M.gguf"),
    "target": ("HauhauCS/Qwen3.5-9B-Uncensored-HauhauCS-Aggressive", "Qwen3.5-9B-Uncensored-HauhauCS-Aggressive-Q4_K_M.gguf"),
    "apply_to": ("Jackrong/Qwen3.5-9B-Claude-4.6-Opus-Reasoning-Distilled-v2-GGUF", "Qwen3.5-9B.Q4_K_M.gguf"),
}

OUTPUT_FILENAME = "result_patched.gguf"

GGUF_MAGIC = 0x46554747

GGUF_TYPE_FIXED_SIZE = {
    0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4,
    6: 4, 7: 1, 10: 8, 11: 8, 12: 8,
}

QUANT_BLOCK = {
    0: (4, 1), 1: (2, 1), 2: (18, 32), 3: (20, 32),
    6: (22, 32), 7: (24, 32), 8: (34, 32), 9: (40, 32),
    10: (84, 256), 11: (110, 256), 12: (144, 256), 13: (176, 256),
    14: (210, 256), 15: (292, 256), 28: (2, 1),
}


class TensorInfo(NamedTuple):
    name: str
    shape: tuple
    dtype: int
    offset: int


def _read_u32(f: BinaryIO) -> int:
    return struct.unpack("<I", f.read(4))[0]


def _read_u64(f: BinaryIO) -> int:
    return struct.unpack("<Q", f.read(8))[0]


def _read_string(f: BinaryIO) -> str:
    length = _read_u64(f)
    return f.read(length).decode("utf-8")


def _skip_metadata_value(f: BinaryIO, vtype: int) -> None:
    if vtype == 8:
        f.seek(_read_u64(f), 1)
    elif vtype == 9:
        elem_type = _read_u32(f)
        count = _read_u64(f)
        for _ in range(count):
            _skip_metadata_value(f, elem_type)
    else:
        f.seek(GGUF_TYPE_FIXED_SIZE[vtype], 1)


def _tensor_byte_size(shape: tuple, dtype: int) -> int:
    n_elements = 1
    for d in shape:
        n_elements *= d
    block_bytes, elems_per_block = QUANT_BLOCK[dtype]
    if elems_per_block == 1:
        return n_elements * block_bytes
    return (n_elements // elems_per_block) * block_bytes


def parse_gguf(filepath: str) -> tuple[dict[str, TensorInfo], int]:
    with open(filepath, "rb") as f:
        magic = _read_u32(f)
        if magic != GGUF_MAGIC:
            raise ValueError(f"Invalid GGUF magic in {filepath}: 0x{magic:08X}")

        version = _read_u32(f)
        if version not in (2, 3):
            raise ValueError(f"Unsupported GGUF version {version} in {filepath}")

        n_tensors = _read_u64(f)
        n_kv = _read_u64(f)

        # Read metadata to find alignment (default 32)
        alignment = 32
        for _ in range(n_kv):
            key = _read_string(f)
            vtype = _read_u32(f)
            if key == "general.alignment" and vtype == 4:
                alignment = _read_u32(f)
            else:
                _skip_metadata_value(f, vtype)

        # Read tensor descriptors
        tensors = {}
        for _ in range(n_tensors):
            name = _read_string(f)
            n_dims = _read_u32(f)
            shape = tuple(_read_u64(f) for _ in range(n_dims))
            dtype = _read_u32(f)
            offset = _read_u64(f)
            tensors[name] = TensorInfo(name, shape, dtype, offset)

        # Data section starts at next aligned boundary
        header_end = f.tell()
        data_start = ((header_end + alignment - 1) // alignment) * alignment

    return tensors, data_start


def fetch(repo_id: str, filename: str) -> str:
    return hf_hub_download(repo_id=repo_id, filename=filename, local_dir=WORK_DIR)


def add_difference(src_path: str, tgt_path: str, app_path: str, out_path: str) -> None:
    print("Parsing GGUF headers...")
    src_tensors, src_data = parse_gguf(src_path)
    tgt_tensors, tgt_data = parse_gguf(tgt_path)
    app_tensors, app_data = parse_gguf(app_path)

    common = set(src_tensors) & set(tgt_tensors) & set(app_tensors)
    if not common:
        raise ValueError("No common tensors across all three files")

    # Full copy of apply_to preserves magic, version, metadata, alignment, tensor table
    print("Copying apply_to as base (preserving header and metadata)...")
    app_size = os.path.getsize(app_path)
    with open(app_path, "rb") as fin, open(out_path, "wb") as fout:
        remaining = app_size
        while remaining > 0:
            n = min(CHUNK_SIZE, remaining)
            fout.write(fin.read(n))
            remaining -= n

    # Verify output header integrity
    with open(out_path, "rb") as f:
        magic = struct.unpack("<I", f.read(4))[0]
        if magic != GGUF_MAGIC:
            raise RuntimeError("Output file header corrupted after copy")

    patched = identical = skipped = 0

    with (
        open(src_path, "rb") as f_src,
        open(tgt_path, "rb") as f_tgt,
        open(app_path, "rb") as f_app,
        open(out_path, "r+b") as f_out,
    ):
        for name in tqdm(sorted(common), desc="Patching tensors"):
            si, ti, ai = src_tensors[name], tgt_tensors[name], app_tensors[name]

            if si.dtype != ti.dtype or si.dtype != ai.dtype:
                skipped += 1
                continue
            if si.shape != ti.shape or si.shape != ai.shape:
                skipped += 1
                continue

            size = _tensor_byte_size(si.shape, si.dtype)
            s_abs = src_data + si.offset
            t_abs = tgt_data + ti.offset
            a_abs = app_data + ai.offset

            # Fast check: skip if source == target for this tensor
            is_same = True
            off = 0
            while off < size:
                n = min(CHUNK_SIZE, size - off)
                f_src.seek(s_abs + off)
                f_tgt.seek(t_abs + off)
                if f_src.read(n) != f_tgt.read(n):
                    is_same = False
                    break
                off += n

            if is_same:
                identical += 1
                continue

            # Apply delta
            off = 0
            while off < size:
                n = min(CHUNK_SIZE, size - off)
                f_src.seek(s_abs + off)
                f_tgt.seek(t_abs + off)
                f_app.seek(a_abs + off)

                s = np.frombuffer(f_src.read(n), dtype=np.uint8).astype(np.int16)
                t = np.frombuffer(f_tgt.read(n), dtype=np.uint8).astype(np.int16)
                a = np.frombuffer(f_app.read(n), dtype=np.uint8).astype(np.int16)

                r = np.clip(a + (t - s), 0, 255).astype(np.uint8)

                f_out.seek(a_abs + off)
                f_out.write(r.tobytes())
                off += n

            patched += 1

    print(f"Done: {patched} patched, {identical} identical, {skipped} skipped")


def main() -> None:
    os.makedirs(WORK_DIR, exist_ok=True)

    paths = {}
    for role, (repo, fname) in MODELS.items():
        print(f"Fetching {role}: {repo}/{fname}")
        paths[role] = fetch(repo, fname)

    output = os.path.join(WORK_DIR, OUTPUT_FILENAME)
    add_difference(paths["source"], paths["target"], paths["apply_to"], output)

    # Final validation
    with open(output, "rb") as f:
        magic = struct.unpack("<I", f.read(4))[0]
        version = struct.unpack("<I", f.read(4))[0]
    print(f"Output: {output} ({os.path.getsize(output) / 1e9:.2f} GB)")
    print(f"GGUF magic: 0x{magic:08X} ({'OK' if magic == GGUF_MAGIC else 'INVALID'}), version: {version}")


if __name__ == "__main__":
    main()