"""
Lossless adaptive compression for .osmoimg v2 payloads.

This module provides per-tensor adaptive compression using zstd over two
transform families:
  - raw bytes + zstd
  - byte-shuffle + zstd

The compressor always preserves exact tensor bytes and records checksums so
decompression can validate payload integrity.
"""

from __future__ import annotations

import hashlib
import io
import json
import zipfile
from typing import Any, Dict, Optional, Tuple

import numpy as np

from model_osmosis.format.osmo_format import ArchitectureSpec, OsmoManifest, OsmoModel


ADAPTIVE_PAYLOAD_CODEC = "adaptive-zstd-v1"
ADAPTIVE_COMPRESSION_ENGINE_VERSION = "0.8.0"
ADAPTIVE_ZSTD_LEVELS = (3, 9, 19)


class CompressionDependencyError(RuntimeError):
    """Raised when an optional compression dependency is missing."""


class PayloadIntegrityError(ValueError):
    """Raised when a v2 payload fails integrity checks."""


def ensure_zstandard():
    """Import zstandard or raise an actionable installation error."""
    try:
        import zstandard as zstd
    except ImportError as exc:
        raise CompressionDependencyError(
            "lossless-max compression requires zstandard. "
            "Install with: pip install model-osmosis[compression] "
            "(or: pip install zstandard)"
        ) from exc
    return zstd


def _sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _array_order(arr: np.ndarray) -> str:
    if arr.flags.f_contiguous and not arr.flags.c_contiguous:
        return "F"
    return "C"


def _byte_shuffle(raw: bytes, itemsize: int) -> bytes:
    if itemsize <= 1 or len(raw) % itemsize != 0:
        return raw
    byte_view = np.frombuffer(raw, dtype=np.uint8).reshape(-1, itemsize)
    return byte_view.T.reshape(-1).tobytes()


def _byte_unshuffle(shuffled: bytes, itemsize: int) -> bytes:
    if itemsize <= 1 or len(shuffled) % itemsize != 0:
        return shuffled
    byte_view = np.frombuffer(shuffled, dtype=np.uint8).reshape(itemsize, -1)
    return byte_view.T.reshape(-1).tobytes()


def _compress_blob(
    raw: bytes,
    *,
    itemsize: int,
    allow_shuffle: bool,
    zstd_levels=ADAPTIVE_ZSTD_LEVELS,
) -> Tuple[Dict[str, Any], bytes]:
    zstd = ensure_zstandard()
    checksum = _sha256(raw)

    best_meta: Dict[str, Any] = {}
    best_payload: Optional[bytes] = None

    for level in zstd_levels:
        compressed_raw = zstd.ZstdCompressor(level=level).compress(raw)
        if best_payload is None or len(compressed_raw) < len(best_payload):
            best_payload = compressed_raw
            best_meta = {"codec": "zstd_raw", "level": level}

        if allow_shuffle and itemsize > 1 and len(raw) % itemsize == 0:
            shuffled = _byte_shuffle(raw, itemsize)
            compressed_shuffled = zstd.ZstdCompressor(level=level).compress(shuffled)
            if len(compressed_shuffled) < len(best_payload):
                best_payload = compressed_shuffled
                best_meta = {"codec": "zstd_byte_shuffle", "level": level}

    assert best_payload is not None  # levels list is never empty
    best_meta.update(
        {
            "raw_len": len(raw),
            "compressed_len": len(best_payload),
            "checksum": checksum,
        }
    )
    return best_meta, best_payload


def _decompress_blob(
    blob_meta: Dict[str, Any],
    compressed: bytes,
    *,
    itemsize: int,
) -> bytes:
    zstd = ensure_zstandard()

    expected_compressed_len = int(blob_meta.get("compressed_len", -1))
    if expected_compressed_len >= 0 and expected_compressed_len != len(compressed):
        raise PayloadIntegrityError(
            f"Compressed length mismatch: expected {expected_compressed_len}, got {len(compressed)}"
        )

    codec = blob_meta.get("codec")
    try:
        raw = zstd.ZstdDecompressor().decompress(compressed)
    except Exception as exc:
        raise PayloadIntegrityError(f"Failed to decompress v2 blob: {exc}") from exc
    if codec == "zstd_byte_shuffle":
        raw = _byte_unshuffle(raw, itemsize)
    elif codec != "zstd_raw":
        raise PayloadIntegrityError(f"Unsupported v2 blob codec: {codec}")

    expected_raw_len = int(blob_meta.get("raw_len", -1))
    if expected_raw_len >= 0 and expected_raw_len != len(raw):
        raise PayloadIntegrityError(
            f"Raw length mismatch: expected {expected_raw_len}, got {len(raw)}"
        )

    expected_checksum = blob_meta.get("checksum")
    actual_checksum = _sha256(raw)
    if expected_checksum and actual_checksum != expected_checksum:
        raise PayloadIntegrityError(
            "Checksum mismatch while decoding v2 payload "
            f"(expected {expected_checksum}, got {actual_checksum})"
        )

    return raw


def _encode_auxiliary(value: Any) -> Tuple[str, bytes]:
    if isinstance(value, bytes):
        return "bytes", value
    if isinstance(value, (dict, list)):
        return "json", json.dumps(value, separators=(",", ":"), default=str).encode("utf-8")
    return "text", str(value).encode("utf-8")


def _decode_auxiliary(kind: str, raw: bytes) -> Any:
    if kind == "bytes":
        return raw
    if kind == "json":
        return json.loads(raw.decode("utf-8"))
    return raw.decode("utf-8", errors="replace")


def _pack_tensor_entry(
    name: str,
    tensor: np.ndarray,
    *,
    zstd_levels=ADAPTIVE_ZSTD_LEVELS,
) -> Tuple[Dict[str, Any], bytes]:
    arr = np.asarray(tensor)
    order = _array_order(arr)
    if order == "F":
        ordered = np.asfortranarray(arr)
    else:
        ordered = np.ascontiguousarray(arr)
    raw = ordered.tobytes(order=order)

    blob_meta, payload = _compress_blob(
        raw,
        itemsize=ordered.dtype.itemsize,
        allow_shuffle=True,
        zstd_levels=zstd_levels,
    )
    tensor_meta = {
        "name": name,
        "dtype": str(ordered.dtype),
        "shape": list(ordered.shape),
        "order": order,
        **blob_meta,
    }
    return tensor_meta, payload


def _unpack_tensor_entry(
    tensor_meta: Dict[str, Any],
    compressed: bytes,
) -> np.ndarray:
    dtype = np.dtype(tensor_meta["dtype"])
    order = tensor_meta.get("order", "C")
    shape = tuple(tensor_meta.get("shape", []))
    raw = _decompress_blob(tensor_meta, compressed, itemsize=dtype.itemsize)

    arr = np.frombuffer(raw, dtype=dtype)
    if shape:
        arr = arr.reshape(shape, order=order)

    if order == "F":
        return np.array(arr, copy=True, order="F")
    return np.array(arr, copy=True, order="C")


def pack_osmo_model_lossless_v2(
    osmo_model: OsmoModel,
    *,
    zstd_levels=ADAPTIVE_ZSTD_LEVELS,
) -> Tuple[bytes, Dict[str, Any]]:
    """
    Serialize an OsmoModel into a v2 adaptive payload.

    Returns:
        (payload_bytes, compression_stats)
    """
    # Validate optional dependency early with actionable message.
    ensure_zstandard()

    descriptor: Dict[str, Any] = {
        "format": "osmoimg-payload-v2",
        "payload_codec": ADAPTIVE_PAYLOAD_CODEC,
        "engine_version": ADAPTIVE_COMPRESSION_ENGINE_VERSION,
        "zstd_levels": [int(level) for level in zstd_levels],
        "manifest": osmo_model.manifest.to_dict(),
        "architecture": osmo_model.architecture.to_dict(),
        "fingerprint": osmo_model.fingerprint,
        "parameters": [],
        "knowledge": [],
        "auxiliary": [],
    }

    stats: Dict[str, Any] = {
        "parameter_tensors": 0,
        "knowledge_tensors": 0,
        "auxiliary_items": 0,
        "raw_tensor_bytes": 0,
        "compressed_tensor_bytes": 0,
        "raw_auxiliary_bytes": 0,
        "compressed_auxiliary_bytes": 0,
        "selected_codecs": {},
    }

    def _count_codec(codec_name: str) -> None:
        stats["selected_codecs"][codec_name] = stats["selected_codecs"].get(codec_name, 0) + 1

    with io.BytesIO() as buffer:
        with zipfile.ZipFile(buffer, "w", compression=zipfile.ZIP_STORED) as zf:
            # Parameters
            for idx, name in enumerate(sorted(osmo_model.parameters.keys())):
                tensor = osmo_model.parameters[name]
                tensor_meta, payload = _pack_tensor_entry(
                    name,
                    tensor,
                    zstd_levels=zstd_levels,
                )
                blob_name = f"parameters/{idx:08d}.bin"
                zf.writestr(blob_name, payload)

                tensor_meta["blob"] = blob_name
                descriptor["parameters"].append(tensor_meta)

                stats["parameter_tensors"] += 1
                stats["raw_tensor_bytes"] += int(tensor_meta["raw_len"])
                stats["compressed_tensor_bytes"] += int(tensor_meta["compressed_len"])
                _count_codec(str(tensor_meta["codec"]))

            # Knowledge tensors
            for idx, name in enumerate(sorted(osmo_model.knowledge.keys())):
                tensor = osmo_model.knowledge[name]
                tensor_meta, payload = _pack_tensor_entry(
                    name,
                    tensor,
                    zstd_levels=zstd_levels,
                )
                blob_name = f"knowledge/{idx:08d}.bin"
                zf.writestr(blob_name, payload)

                tensor_meta["blob"] = blob_name
                descriptor["knowledge"].append(tensor_meta)

                stats["knowledge_tensors"] += 1
                stats["raw_tensor_bytes"] += int(tensor_meta["raw_len"])
                stats["compressed_tensor_bytes"] += int(tensor_meta["compressed_len"])
                _count_codec(str(tensor_meta["codec"]))

            # Auxiliary files
            for idx, name in enumerate(sorted(osmo_model.auxiliary_files.keys())):
                value = osmo_model.auxiliary_files[name]
                kind, raw = _encode_auxiliary(value)
                blob_meta, payload = _compress_blob(
                    raw,
                    itemsize=1,
                    allow_shuffle=False,
                    zstd_levels=zstd_levels,
                )

                blob_name = f"auxiliary/{idx:08d}.bin"
                zf.writestr(blob_name, payload)

                descriptor["auxiliary"].append(
                    {
                        "name": name,
                        "kind": kind,
                        "blob": blob_name,
                        **blob_meta,
                    }
                )

                stats["auxiliary_items"] += 1
                stats["raw_auxiliary_bytes"] += int(blob_meta["raw_len"])
                stats["compressed_auxiliary_bytes"] += int(blob_meta["compressed_len"])
                _count_codec(str(blob_meta["codec"]))

            descriptor_json = json.dumps(descriptor, separators=(",", ":"), sort_keys=True)
            zf.writestr("meta/descriptor.json", descriptor_json.encode("utf-8"))

        payload_bytes = buffer.getvalue()

    compressed_total = (
        stats["compressed_tensor_bytes"] + stats["compressed_auxiliary_bytes"]
    )
    raw_total = stats["raw_tensor_bytes"] + stats["raw_auxiliary_bytes"]
    stats["tensor_count"] = stats["parameter_tensors"] + stats["knowledge_tensors"]
    stats["payload_bytes"] = len(payload_bytes)
    stats["raw_total_bytes"] = raw_total
    stats["compressed_total_bytes"] = compressed_total
    stats["tensor_compression_ratio"] = (
        round(raw_total / compressed_total, 4) if compressed_total > 0 else 1.0
    )
    stats["payload_overhead_bytes"] = max(0, len(payload_bytes) - compressed_total)
    stats["zstd_levels"] = [int(level) for level in zstd_levels]

    return payload_bytes, stats


def unpack_osmo_model_lossless_v2(payload: bytes) -> Tuple[OsmoModel, Dict[str, Any]]:
    """
    Deserialize a v2 adaptive payload into an OsmoModel.

    Returns:
        (osmo_model, compression_stats)
    """
    ensure_zstandard()

    try:
        with zipfile.ZipFile(io.BytesIO(payload), "r") as zf:
            descriptor = json.loads(zf.read("meta/descriptor.json").decode("utf-8"))

            if descriptor.get("format") != "osmoimg-payload-v2":
                raise PayloadIntegrityError("Invalid v2 payload descriptor format")
            if descriptor.get("payload_codec") != ADAPTIVE_PAYLOAD_CODEC:
                raise PayloadIntegrityError(
                    f"Unexpected payload codec {descriptor.get('payload_codec')!r}"
                )

            parameters: Dict[str, np.ndarray] = {}
            for tensor_meta in descriptor.get("parameters", []):
                blob_name = tensor_meta["blob"]
                compressed = zf.read(blob_name)
                parameters[tensor_meta["name"]] = _unpack_tensor_entry(tensor_meta, compressed)

            knowledge: Dict[str, np.ndarray] = {}
            for tensor_meta in descriptor.get("knowledge", []):
                blob_name = tensor_meta["blob"]
                compressed = zf.read(blob_name)
                knowledge[tensor_meta["name"]] = _unpack_tensor_entry(tensor_meta, compressed)

            auxiliary: Dict[str, Any] = {}
            for aux_meta in descriptor.get("auxiliary", []):
                blob_name = aux_meta["blob"]
                compressed = zf.read(blob_name)
                raw = _decompress_blob(aux_meta, compressed, itemsize=1)
                auxiliary[aux_meta["name"]] = _decode_auxiliary(aux_meta.get("kind", "text"), raw)

    except PayloadIntegrityError:
        raise
    except zipfile.BadZipFile as exc:
        raise PayloadIntegrityError(f"Invalid v2 payload archive: {exc}") from exc
    except KeyError as exc:
        raise PayloadIntegrityError(f"Missing v2 payload entry: {exc}") from exc
    except json.JSONDecodeError as exc:
        raise PayloadIntegrityError(f"Invalid v2 descriptor JSON: {exc}") from exc

    manifest = OsmoManifest.from_dict(descriptor["manifest"])
    architecture = ArchitectureSpec.from_dict(descriptor["architecture"])
    fingerprint = descriptor.get("fingerprint", "")

    osmo = OsmoModel(
        manifest=manifest,
        architecture=architecture,
        parameters=parameters,
        knowledge=knowledge,
        auxiliary_files=auxiliary,
        fingerprint=fingerprint,
    )

    if fingerprint and osmo.fingerprint != fingerprint:
        raise PayloadIntegrityError(
            "Model fingerprint mismatch after v2 reconstruction "
            f"(expected {fingerprint}, got {osmo.fingerprint})"
        )

    stats = {
        "tensor_count": len(descriptor.get("parameters", [])) + len(descriptor.get("knowledge", [])),
        "parameter_tensors": len(descriptor.get("parameters", [])),
        "knowledge_tensors": len(descriptor.get("knowledge", [])),
        "auxiliary_items": len(descriptor.get("auxiliary", [])),
        "engine_version": descriptor.get("engine_version", ADAPTIVE_COMPRESSION_ENGINE_VERSION),
        "payload_bytes": len(payload),
    }
    return osmo, stats


def inspect_osmo_model_lossless_v2(payload: bytes) -> Tuple[OsmoModel, Dict[str, Any]]:
    """
    Read v2 payload metadata without decompressing tensors.

    Returns:
        (osmo_model_with_metadata_only, compression_stats)
    """
    try:
        with zipfile.ZipFile(io.BytesIO(payload), "r") as zf:
            descriptor = json.loads(zf.read("meta/descriptor.json").decode("utf-8"))
    except zipfile.BadZipFile as exc:
        raise PayloadIntegrityError(f"Invalid v2 payload archive: {exc}") from exc
    except KeyError as exc:
        raise PayloadIntegrityError(f"Missing v2 payload entry: {exc}") from exc
    except json.JSONDecodeError as exc:
        raise PayloadIntegrityError(f"Invalid v2 descriptor JSON: {exc}") from exc

    if descriptor.get("format") != "osmoimg-payload-v2":
        raise PayloadIntegrityError("Invalid v2 payload descriptor format")
    if descriptor.get("payload_codec") != ADAPTIVE_PAYLOAD_CODEC:
        raise PayloadIntegrityError(
            f"Unexpected payload codec {descriptor.get('payload_codec')!r}"
        )

    manifest = OsmoManifest.from_dict(descriptor["manifest"])
    architecture = ArchitectureSpec.from_dict(descriptor["architecture"])
    osmo = OsmoModel(
        manifest=manifest,
        architecture=architecture,
        parameters={},
        knowledge={},
        auxiliary_files={},
        fingerprint=descriptor.get("fingerprint", ""),
    )

    compressed_total = sum(
        int(item.get("compressed_len", 0))
        for item in descriptor.get("parameters", [])
    ) + sum(
        int(item.get("compressed_len", 0))
        for item in descriptor.get("knowledge", [])
    ) + sum(
        int(item.get("compressed_len", 0))
        for item in descriptor.get("auxiliary", [])
    )
    raw_total = sum(
        int(item.get("raw_len", 0))
        for item in descriptor.get("parameters", [])
    ) + sum(
        int(item.get("raw_len", 0))
        for item in descriptor.get("knowledge", [])
    ) + sum(
        int(item.get("raw_len", 0))
        for item in descriptor.get("auxiliary", [])
    )

    stats = {
        "tensor_count": len(descriptor.get("parameters", [])) + len(descriptor.get("knowledge", [])),
        "parameter_tensors": len(descriptor.get("parameters", [])),
        "knowledge_tensors": len(descriptor.get("knowledge", [])),
        "auxiliary_items": len(descriptor.get("auxiliary", [])),
        "engine_version": descriptor.get("engine_version", ADAPTIVE_COMPRESSION_ENGINE_VERSION),
        "payload_bytes": len(payload),
        "raw_total_bytes": raw_total,
        "compressed_total_bytes": compressed_total,
        "tensor_compression_ratio": (
            round(raw_total / compressed_total, 4) if compressed_total > 0 else 1.0
        ),
    }
    return osmo, stats
