#!/usr/bin/env python3
"""Lossless storage codec for portable DesignIR 3.0 documents.

The semantic envelope stays ordinary JSON so an agent can inspect parameters,
features, constraints, and edit bindings without consuming the dense SurfaceIR
coordinate payload. SurfaceIR is compressed losslessly inside the same JSON
document, preserving single-file source-independent rebuilds.
"""

from __future__ import annotations

import base64
import copy
import gzip
import hashlib
import json
from pathlib import Path
from typing import Any


SURFACE_LAYER_ENCODING = "gzip+base64+json"
SURFACE_LAYER_FORMAT = "surfaceir-3.0"


class DesignIRCodecError(ValueError):
    """Raised when a compact DesignIR payload is corrupt or unsupported."""


def is_encoded_surface_layer(value: Any) -> bool:
    return (
        isinstance(value, dict)
        and value.get("encoding") == SURFACE_LAYER_ENCODING
        and value.get("format") == SURFACE_LAYER_FORMAT
        and isinstance(value.get("data"), str)
    )


def encode_surface_layer(payload: dict[str, Any]) -> dict[str, Any]:
    """Return a copy with SurfaceIR losslessly compressed in-place."""
    surface_layer = payload.get("surface_layer")
    if surface_layer is None or is_encoded_surface_layer(surface_layer):
        return copy.deepcopy(payload)
    if not isinstance(surface_layer, dict):
        raise DesignIRCodecError("surface_layer must be a JSON object")

    raw = json.dumps(
        surface_layer,
        ensure_ascii=False,
        separators=(",", ":"),
    ).encode("utf-8")
    compressed = gzip.compress(raw, compresslevel=9, mtime=0)
    encoded = copy.copy(payload)
    encoded["surface_layer"] = {
        "format": SURFACE_LAYER_FORMAT,
        "encoding": SURFACE_LAYER_ENCODING,
        "uncompressed_bytes": len(raw),
        "compressed_bytes": len(compressed),
        "sha256": hashlib.sha256(raw).hexdigest(),
        "data": base64.b64encode(compressed).decode("ascii"),
    }
    return encoded


def decode_surface_layer(payload: dict[str, Any]) -> dict[str, Any]:
    """Return a copy with an encoded SurfaceIR layer expanded and verified."""
    envelope = payload.get("surface_layer")
    if not is_encoded_surface_layer(envelope):
        return copy.deepcopy(payload)
    try:
        compressed = base64.b64decode(envelope["data"], validate=True)
        raw = gzip.decompress(compressed)
    except (ValueError, OSError) as exc:
        raise DesignIRCodecError(
            f"Cannot decode compact SurfaceIR payload: {exc}"
        ) from exc

    expected_compressed = envelope.get("compressed_bytes")
    if expected_compressed is not None and len(compressed) != expected_compressed:
        raise DesignIRCodecError("Compact SurfaceIR compressed size mismatch")
    expected_uncompressed = envelope.get("uncompressed_bytes")
    if expected_uncompressed is not None and len(raw) != expected_uncompressed:
        raise DesignIRCodecError("Compact SurfaceIR uncompressed size mismatch")
    expected_digest = envelope.get("sha256")
    if expected_digest and hashlib.sha256(raw).hexdigest() != expected_digest:
        raise DesignIRCodecError("Compact SurfaceIR SHA-256 mismatch")
    try:
        surface_layer = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise DesignIRCodecError(
            f"Decoded SurfaceIR is not valid JSON: {exc}"
        ) from exc
    if not isinstance(surface_layer, dict):
        raise DesignIRCodecError("Decoded SurfaceIR must be a JSON object")

    decoded = copy.copy(payload)
    decoded["surface_layer"] = surface_layer
    return decoded


def read_designir(path: Path) -> dict[str, Any]:
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise DesignIRCodecError(f"Cannot read DesignIR JSON {path}: {exc}") from exc
    if not isinstance(payload, dict):
        raise DesignIRCodecError("DesignIR must be a JSON object")
    return decode_surface_layer(payload)


def write_designir(
    path: Path,
    payload: dict[str, Any],
    *,
    encode_surface: bool = True,
) -> None:
    stored = encode_surface_layer(payload) if encode_surface else payload
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(stored, ensure_ascii=False, separators=(",", ":")) + "\n",
        encoding="utf-8",
    )