"""Minimal, dependency-light reader for the GGUF files shipped with AksharaLP-MultiIndic-8M.

GGUF here is used as what it is: a portable, single-file tensor container -- a header, typed
key/value metadata, a tensor table, and aligned little-endian tensor data. It is NOT a
llama.cpp-loadable model: AksharaLP is a custom character-level encoder-decoder transformer with a
bespoke decoding stack (lexicon-trie candidate beam, hand-built candidate features, gates, MLP
reranker), and llama.cpp has no such architecture. This reader is the decoder-side tool: it returns
the weights our reference implementation (python/bhc/xlit_*.py) expects.

Only numpy is required. No torch, no llama.cpp, no gguf/gguf-py.

    python gguf_read.py aksharalp-1m-v2-f16.gguf          # summary
    python gguf_read.py --json aksharalp-1m-v2-q4.gguf     # metadata + tensor table

    import gguf_read
    w = gguf_read.load("aksharalp-1m-v2-f16.gguf")          # -> Weights(cfg, state_dict, metadata)
    w.state_dict["semb.weight"].shape                      # (30, 128) float32

Two container layouts are understood, both GGUF v3:

  * aksharalp-1m-v2-q4.gguf   one I8 tensor, ``aksharalp.q4.blob``: the byte-exact xlit.q4 file
                              (4-byte header length, JSON tensor table, then the payload). Decoded
                              here exactly as bhc/xlit_q.py:load_q does.
  * aksharalp-1m-v2-f16.gguf  every weight under its state_dict name at F16.
"""
import json
import struct
import sys

import numpy as np

MAGIC = b"GGUF"
SUPPORTED_VERSIONS = (3,)

# GGUF metadata value types
GGUF_UINT8, GGUF_INT8, GGUF_UINT16, GGUF_INT16 = 0, 1, 2, 3
GGUF_UINT32, GGUF_INT32, GGUF_FLOAT32, GGUF_BOOL = 4, 5, 6, 7
GGUF_STRING, GGUF_ARRAY, GGUF_UINT64, GGUF_INT64, GGUF_FLOAT64 = 8, 9, 10, 11, 12

# GGML tensor types we can read: F32, F16, I8 (the three this model uses)
_GGML_F32, _GGML_F16, _GGML_I8 = 0, 1, 24
_TENSOR_DTYPES = {_GGML_F32: np.dtype("<f4"), _GGML_F16: np.dtype("<f2"), _GGML_I8: np.dtype("i1")}
_TENSOR_NAMES = {_GGML_F32: "F32", _GGML_F16: "F16", _GGML_I8: "I8"}
_KV_FMT = {GGUF_UINT8: "<B", GGUF_INT8: "<b", GGUF_UINT16: "<H", GGUF_INT16: "<h",
           GGUF_UINT32: "<I", GGUF_INT32: "<i", GGUF_FLOAT32: "<f", GGUF_BOOL: "?",
           GGUF_UINT64: "<Q", GGUF_INT64: "<q", GGUF_FLOAT64: "<d"}


class GGUFError(ValueError):
    """Raised for anything that is not a well-formed GGUF v3 file we can read."""


class _Cursor:
    def __init__(self, buf, pos=0):
        self.buf, self.pos = buf, pos

    def take(self, n):
        if n < 0 or self.pos + n > len(self.buf):
            raise GGUFError("truncated file: wanted %d bytes at %d, file is %d" % (n, self.pos, len(self.buf)))
        b = self.buf[self.pos:self.pos + n]
        self.pos += n
        return b

    def unpack(self, fmt):
        n = struct.calcsize(fmt)
        return struct.unpack(fmt, self.take(n))[0]

    def string(self):
        """GGUF strings: u64 length + UTF-8 bytes (not NUL-terminated)."""
        n = self.unpack("Q")
        return self.take(n).decode("utf-8")


class Tensor:
    """One tensor: name, logical shape (GGML order reversed back to row-major), dtype, raw bytes."""

    __slots__ = ("name", "shape", "type_id", "offset", "data")

    def __init__(self, name, shape, type_id, offset, data):
        self.name, self.shape, self.type_id, self.offset, self.data = name, shape, type_id, offset, data

    @property
    def dtype(self):
        return _TENSOR_DTYPES[self.type_id]

    @property
    def type_name(self):
        return _TENSOR_NAMES[self.type_id]

    def array(self):
        return np.frombuffer(self.data, dtype=self.dtype).reshape(self.shape)

    def __repr__(self):
        return "<Tensor %s %s %s>" % (self.name, self.type_name, tuple(self.shape))


class GGUF:
    """Parsed GGUF file: header fields, typed metadata, tensor table, tensor data."""

    def __init__(self, path):
        self.path = str(path)
        with open(self.path, "rb") as f:
            self.buf = f.read()
        self._parse()

    # ---------------------------------------------------------------- parsing
    def _parse(self):
        c = _Cursor(self.buf)
        self.magic = c.take(4)
        if self.magic != MAGIC:
            raise GGUFError("%s: not a GGUF file (magic %r)" % (self.path, self.magic))
        self.version = c.unpack("I")
        if self.version not in SUPPORTED_VERSIONS:
            raise GGUFError("%s: GGUF version %d not supported (this reader knows %s)"
                            % (self.path, self.version, SUPPORTED_VERSIONS))
        self.tensor_count = c.unpack("Q")
        self.kv_count = c.unpack("Q")

        self.metadata = {}
        for _ in range(self.kv_count):
            key = c.string()
            self.metadata[key] = self._value(c)

        self.alignment = int(self.metadata.get("general.alignment", 32))
        if self.alignment <= 0 or self.alignment & (self.alignment - 1):
            raise GGUFError("%s: general.alignment %d is not a power of two" % (self.path, self.alignment))

        infos = []
        for _ in range(self.tensor_count):
            name = c.string()
            n_dims = c.unpack("I")
            if n_dims > 8:
                raise GGUFError("%s: tensor %s declares %d dims" % (self.path, name, n_dims))
            # GGML stores dims fastest-varying first; reverse to get the row-major shape.
            dims = [c.unpack("Q") for _ in range(n_dims)]
            type_id = c.unpack("I")
            offset = c.unpack("Q")
            infos.append((name, list(reversed(dims)), type_id, offset))

        pad = (-c.pos) % self.alignment
        self.data_start = c.pos + pad
        if self.data_start > len(self.buf):
            raise GGUFError("%s: tensor info block overruns the file" % self.path)

        self.tensors = {}
        self.order = []
        for name, shape, type_id, offset in infos:
            if type_id not in _TENSOR_DTYPES:
                raise GGUFError("%s: tensor %s has GGML type %d, which this reader does not decode "
                                "(only F32/F16/I8 are used by these files)" % (self.path, name, type_id))
            if name in self.tensors:
                raise GGUFError("%s: duplicate tensor name %r" % (self.path, name))
            numel = int(np.prod(shape, dtype=np.int64)) if shape else 1
            nbytes = numel * _TENSOR_DTYPES[type_id].itemsize
            start = self.data_start + offset
            end = start + nbytes
            if start < self.data_start or end > len(self.buf):
                raise GGUFError("%s: tensor %s spans [%d, %d) outside the data section [%d, %d)"
                                % (self.path, name, start, end, self.data_start, len(self.buf)))
            if offset % self.alignment:
                raise GGUFError("%s: tensor %s offset %d is not aligned to %d"
                                % (self.path, name, offset, self.alignment))
            self.tensors[name] = Tensor(name, shape, type_id, offset, self.buf[start:end])
            self.order.append(name)

    def _value(self, c, vtype=None):
        if vtype is None:
            vtype = c.unpack("I")
        if vtype == GGUF_STRING:
            return c.string()
        if vtype == GGUF_ARRAY:
            sub = c.unpack("I")
            n = c.unpack("Q")
            return [self._value(c, sub) for _ in range(n)]
        if vtype in _KV_FMT:
            return c.unpack(_KV_FMT[vtype])
        raise GGUFError("%s: unknown metadata value type %d" % (self.path, vtype))

    # ---------------------------------------------------------------- checks
    def validate(self):
        """Structural checks; returns a list of human-readable results, raises GGUFError on failure."""
        notes = []
        if len(self.tensors) != self.tensor_count:
            raise GGUFError("tensor count mismatch: header says %d, parsed %d" % (self.tensor_count, len(self.tensors)))
        notes.append("magic GGUF, version %d" % self.version)
        notes.append("%d tensors, %d metadata keys" % (self.tensor_count, len(self.metadata)))
        notes.append("alignment %d, data section at byte %d" % (self.alignment, self.data_start))
        declared = self.metadata.get("general.architecture")
        if declared is None:
            raise GGUFError("general.architecture is missing")
        for name in self.order:
            t = self.tensors[name]
            if t.offset + len(t.data) > len(self.buf) - self.data_start:
                raise GGUFError("tensor %s overruns the data section" % name)
        notes.append("all tensor offsets aligned and inside the data section")
        arch = self.metadata.get("general.architecture")
        notes.append("general.architecture = %r (custom; not a llama.cpp architecture)" % arch)
        return notes

    # ---------------------------------------------------------------- weights
    def state_dict(self):
        """Return {state_dict name: float32 ndarray} for the reference XlitTransformer."""
        cfg, sd = self._weights()
        return cfg, sd

    def _weights(self):
        if "aksharalp.nlm.blob" in self.tensors:
            return _weights_from_nlm_blob(bytes(self.tensors["aksharalp.nlm.blob"].data))
        if _BLOB_TENSOR in self.tensors:
            return _weights_from_q4_blob(bytes(self.tensors[_BLOB_TENSOR].data))
        cfg = _cfg_from_metadata(self.metadata)
        sd = {}
        for name in self.order:
            if name.startswith(RERANKER_PREFIX):
                continue            # reranker matrices are carried separately; see reranker()
            a = self.tensors[name].array()
            sd[name] = np.ascontiguousarray(a, dtype=np.float32)
        return cfg, sd

    def reranker(self):
        """{('core'|'pack', 'w1'|'b1'|'w2'|'b2'): float32 ndarray} from the F16 variant, else {}."""
        out = {}
        for name in self.order:
            if not name.startswith(RERANKER_PREFIX):
                continue
            _, part, param = name.split(".")
            out[(part, param)] = np.ascontiguousarray(self.tensors[name].array(), dtype=np.float32)
        return out


_BLOB_TENSOR = "aksharalp.q4.blob"
RERANKER_PREFIX = "reranker."


def _cfg_from_metadata(md):
    if "aksharalp.config_json" in md: return json.loads(md["aksharalp.config_json"])
    def g(key, default=None):
        v = md.get("aksharalp." + key, default)
        if v is None:
            raise GGUFError("metadata is missing aksharalp.%s" % key)
        return v
    cfg = {"d": int(g("d_model")), "h": int(g("n_heads")), "enc": int(g("encoder_layers")),
           "dec": int(g("decoder_layers")), "ff": int(g("ff_dim")), "p": float(g("dropout", 0.1))}
    return cfg


def _weights_from_q4_blob(blob):
    """Decode the carried xlit.q4 bytes exactly as python/bhc/xlit_q.py:load_q does."""
    n = int.from_bytes(blob[:4], "little")
    meta = json.loads(blob[4:4 + n])
    data = blob[4 + n:]
    bits_default, G = int(meta["bits"]), int(meta["G"])
    cfg = dict(meta["cfg"])
    sd = {}
    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
            elif bits == 8:
                q = np.frombuffer(data, np.int8, ln, off).astype(np.int16)
            else:
                raise GGUFError("unsupported q width %d for %s" % (bits, k))
            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]
            sd[k] = w
        elif kind == "f16":
            sd[k] = np.frombuffer(data, np.float16, cnt, off).astype(np.float32).reshape(shape).copy()
        else:
            raise GGUFError("unsupported tensor kind %r for %s" % (kind, k))
    return cfg, sd


class Weights:
    """Weights + config, in the shape the reference implementation consumes.

    ``state_dict`` maps state_dict names (as in python/bhc/xlit_q.py:load_q) to float32 ndarrays;
    ``cfg`` is the XlitTransformer config. To build the torch model the reference uses:

        import torch
        from bhc.xlit_model import XlitTransformer
        m = XlitTransformer(**w.cfg)
        m.load_state_dict({k: torch.from_numpy(v) for k, v in w.state_dict.items()})
        m.eval()
    """

    def __init__(self, cfg, state_dict, metadata=None, source=None, reranker=None):
        self.cfg, self.state_dict, self.metadata, self.source = cfg, state_dict, metadata or {}, source
        self.reranker = reranker or {}

    @property
    def n_params(self):
        return int(sum(v.size for v in self.state_dict.values()))

    def __repr__(self):
        return "<Weights %s: %d tensors, %d params>" % (self.source, len(self.state_dict), self.n_params)


def load(path):
    """Parse either GGUF variant and return Weights (cfg + state_dict of float32 arrays)."""
    g = GGUF(path)
    cfg, sd = g.state_dict()
    return Weights(cfg, sd, g.metadata, g.path, g.reranker())


def tensors(path):
    """Yield (name, shape, type_name, float32 array) for every tensor in the file."""
    g = GGUF(path)
    for name in g.order:
        t = g.tensors[name]
        yield name, tuple(t.shape), t.type_name, t.array()


def q4_blob(path):
    """Bytes of the carried xlit.q4 file from the compact GGUF (byte-identical to the original)."""
    g = GGUF(path)
    if _BLOB_TENSOR not in g.tensors:
        raise GGUFError("%s carries no %s tensor (that is the F16 variant)" % (path, _BLOB_TENSOR))
    return bytes(g.tensors[_BLOB_TENSOR].data)


def main(argv):
    args = [a for a in argv[1:] if not a.startswith("-")]
    flags = {a for a in argv[1:] if a.startswith("-")}
    if not args:
        print(__doc__)
        return 2
    g = GGUF(args[0])
    if "--json" in flags:
        print(json.dumps({"version": g.version, "tensors": g.tensor_count,
                          "metadata": g.metadata,
                          "tensor_table": [{"name": n, "shape": g.tensors[n].shape,
                                            "type": g.tensors[n].type_name,
                                            "offset": g.tensors[n].offset} for n in g.order]},
                         indent=1, ensure_ascii=False))
        return 0
    for line in g.validate():
        print(line)
    print("general.name = %r" % g.metadata.get("general.name"))
    print("-- tensors --")
    for n in g.order:
        t = g.tensors[n]
        print("  %-28s %-4s %-18s %d B" % (n, t.type_name, tuple(t.shape), len(t.data)))
    w = load(args[0])
    print("-- weights --")
    print("cfg = %s" % (w.cfg,))
    print("%d tensors, %d parameters" % (len(w.state_dict), w.n_params))
    return 0





def _weights_from_nlm_blob(blob):
    n=int.from_bytes(blob[:4],"little");md=json.loads(blob[4:4+n]);data=blob[4+n:];W={};G=md["G"]
    for k,shape,kind,off,ln in md["tensors"]:
        count=int(np.prod(shape))
        if kind in ("f16","f32"):
            W[k]=np.frombuffer(data,np.float16 if kind=="f16" else np.float32,count,off).astype(np.float32).reshape(shape).copy()
        else:
            r,c=shape;ng=(c+G-1)//G
            if kind=="q8":q=np.frombuffer(data,np.int8,r*c,off).astype(np.float32).reshape(r,c);end=off+r*c
            else:
                nb=(r*c+1)//2;u=np.frombuffer(data,np.uint8,nb,off);q=(np.stack([u>>4,u&15],1).reshape(-1)[:r*c].astype(np.int16)-8).astype(np.float32).reshape(r,c);end=off+nb
            end+=(-end)%4;sc=np.frombuffer(data,np.float16,r*ng,end).astype(np.float32).reshape(r,ng);W[k]=q*np.repeat(sc,G,1)[:,:c]
    return md["cfg"],W

def load_word_model(path):
    import torch
    from bhc.xlit_model import XlitTransformer, set_vocab
    w=load(path);set_vocab(w.cfg.get("vocab","hi"));m=XlitTransformer(**w.cfg);m.load_state_dict({k:torch.from_numpy(v.copy()) for k,v in w.state_dict.items()});return m.eval()

def load_sentence_lm(path):
    from bhc.nlm import NLM
    w=load(path);vocab=json.loads(w.metadata["aksharalp.tokenizer_json"])
    return NLM(w.cfg,w.state_dict,vocab,w.cfg.get("norm","self"))

if __name__ == "__main__":
    sys.exit(main(sys.argv))