"""Serialization of caches (safetensors) with a compatibility manifest; the loader rejects incompatible manifests
before any model execution. Independent of mlx_lm's save_prompt_cache so that the schema is under WP30 control."""
from __future__ import annotations
import json, os, gzip, tempfile
import mlx.core as mx
from mlx_lm.models.cache import KVCache, QuantizedKVCache
from .sparse import SparseKVCache, SparseQuantizedKVCache, _arrays_of
from ..util import sha256_text, utc_now

SCHEMA_VERSION = "wp30-cache-2"


class IncompatibleCache(Exception):
    pass


def _layer_desc(i, c):
    k, v = c.keys_and_values()
    if hasattr(c, "bits"):
        return {"layer": i, "cls": type(c).__name__, "k_shapes": [list(x.shape) for x in k], "v_shapes": [list(x.shape) for x in v],
                "dtypes": [str(x.dtype) for x in k], "bits": c.bits, "group_size": c.group_size,
                "stored_positions": int(k[0].shape[2]), "logical_offset": int(c.offset)}
    return {"layer": i, "cls": type(c).__name__, "k_shape": list(k.shape), "v_shape": list(v.shape), "dtype": str(k.dtype),
            "stored_positions": int(k.shape[2]), "logical_offset": int(c.offset)}


def build_manifest(model_info, prefix_token_ids: list[int], cache: list, document_id: str, method_id: str,
                   compression_config: dict, extra: dict | None = None) -> dict:
    m = {
        "schema_version": SCHEMA_VERSION, "created_utc": utc_now(),
        "model_repo_id": model_info.repo_id, "model_revision": model_info.revision, "config_hash": model_info.config_hash,
        "weight_files_sha256": model_info.weight_files_sha256, "tokenizer_hash": model_info.tokenizer_hash,
        "chat_template_hash": model_info.chat_template_hash,
        "prefix_token_hash": sha256_text(",".join(map(str, prefix_token_ids))), "prefix_token_count": len(prefix_token_ids),
        "num_layers": len(cache), "num_kv_heads": model_info.num_kv_heads, "head_dim": model_info.head_dim,
        "rope_theta": model_info.rope_theta, "rope_scaling": model_info.rope_scaling, "sliding_window": model_info.sliding_window,
        "positions": "keys stored post-RoPE at original positions; logical_offset = original processed-token count",
        "causal_mask_semantics": "retained positions fully visible to new tokens; new tokens causal among themselves",
        "mlx_version": model_info.mlx_version, "mlx_lm_version": model_info.mlx_lm_version,
        "document_id": document_id, "method_id": method_id, "compression_config": compression_config,
        "layers": [_layer_desc(i, c) for i, c in enumerate(cache)],
    }
    if extra:
        m.update(extra)
    return m


def _export(cache: list) -> tuple[dict, list]:
    arrays, metas = {}, []
    for i, c in enumerate(cache):
        k, v = c.keys_and_values()
        if hasattr(c, "bits"):
            for j, x in enumerate(k): arrays[f"L{i}.k{j}"] = x
            for j, x in enumerate(v): arrays[f"L{i}.v{j}"] = x
            metas.append({"cls": type(c).__name__, "offset": int(c.offset), "group_size": int(c.group_size), "bits": int(c.bits),
                          "n_kept": int(k[0].shape[2])})
        else:
            arrays[f"L{i}.k"] = k; arrays[f"L{i}.v"] = v
            metas.append({"cls": type(c).__name__, "offset": int(c.offset), "n_kept": int(k.shape[2])})
    return arrays, metas


def serialize_cache(cache: list, manifest: dict) -> bytes:
    arrays, metas = _export(cache)
    md = {"wp30_manifest": json.dumps(manifest, separators=(",", ":")), "wp30_layers": json.dumps(metas)}
    with tempfile.TemporaryDirectory() as d:
        p = os.path.join(d, "c.safetensors")
        mx.save_safetensors(p, arrays, md)
        with open(p, "rb") as f:
            return f.read()


def check_manifest(manifest: dict, expected_model_info=None, expected_prefix_hash: str | None = None):
    if manifest.get("schema_version") != SCHEMA_VERSION:
        raise IncompatibleCache(f"schema {manifest.get('schema_version')} != {SCHEMA_VERSION}")
    if expected_model_info is not None:
        pairs = [("config_hash", "config_hash"), ("tokenizer_hash", "tokenizer_hash"), ("chat_template_hash", "chat_template_hash"), ("model_revision", "revision")]
        for mk, ak in pairs:
            if manifest[mk] != getattr(expected_model_info, ak):
                raise IncompatibleCache(f"{mk} mismatch: cache from {manifest['model_repo_id']}@{manifest['model_revision']}, receiver {expected_model_info.repo_id}@{expected_model_info.revision}")
        if (manifest["num_layers"], manifest["num_kv_heads"], manifest["head_dim"]) != (expected_model_info.num_layers, expected_model_info.num_kv_heads, expected_model_info.head_dim):
            raise IncompatibleCache("layer/head geometry mismatch")
    if expected_prefix_hash is not None and manifest["prefix_token_hash"] != expected_prefix_hash:
        raise IncompatibleCache("prefix token hash mismatch")


def deserialize_cache(payload: bytes, expected_model_info=None, expected_prefix_hash: str | None = None, check: bool = True) -> tuple[list, dict]:
    with tempfile.TemporaryDirectory() as d:
        p = os.path.join(d, "c.safetensors")
        with open(p, "wb") as f:
            f.write(payload)
        arrays, md = mx.load(p, return_metadata=True)
        mx.eval(list(arrays.values()))
    manifest = json.loads(md["wp30_manifest"]); metas = json.loads(md["wp30_layers"])
    if check:
        check_manifest(manifest, expected_model_info, expected_prefix_hash)
    cache = []
    for i, m in enumerate(metas):
        cls = m["cls"]
        if cls in ("SparseQuantizedKVCache", "QuantizedKVCache"):
            c = SparseQuantizedKVCache(group_size=m["group_size"], bits=m["bits"])
            c.keys = tuple(arrays[f"L{i}.k{j}"] for j in range(3)); c.values = tuple(arrays[f"L{i}.v{j}"] for j in range(3))
        else:
            c = SparseKVCache(); c.keys, c.values = arrays[f"L{i}.k"], arrays[f"L{i}.v"]
        c.offset, c.n_kept = m["offset"], m["n_kept"]
        cache.append(c)
    return cache, manifest


def codec_compress(b: bytes, codec: str, level: int | None = None) -> tuple[bytes, float]:
    import time as _t
    t0 = _t.perf_counter()
    if codec == "none":
        out = b
    elif codec == "gzip":
        out = gzip.compress(b, compresslevel=6 if level is None else level)
    elif codec == "zstd":
        import zstandard as zstd
        out = zstd.ZstdCompressor(level=3 if level is None else level).compress(b)
    else:
        raise ValueError(codec)
    return out, (_t.perf_counter() - t0) * 1000.0


def codec_decompress(b: bytes, codec: str) -> tuple[bytes, float]:
    import time as _t
    t0 = _t.perf_counter()
    if codec == "none":
        out = b
    elif codec == "gzip":
        out = gzip.decompress(b)
    elif codec == "zstd":
        import zstandard as zstd
        out = zstd.ZstdDecompressor().decompress(b)
    else:
        raise ValueError(codec)
    return out, (_t.perf_counter() - t0) * 1000.0