from __future__ import annotations

"""QDiff: compact quantized delta/checkpoint utilities.

This module provides a minimal, well-tested implementation of a
quantized residual delta suitable for toy workloads and unit tests.

Features implemented:
- 8-bit symmetric quantization of residuals
- optional top-k sparsification (keep largest magnitude residuals)
- chunking and Merkle-style manifest (SHA256 per-chunk)

This is intentionally small: it demonstrates the QDiff concept and is
meant to be extended by other agents (e.g. block-sparse, file-backed,
and resumable manifests).
"""

from dataclasses import dataclass, asdict
import hashlib
import json
from typing import Dict, List, Optional, Tuple

import numpy as np


@dataclass
class QDiffBundle:
    header: Dict
    chunks: Dict[str, bytes]


def _quantize_residuals(base: np.ndarray, new: np.ndarray) -> Tuple[np.ndarray, float]:
    # compute residuals and a symmetric scale to fit in int8
    resid = new - base
    max_abs = float(np.max(np.abs(resid)))
    if max_abs == 0.0:
        return np.zeros_like(resid, dtype=np.int8), 1.0
    scale = max_abs / 127.0
    q = np.round(resid / scale).astype(np.int8)
    return q, scale


def _dequantize(q: np.ndarray, scale: float) -> np.ndarray:
    return q.astype(np.float32) * scale


def _sparsify(q: np.ndarray, top_k: Optional[int]) -> Tuple[np.ndarray, Optional[List[int]]]:
    if top_k is None or top_k >= q.size:
        return q, None
    # keep top_k by absolute magnitude
    idx = np.argsort(np.abs(q))[-top_k:]
    mask = np.zeros(q.size, dtype=bool)
    mask[idx] = True
    sparse_q = np.zeros_like(q)
    sparse_q[mask] = q[mask]
    return sparse_q, idx.tolist()


def _chunk_bytes(b: bytes, chunk_size: int = 1024) -> List[bytes]:
    return [b[i : i + chunk_size] for i in range(0, len(b), chunk_size)]


def _sha256_hex(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


def create_qdiff(
    base: List[float],
    new: List[float],
    top_k: Optional[int] = None,
    chunk_size: int = 1024,
) -> QDiffBundle:
    """Create a QDiff bundle describing new relative to base.

    base and new are numeric arrays (lists or numpy-compatible). Returns
    a QDiffBundle with a header and chunk map (sha256->bytes).
    """
    a = np.asarray(base, dtype=np.float32)
    b = np.asarray(new, dtype=np.float32)
    if a.shape != b.shape:
        raise ValueError("base and new must have same shape")

    q, scale = _quantize_residuals(a, b)
    sparse_q, indices = _sparsify(q, top_k)

    # serialize: header JSON, and the quantized bytes as raw int8
    q_bytes = sparse_q.tobytes()

    chunks = {}
    chunk_list = _chunk_bytes(q_bytes, chunk_size=chunk_size)
    manifest = []
    for c in chunk_list:
        h = _sha256_hex(c)
        chunks[h] = c
        manifest.append(h)

    header = {
        "shape": list(a.shape),
        "dtype": "float32",
        "quant": "int8",
        "scale": float(scale),
        "top_k_indices": indices,
        "manifest": manifest,
        "chunk_size": chunk_size,
    }

    return QDiffBundle(header=header, chunks=chunks)


def apply_qdiff(base: List[float], bundle: QDiffBundle) -> List[float]:
    """Apply QDiff bundle to a base array and return reconstructed new array."""
    a = np.asarray(base, dtype=np.float32)
    shape = tuple(bundle.header["shape"])
    if a.shape != shape:
        raise ValueError("base shape does not match bundle header")

    # reconstruct bytes from manifest
    manifest = bundle.header["manifest"]
    parts = [bundle.chunks[h] for h in manifest]
    q_bytes = b"".join(parts)

    # ensure length matches
    expected_elems = int(np.prod(shape))
    q = np.frombuffer(q_bytes, dtype=np.int8, count=expected_elems)

    scale = float(bundle.header["scale"])
    deq = _dequantize(q, scale)
    new = a + deq.reshape(shape)
    return new.tolist()


def manifest_proofs(bundle: QDiffBundle) -> Dict[str, str]:
    """Return a mapping of chunk hash -> hex digest (Merkle leaf hashes).

    This is a thin helper used by tests and auditors.
    """
    return {h: _sha256_hex(bundle.chunks[h]) for h in bundle.header["manifest"]}