"""GGUF v3 writer and q4 tensor decoder reused from the v2 release."""
import argparse
import hashlib
import json
import os
import struct
import sys

import numpy as np

MAGIC = b"GGUF"
VERSION = 3
ALIGNMENT = 32

# GGUF metadata value types
T_UINT8, T_INT8, T_UINT16, T_INT16 = 0, 1, 2, 3
T_UINT32, T_INT32, T_FLOAT32, T_BOOL = 4, 5, 6, 7
T_STRING, T_ARRAY, T_UINT64 = 8, 9, 10
# GGML tensor types
GGML_F16, GGML_I8 = 1, 24


# ---------------------------------------------------------------------------- q4 decoding (mirrors bhc/xlit_q.py)
def read_q4(path):
    blob = open(path, "rb").read()
    n = int.from_bytes(blob[:4], "little")
    header = blob[4:4 + n]
    meta = json.loads(header)
    return blob, header, meta, blob[4 + n:]


def dequantise(meta, data):
    """Dequantise every tensor of an xlit.q4 payload to float32 (same arithmetic as bhc/xlit_q.py)."""
    bits_default, G = int(meta["bits"]), int(meta["G"])
    out = []
    for k, shape, kind, off, ln in meta["tensors"]:
        cnt = int(np.prod(shape))
        if kind.startswith("q"):
            bits = int(kind[1:]) if len(kind) > 1 else bits_default
            if bits == 4:
                raw = np.frombuffer(data, np.uint8, ln, off)
                q = np.stack([raw >> 4, raw & 15], 1).reshape(-1)[:cnt].astype(np.int16) - 8
            else:
                q = np.frombuffer(data, np.int8, ln, off).astype(np.int16)
            r, c = shape
            ng = (c + G - 1) // G
            sc = np.frombuffer(data, np.float16, r * ng, off + ln).astype(np.float32).reshape(r, ng)
            w = q.reshape(r, c).astype(np.float32) * np.repeat(sc, G, 1)[:, :c]
        else:
            w = np.frombuffer(data, np.float16, cnt, off).astype(np.float32).reshape(shape).copy()
        out.append((k, list(shape), w))
    return out


# ---------------------------------------------------------------------------- GGUF writer (spec-conformant, v3)
class Writer:
    def __init__(self):
        self.kv = []

    # -- typed metadata helpers
    def add(self, key, vtype, value):
        self.kv.append((key, vtype, value))
        return self

    def string(self, k, v): return self.add(k, T_STRING, v)
    def uint32(self, k, v): return self.add(k, T_UINT32, int(v))
    def uint64(self, k, v): return self.add(k, T_UINT64, int(v))
    def float32(self, k, v): return self.add(k, T_FLOAT32, float(v))
    def boolean(self, k, v): return self.add(k, T_BOOL, bool(v))
    def strings(self, k, v): return self.add(k, T_ARRAY, (T_STRING, list(v)))

    @staticmethod
    def _kv_bytes(key, vtype, value):
        out = bytearray()

        def s(x):
            b = x.encode("utf-8")
            return struct.pack("<Q", len(b)) + b

        out += s(key)
        out += struct.pack("<I", vtype)
        if vtype == T_STRING:
            out += s(value)
        elif vtype == T_ARRAY:
            sub, vals = value
            out += struct.pack("<I", sub) + struct.pack("<Q", len(vals))
            for v in vals:
                if sub == T_STRING:
                    out += s(v)
                else:
                    raise ValueError("unsupported array element type %d" % sub)
        elif vtype == T_BOOL:
            out += struct.pack("<?", value)
        elif vtype == T_UINT32:
            out += struct.pack("<I", value)
        elif vtype == T_UINT64:
            out += struct.pack("<Q", value)
        elif vtype == T_FLOAT32:
            out += struct.pack("<f", value)
        else:
            raise ValueError("unsupported metadata type %d" % vtype)
        return bytes(out)

    def write(self, path, tensors):
        """tensors: list of (name, np.ndarray, ggml_type_int)."""
        # tensor table: dims are stored fastest-varying first
        info = bytearray()
        payload = bytearray()
        for name, arr, gtype in tensors:
            raw = arr.tobytes(order="C")
            payload += raw
            payload += b"\0" * ((-len(payload)) % ALIGNMENT)
        offsets = []
        off = 0
        for name, arr, gtype in tensors:
            offsets.append(off)
            nbytes = arr.size * arr.dtype.itemsize
            off += nbytes + ((-nbytes) % ALIGNMENT)
        for (name, arr, gtype), o in zip(tensors, offsets):
            b = name.encode("utf-8")
            info += struct.pack("<Q", len(b)) + b
            dims = list(arr.shape)
            info += struct.pack("<I", len(dims))
            for d in reversed(dims):
                info += struct.pack("<Q", int(d))
            info += struct.pack("<I", gtype)
            info += struct.pack("<Q", o)

        kv = bytearray()
        for key, vtype, value in self.kv:
            kv += self._kv_bytes(key, vtype, value)

        head = MAGIC + struct.pack("<I", VERSION) + struct.pack("<Q", len(tensors)) + struct.pack("<Q", len(self.kv))
        body = bytes(head) + bytes(kv) + bytes(info)
        body += b"\0" * ((-len(body)) % ALIGNMENT)
        with open(path, "wb") as f:
            f.write(body)
            f.write(bytes(payload))
        return len(body) + len(payload), len(body)