"""
OsmoImage — Docker-like portable container for ML models.

Build any model into a compressed, encrypted .osmoimg image that runs anywhere.

    # Build
    image = OsmoImage.build("user/ModelName", encrypt=True)
    image.save("model.osmoimg")

    # Run
    image = OsmoImage.load("model.osmoimg")
    output = image.run("Hello world")

    # Inspect
    info = OsmoImage.inspect("model.osmoimg")

    # Export
    image.export("pytorch", output_path="model.pt")
"""

from __future__ import annotations

import json
import os
import tempfile
import time
import hashlib
import threading
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import numpy as np

from model_osmosis.format.osmo_format import OsmoModel, OsmoManifest
from model_osmosis.format.serialization import (
    OsmoSerializer,
    OsmoLoader,
)
from model_osmosis.format.adaptive_compression import (
    ADAPTIVE_COMPRESSION_ENGINE_VERSION,
    ADAPTIVE_PAYLOAD_CODEC,
    ADAPTIVE_ZSTD_LEVELS,
    CompressionDependencyError,
    PayloadIntegrityError,
    ensure_zstandard,
    inspect_osmo_model_lossless_v2,
    pack_osmo_model_lossless_v2,
    unpack_osmo_model_lossless_v2,
)
from model_osmosis.optimizer.compression_engine import (
    resolve_compression_engine_mode,
    select_compression_selection,
)
from model_osmosis.format.file_format import (
    OSMOIMG_EXTENSION,
    write_osmoimg_header,
    read_osmoimg_header,
    is_osmoimg,
)
from model_osmosis.format.trust import (
    DEFAULT_SIGNER_ID,
    DEFAULT_SIGNER_BACKEND,
    deterministic_zip_bytes,
    ensure_default_signer,
    get_signer_backend,
    publish_transparency_event,
    sign_statement_with_backend,
    verify_trust_block_with_backend,
)


# ─── Runtime Metadata ────────────────────────────────────────────────────

@dataclass
class OsmoRuntime:
    """Runtime configuration bundled inside an .osmo image.

    Captures everything needed to reconstruct and run the model
    on any machine, regardless of the original framework.
    """

    # Framework the model was built from
    source_framework: str = "unknown"

    # Preferred framework for inference
    preferred_framework: str = "pytorch"

    # Quantization info
    quantization: Optional[str] = None  # "int8", "int4", "fp16", etc.

    # Device requirements
    min_ram_gb: float = 0.0
    recommended_gpu_vram_gb: float = 0.0
    supports_cpu: bool = True

    # Model type for inference
    model_type: str = "generic"  # "causal_lm", "seq2seq", "classifier", etc.

    # HuggingFace model ID (for architecture reconstruction)
    hf_model_id: Optional[str] = None

    # Generation config
    default_max_tokens: int = 256
    default_temperature: float = 0.7

    # Image metadata
    image_version: str = "1.0.0"
    built_at: str = ""
    built_by: str = "osmosis"

    # Compression stats
    original_size_mb: float = 0.0
    compressed_size_mb: float = 0.0
    compression_ratio: float = 1.0
    payload_codec: str = "zip-deflate"
    compression_profile: str = "legacy"
    compression_engine_version: str = "legacy-zip-v1"
    compression_stats: Dict[str, Any] = field(default_factory=dict)

    # Obfuscation
    is_encrypted: bool = False
    preserve_quality: bool = True

    # Trusted artifact metadata
    trust: Dict[str, Any] = field(default_factory=dict)
    trust_status: str = "unsigned"
    signer_id: Optional[str] = None
    key_id: Optional[str] = None
    signature_alg: Optional[str] = None
    signer_backend: str = DEFAULT_SIGNER_BACKEND
    payload_sha256: str = ""
    artifact_id: str = ""
    manifest_sha256: str = ""
    tensor_index_sha256: str = ""
    verified_at: Optional[str] = None
    transparency_status: Optional[str] = None

    # Provenance + reproducibility
    provenance_mode: str = "private"
    provenance: Dict[str, Any] = field(default_factory=dict)
    repro_mode: str = "strict"
    repro_manifest: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "OsmoRuntime":
        known = {f.name for f in cls.__dataclass_fields__.values()}
        filtered = {k: v for k, v in d.items() if k in known}
        return cls(**filtered)


# ─── OsmoImage ────────────────────────────────────────────────────────────

class OsmoImage:
    """Docker-like portable container for ML models.

    An OsmoImage bundles:
    - Model weights (compressed, optionally encrypted)
    - Architecture specification
    - Tokenizer files
    - Runtime metadata
    - Auto-generated model card

    Like Docker images, OsmoImages are:
    - Portable: run on any machine with `osmosis run`
    - Compressed: smaller than raw model files
    - Encrypted: weights can't be traced to source models
    - Self-contained: everything needed to run inference
    """

    SUPPORTED_COMPRESSION_PROFILES = {"legacy", "lossless-max"}
    SUPPORTED_COMPRESSION_ENGINE_MODES = {"auto", "off"}
    SUPPORTED_PROVENANCE_MODES = {"private", "full"}
    SUPPORTED_REPRO_MODES = {"best-effort", "strict"}
    SUPPORTED_RUNTIME_PROFILES = {"default", "cpu-lowram", "gpu-fast"}
    SUPPORTED_RUNTIME_DEVICES = {"auto", "cpu", "cuda"}
    SUPPORTED_RUNTIME_QUANTIZATION = {"none", "int8"}
    TOKENIZER_MANIFEST_FILE = "tokenizer_manifest.json"
    TOKENIZER_PLACEHOLDER_FILE = "tokenizer_placeholder.txt"
    TOKENIZER_TEXT_EXTENSIONS = {
        ".txt",
        ".json",
        ".md",
        ".yaml",
        ".yml",
        ".ini",
        ".cfg",
        ".conf",
        ".csv",
        ".tsv",
        ".template",
    }
    TEXT_MODEL_TYPES = {"causal_lm", "seq2seq", "language_model"}

    def __init__(
        self,
        osmo_model: OsmoModel,
        runtime: OsmoRuntime,
        tokenizer_data: Optional[Dict[str, Any]] = None,
        readme: str = "",
    ):
        self.osmo_model = osmo_model
        self.runtime = runtime
        self.tokenizer_data = tokenizer_data or {}
        self.readme = readme
        self._runtime_cache: Dict[str, Dict[str, Any]] = {}
        self._runtime_cache_lock = threading.RLock()

    # ── Build ──────────────────────────────────────────────────────────

    @classmethod
    def build(
        cls,
        model: Any,
        name: str = "",
        quantize: Optional[str] = None,
        compression_profile: str = "legacy",
        compression_engine: str = "auto",
        trust_sign: bool = True,
        provenance_mode: str = "private",
        repro_mode: str = "strict",
        signer_id: Optional[str] = None,
        signer_backend: Optional[str] = None,
    ) -> "OsmoImage":
        """Build an OsmoImage from any model.

        Every .osmoimg is ALWAYS:
        - Encrypted (invisible weight obfuscation, lossless)
        - Compressed (lossless zip compression)
        - Quality-preserved (exact same outputs as original)

        This is enforced by the format — there is no way to create
        an unencrypted or uncompressed .osmoimg.

        Accepts:
        - HuggingFace model ID (str): "user/ModelName"
        - PyTorch nn.Module
        - TensorFlow/Keras model
        - JAX params dict
        - File path to saved model
        - Existing OsmoModel

        Args:
            model: The model to package.
            name: Name for the image.
            quantize: Optional quantization ("int8", "fp16", None).
            compression_profile: Compression profile to persist in image metadata.
                One of "legacy" or "lossless-max".
            compression_engine: Compression optimizer engine mode ("auto" or "off").
            trust_sign: Whether to sign the artifact trust envelope on save.
            provenance_mode: One of "private" or "full".
            repro_mode: One of "best-effort" or "strict".
            signer_id: Local signer key id for trusted signing.
            signer_backend: Trust signer backend name (local-keyring or aws-kms).

        Returns:
            OsmoImage ready to save or run.
        """
        from model_osmosis.format.converters import UniversalConverter

        converter = UniversalConverter()
        runtime = OsmoRuntime(
            built_at=cls._compute_built_at(repro_mode),
            provenance_mode=cls._normalize_provenance_mode(provenance_mode),
            repro_mode=cls._normalize_repro_mode(repro_mode),
        )
        runtime.signer_id = (signer_id or DEFAULT_SIGNER_ID).strip() or DEFAULT_SIGNER_ID
        if signer_backend:
            runtime.signer_backend = signer_backend.strip() or DEFAULT_SIGNER_BACKEND
        else:
            runtime.signer_backend = get_signer_backend().backend_name
        runtime.trust = {"enabled": bool(trust_sign)}
        runtime.compression_profile = cls._normalize_compression_profile(
            compression_profile
        )
        runtime.compression_stats = dict(runtime.compression_stats or {})
        runtime.compression_stats["optimizer"] = {
            "optimizer_mode": resolve_compression_engine_mode(compression_engine),
            "profile_id": "",
            "pack_version": "",
            "selector_version": "",
            "canary_status": "shadow",
            "rollback_reason": "",
            "selection_latency_ms": 0.0,
            "fallback_reason": "",
            "zstd_levels": [],
        }
        if runtime.compression_profile == "lossless-max":
            runtime.payload_codec = ADAPTIVE_PAYLOAD_CODEC
            runtime.compression_engine_version = ADAPTIVE_COMPRESSION_ENGINE_VERSION
        else:
            runtime.payload_codec = "zip-deflate"
            runtime.compression_engine_version = "legacy-zip-v1"
        tokenizer_data = {}

        # ── Load / convert to OsmoModel ──

        if isinstance(model, OsmoModel):
            osmo = model
            runtime.source_framework = "osmo"

        elif isinstance(model, str):
            model_path = Path(model)

            # Check if it's an existing .osmoimg file
            if model_path.suffix == OSMOIMG_EXTENSION or (
                model_path.exists() and is_osmoimg(model_path)
            ):
                return cls.load(model)

            # Check if it's a local directory with a model
            if model_path.is_dir():
                osmo, tokenizer_data, runtime = cls._build_from_local(
                    model_path, name, runtime
                )
            else:
                # Assume HuggingFace model ID
                osmo, tokenizer_data, runtime = cls._build_from_hf(
                    model, name, runtime
                )

        else:
            # Native framework model (PyTorch, TF, JAX)
            if not name:
                name = type(model).__name__
            osmo = converter.to_osmo(model, name=name)
            runtime.source_framework = osmo.manifest.framework_origin.value

        tokenizer_data = cls._merge_tokenizer_from_auxiliary(osmo, tokenizer_data)
        tokenizer_data = cls._ensure_tokenizer_bundle(
            tokenizer_data,
            runtime,
            enforce_text_tokenizer=True,
        )

        # ── MANDATORY: Invisible Encryption (lossless obfuscation) ──
        # Every .osmoimg is encrypted with preserve_quality=True,
        # meaning the obfuscated model produces IDENTICAL outputs
        # to the original — zero quality loss.

        from model_osmosis.obfuscation import WeightObfuscator, ObfuscationConfig
        from model_osmosis.core.exceptions import VerificationError
        import logging

        logger = logging.getLogger("model_osmosis.obfuscation")

        config = ObfuscationConfig.auto(osmo)
        config.seed = 42

        # HuggingFace chat models must keep original parameter/layer names so they
        # can be loaded back into the same architecture for `run --prompt`.
        # Also skip generic verification here: the verifier reconstructs an MLP-like
        # forward pass and can produce false failures on large transformer checkpoints.
        if runtime.hf_model_id:
            config.obfuscate_param_names = False
            config.obfuscate_layer_names = False
            config.verify = False

        # Progressive Fallback Loop — each attempt works on a fresh copy
        # so failed transforms don't corrupt the model for the next attempt.
        original_osmo = osmo
        success = False
        while not success:
            try:
                obfuscator = WeightObfuscator(config)
                logger.info(
                    "Attempting obfuscation (permutation=%s, scaling=%s)",
                    config.enable_permutation, config.enable_scaling,
                )
                osmo = obfuscator.obfuscate(original_osmo)
                success = True
            except (VerificationError, RuntimeError, ValueError) as e:
                logger.warning("Obfuscation verification failed: %s", e)

                # Downgrade config progressively
                if config.enable_permutation:
                    logger.info("Downgrading: Disabling permutation...")
                    config.enable_permutation = False
                elif config.enable_scaling:
                    logger.info("Downgrading: Disabling scaling...")
                    config.enable_scaling = False
                elif config.obfuscate_param_names or config.obfuscate_layer_names:
                    logger.info("Downgrading: Disabling name obfuscation...")
                    config.obfuscate_param_names = False
                    config.obfuscate_layer_names = False
                else:
                    logger.warning(
                        "All obfuscation disabled. Packaging original weights."
                    )
                    # Ensure we use the clean original, not a partially-modified copy
                    osmo = original_osmo
                    break

        # If the loop exited without success, osmo is the unmodified original
        if not success:
            runtime.is_encrypted = False
        else:
            has_weight_transforms = (
                config.enable_permutation or config.enable_scaling or config.enable_rotation
            )
            has_name_obfuscation = (
                config.obfuscate_param_names or config.obfuscate_layer_names
            )
            runtime.is_encrypted = has_weight_transforms or has_name_obfuscation
        runtime.preserve_quality = config.preserve_quality

        # ── Calculate size stats ──

        total_bytes = sum(
            arr.nbytes for arr in osmo.parameters.values()
        )
        runtime.original_size_mb = total_bytes / (1024 * 1024)

        # ── Set name ──

        if name and not osmo.manifest.name:
            osmo.manifest.name = name

        # ── Generate README ──

        readme = cls._generate_readme(osmo, runtime)

        return cls(
            osmo_model=osmo,
            runtime=runtime,
            tokenizer_data=tokenizer_data,
            readme=readme,
        )

    @classmethod
    def _build_from_hf(
        cls,
        model_id: str,
        name: str,
        runtime: OsmoRuntime,
    ) -> tuple:
        """Build from a HuggingFace model ID."""
        from model_osmosis.format.converters import HuggingFaceConverter

        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

            print(f"  📥 Downloading {model_id}...")
            hf_model = AutoModelForCausalLM.from_pretrained(
                model_id, trust_remote_code=True
            )

            converter = HuggingFaceConverter()
            osmo = converter.to_osmo(
                hf_model,
                name=name or model_id.split("/")[-1],
                include_tokenizer=True,
                include_config=True,
            )

            runtime.source_framework = "huggingface"
            runtime.hf_model_id = model_id
            runtime.model_type = "causal_lm"

            # Capture config
            config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
            config_dict = config.to_dict()
            if "max_position_embeddings" in config_dict:
                runtime.default_max_tokens = min(
                    256, config_dict["max_position_embeddings"]
                )

            # Estimate RAM requirement
            param_count = sum(
                arr.size for arr in osmo.parameters.values()
            )
            runtime.min_ram_gb = round(param_count * 4 / (1024**3), 1)  # fp32
            runtime.recommended_gpu_vram_gb = round(param_count * 2 / (1024**3), 1)

            # Capture tokenizer data
            tokenizer_data = {}
            try:
                tokenizer = AutoTokenizer.from_pretrained(
                    model_id, trust_remote_code=True
                )
                import tempfile
                with tempfile.TemporaryDirectory() as tmpdir:
                    tokenizer.save_pretrained(tmpdir)
                    for f in sorted(Path(tmpdir).rglob("*")):
                        if f.is_file():
                            rel_name = f.relative_to(tmpdir).as_posix()
                            tokenizer_data[rel_name] = cls._read_tokenizer_asset(f)
            except Exception:
                pass

            return osmo, tokenizer_data, runtime

        except ImportError:
            raise ImportError(
                "HuggingFace transformers required: pip install transformers"
            )

    @classmethod
    def _build_from_local(
        cls,
        path: Path,
        name: str,
        runtime: OsmoRuntime,
    ) -> tuple:
        """Build from a local model directory."""
        from model_osmosis.format.converters import UniversalConverter

        tokenizer_data = {}

        # Check if it's a HuggingFace-format directory
        if (path / "config.json").exists():
            return cls._build_from_hf(str(path), name, runtime)

        # Check for PyTorch files
        pt_files = list(path.glob("*.pt")) + list(path.glob("*.pth"))
        if pt_files:
            import torch
            model = torch.load(pt_files[0], map_location="cpu", weights_only=False)
            converter = UniversalConverter()
            osmo = converter.to_osmo(model, name=name or path.stem)
            runtime.source_framework = "pytorch"
            return osmo, tokenizer_data, runtime

        raise ValueError(f"Cannot detect model format in {path}")

    # ── Save ───────────────────────────────────────────────────────────

    @classmethod
    def _normalize_compression_profile(cls, compression_profile: str) -> str:
        normalized = (compression_profile or "legacy").strip().lower()
        if normalized not in cls.SUPPORTED_COMPRESSION_PROFILES:
            raise ValueError(
                f"Unsupported compression profile '{compression_profile}'. "
                f"Expected one of: {sorted(cls.SUPPORTED_COMPRESSION_PROFILES)}"
            )
        return normalized

    @classmethod
    def _normalize_provenance_mode(cls, provenance_mode: str) -> str:
        normalized = (provenance_mode or "private").strip().lower()
        if normalized not in cls.SUPPORTED_PROVENANCE_MODES:
            raise ValueError(
                f"Unsupported provenance mode '{provenance_mode}'. "
                f"Expected one of: {sorted(cls.SUPPORTED_PROVENANCE_MODES)}"
            )
        return normalized

    @classmethod
    def _normalize_repro_mode(cls, repro_mode: str) -> str:
        normalized = (repro_mode or "strict").strip().lower()
        if normalized not in cls.SUPPORTED_REPRO_MODES:
            raise ValueError(
                f"Unsupported repro mode '{repro_mode}'. "
                f"Expected one of: {sorted(cls.SUPPORTED_REPRO_MODES)}"
            )
        return normalized

    @classmethod
    def _normalize_runtime_profile(cls, runtime_profile: str) -> str:
        normalized = (runtime_profile or "default").strip().lower()
        if normalized not in cls.SUPPORTED_RUNTIME_PROFILES:
            raise ValueError(
                f"Unsupported runtime profile '{runtime_profile}'. "
                f"Expected one of: {sorted(cls.SUPPORTED_RUNTIME_PROFILES)}"
            )
        return normalized

    @classmethod
    def _normalize_runtime_device(cls, device: str) -> str:
        normalized = (device or "auto").strip().lower()
        if normalized not in cls.SUPPORTED_RUNTIME_DEVICES:
            raise ValueError(
                f"Unsupported runtime device '{device}'. "
                f"Expected one of: {sorted(cls.SUPPORTED_RUNTIME_DEVICES)}"
            )
        return normalized

    @classmethod
    def _normalize_runtime_quantization(cls, runtime_quantization: str) -> str:
        normalized = (runtime_quantization or "none").strip().lower()
        if normalized not in cls.SUPPORTED_RUNTIME_QUANTIZATION:
            raise ValueError(
                f"Unsupported runtime quantization '{runtime_quantization}'. "
                f"Expected one of: {sorted(cls.SUPPORTED_RUNTIME_QUANTIZATION)}"
            )
        return normalized

    @classmethod
    def _is_text_model_type(cls, model_type: Optional[str]) -> bool:
        return str(model_type or "").strip().lower() in cls.TEXT_MODEL_TYPES

    @classmethod
    def _has_real_tokenizer_assets(cls, tokenizer_data: Dict[str, Any]) -> bool:
        keys = {str(k).strip().lower() for k in (tokenizer_data or {}).keys()}
        if "tokenizer.json" in keys:
            return True
        if "spiece.model" in keys or "tokenizer.model" in keys:
            return True
        if "vocab.txt" in keys:
            return True
        if "vocab.json" in keys and "merges.txt" in keys:
            return True
        return False

    @classmethod
    def _tokenizer_entry_kind(cls, name: str, data: Any) -> str:
        if isinstance(data, (dict, list)):
            return "json"
        if isinstance(data, bytes):
            return "binary"
        return "text"

    @classmethod
    def _tokenizer_entry_bytes(cls, name: str, data: Any) -> bytes:
        if isinstance(data, (dict, list)):
            return cls._canonical_json_bytes(data)
        if isinstance(data, bytes):
            return data
        if isinstance(data, str):
            return data.encode("utf-8")
        return str(data).encode("utf-8")

    @classmethod
    def _read_tokenizer_asset(cls, path: Path) -> Any:
        suffix = path.suffix.lower()
        if suffix == ".json":
            with open(path, "r", encoding="utf-8") as fh:
                return json.load(fh)
        if suffix in cls.TOKENIZER_TEXT_EXTENSIONS:
            with open(path, "r", encoding="utf-8") as fh:
                return fh.read()
        with open(path, "rb") as fh:
            return fh.read()

    @classmethod
    def _write_tokenizer_asset(cls, path: Path, data: Any) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        if isinstance(data, (dict, list)):
            with open(path, "w", encoding="utf-8") as fh:
                json.dump(data, fh, sort_keys=True, separators=(",", ":"))
            return
        if isinstance(data, bytes):
            with open(path, "wb") as fh:
                fh.write(data)
            return
        with open(path, "w", encoding="utf-8") as fh:
            fh.write(data if isinstance(data, str) else str(data))

    @classmethod
    def validate_tokenizer_bundle(cls, tokenizer_data: Optional[Dict[str, Any]]) -> Dict[str, Any]:
        bundle = dict(tokenizer_data or {})
        errors: List[str] = []
        manifest = bundle.get(cls.TOKENIZER_MANIFEST_FILE)
        files_checked = 0

        if manifest is not None and not isinstance(manifest, dict):
            errors.append("tokenizer_manifest_invalid_type")
            manifest = {}

        manifest_files = []
        if isinstance(manifest, dict):
            raw_files = manifest.get("files", [])
            if raw_files is None:
                raw_files = []
            if isinstance(raw_files, list):
                manifest_files = raw_files
            else:
                errors.append("tokenizer_manifest_files_invalid")

        for idx, entry in enumerate(manifest_files):
            if not isinstance(entry, dict):
                errors.append(f"tokenizer_manifest_entry_{idx}_invalid")
                continue
            name = str(entry.get("name") or "").strip()
            if not name:
                errors.append(f"tokenizer_manifest_entry_{idx}_missing_name")
                continue
            if name not in bundle:
                errors.append(f"tokenizer_manifest_missing_file:{name}")
                continue

            data = bundle[name]
            kind = cls._tokenizer_entry_kind(name, data)
            raw = cls._tokenizer_entry_bytes(name, data)
            files_checked += 1

            expected_kind = str(entry.get("kind") or "").strip()
            if expected_kind and expected_kind != kind:
                errors.append(
                    f"tokenizer_manifest_kind_mismatch:{name}:expected={expected_kind}:actual={kind}"
                )
            expected_size = entry.get("size_bytes")
            if expected_size is not None:
                try:
                    if int(expected_size) != len(raw):
                        errors.append(
                            f"tokenizer_manifest_size_mismatch:{name}:expected={int(expected_size)}:actual={len(raw)}"
                        )
                except Exception:
                    errors.append(f"tokenizer_manifest_size_invalid:{name}")
            expected_sha = str(entry.get("sha256") or "").strip()
            if expected_sha and cls._sha256_bytes(raw) != expected_sha:
                errors.append(f"tokenizer_manifest_sha256_mismatch:{name}")

        has_real_assets = cls._has_real_tokenizer_assets(bundle)
        has_placeholder = cls.TOKENIZER_PLACEHOLDER_FILE in bundle
        if errors:
            status = "missing_or_corrupt"
        elif has_real_assets:
            status = "real_assets"
        elif has_placeholder:
            status = "placeholder_only"
        else:
            status = "missing_or_corrupt"

        return {
            "status": status,
            "has_manifest": isinstance(manifest, dict),
            "manifest_schema_version": int(manifest.get("schema_version", 0))
            if isinstance(manifest, dict)
            else 0,
            "has_real_tokenizer_assets": has_real_assets,
            "files_checked": files_checked,
            "errors": sorted(errors),
        }

    @classmethod
    def _merge_tokenizer_from_auxiliary(
        cls,
        osmo: OsmoModel,
        tokenizer_data: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        merged = dict(tokenizer_data or {})
        for key in list(osmo.auxiliary_files.keys()):
            if key.startswith("tokenizer/"):
                fname = key.split("/", 1)[1]
                merged.setdefault(fname, osmo.auxiliary_files.pop(key))
        return merged

    @classmethod
    def _ensure_tokenizer_bundle(
        cls,
        tokenizer_data: Optional[Dict[str, Any]],
        runtime: OsmoRuntime,
        *,
        enforce_text_tokenizer: bool = False,
    ) -> Dict[str, Any]:
        bundle = dict(tokenizer_data or {})
        has_real_assets = cls._has_real_tokenizer_assets(bundle)

        if not has_real_assets:
            bundle.setdefault(
                cls.TOKENIZER_PLACEHOLDER_FILE,
                "No native tokenizer assets were available during packaging.",
            )

        file_entries: List[Dict[str, Any]] = []
        for fname in sorted(bundle.keys()):
            if fname == cls.TOKENIZER_MANIFEST_FILE:
                continue
            data = bundle[fname]
            raw = cls._tokenizer_entry_bytes(fname, data)
            file_entries.append(
                {
                    "name": fname,
                    "kind": cls._tokenizer_entry_kind(fname, data),
                    "sha256": cls._sha256_bytes(raw),
                    "size_bytes": len(raw),
                }
            )

        manifest = {
            "schema_version": 2,
            "embedded": True,
            "has_real_tokenizer_assets": has_real_assets,
            "model_type": runtime.model_type or "generic",
            "source_framework": runtime.source_framework or "unknown",
            "hf_model_id": runtime.hf_model_id or "",
            "files": file_entries,
        }
        bundle[cls.TOKENIZER_MANIFEST_FILE] = manifest

        if enforce_text_tokenizer and cls._is_text_model_type(runtime.model_type) and not has_real_assets:
            raise ValueError(
                "Tokenizer packaging failed: text model artifacts must include real tokenizer assets "
                "(e.g. tokenizer.json or vocab files)."
            )

        return bundle

    @staticmethod
    def _resolve_runtime_device(device: str) -> str:
        if device == "cpu":
            return "cpu"
        if device == "cuda":
            try:
                import torch

                return "cuda" if torch.cuda.is_available() else "cpu"
            except Exception:
                return "cpu"
        try:
            import torch

            return "cuda" if torch.cuda.is_available() else "cpu"
        except Exception:
            return "cpu"

    @staticmethod
    def _compute_built_at(repro_mode: str) -> str:
        if str(repro_mode).strip().lower() == "strict":
            source_epoch = os.environ.get("SOURCE_DATE_EPOCH", "").strip()
            if source_epoch:
                try:
                    epoch_i = int(source_epoch)
                    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(epoch_i))
                except Exception:
                    pass
            return "1970-01-01T00:00:00Z"
        return time.strftime("%Y-%m-%dT%H:%M:%SZ")

    @staticmethod
    def _sha256_bytes(data: bytes) -> str:
        return hashlib.sha256(data).hexdigest()

    def _build_provenance_record(self) -> Dict[str, Any]:
        mode = self.runtime.provenance_mode
        if mode == "full":
            return {
                "source_framework": self.runtime.source_framework,
                "hf_model_id": self.runtime.hf_model_id,
                "built_by": self.runtime.built_by,
                "built_at": self.runtime.built_at,
            }
        # private mode: avoid leaking direct source identities
        source_lineage = f"{self.runtime.source_framework}:{self.runtime.hf_model_id or ''}"
        return {
            "lineage_hash": self._sha256_bytes(source_lineage.encode("utf-8")),
            "builder_class": self.runtime.built_by,
        }

    @staticmethod
    def _canonical_json_bytes(data: Any) -> bytes:
        return json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8")

    def _compute_manifest_sha256(self) -> str:
        manifest_data = self.osmo_model.manifest.to_dict()
        return self._sha256_bytes(self._canonical_json_bytes(manifest_data))

    def _compute_tensor_index_sha256(self) -> str:
        tensor_index: List[Dict[str, Any]] = []
        for name in sorted(self.osmo_model.parameters.keys()):
            arr = self.osmo_model.parameters[name]
            np_arr = np.asarray(arr)
            order = "F" if np_arr.flags.f_contiguous and not np_arr.flags.c_contiguous else "C"
            raw = np_arr.tobytes(order=order)
            tensor_index.append(
                {
                    "name": name,
                    "dtype": str(np_arr.dtype),
                    "shape": list(np_arr.shape),
                    "order": order,
                    "sha256": self._sha256_bytes(raw),
                    "raw_len": len(raw),
                }
            )
        return self._sha256_bytes(self._canonical_json_bytes(tensor_index))

    def _compute_artifact_identity(self) -> None:
        self.runtime.manifest_sha256 = self._compute_manifest_sha256()
        self.runtime.tensor_index_sha256 = self._compute_tensor_index_sha256()
        artifact_descriptor = {
            "payload_sha256": self.runtime.payload_sha256,
            "manifest_sha256": self.runtime.manifest_sha256,
            "tensor_index_sha256": self.runtime.tensor_index_sha256,
            "model_fingerprint": self.osmo_model.fingerprint,
            "payload_codec": self.runtime.payload_codec,
            "compression_profile": self.runtime.compression_profile,
        }
        self.runtime.artifact_id = self._sha256_bytes(
            self._canonical_json_bytes(artifact_descriptor)
        )

    def _sort_model_maps_for_repro(self) -> None:
        self.osmo_model.parameters = {
            k: self.osmo_model.parameters[k] for k in sorted(self.osmo_model.parameters.keys())
        }
        self.osmo_model.knowledge = {
            k: self.osmo_model.knowledge[k] for k in sorted(self.osmo_model.knowledge.keys())
        }
        self.osmo_model.auxiliary_files = {
            k: self.osmo_model.auxiliary_files[k] for k in sorted(self.osmo_model.auxiliary_files.keys())
        }

    def _normalize_manifest_for_repro(self) -> None:
        deterministic_ts = self._compute_built_at("strict")
        self.osmo_model.manifest.created_at = deterministic_ts

    def _prepare_repro_manifest(self) -> None:
        import model_osmosis

        strict = self.runtime.repro_mode == "strict"
        self.runtime.repro_manifest = {
            "repro_mode": self.runtime.repro_mode,
            "seed": 42,
            "deterministic_ordering": True,
            "normalized_zip_timestamps": bool(strict),
            "tool_version": model_osmosis.__version__,
            "hash_algo": "sha256",
        }

    def _prepare_auxiliary_payload(self) -> None:
        """Persist runtime, tokenizer, and README into auxiliary files."""
        self.tokenizer_data = self._ensure_tokenizer_bundle(
            self.tokenizer_data,
            self.runtime,
            enforce_text_tokenizer=False,
        )
        tokenizer_validation = self.validate_tokenizer_bundle(self.tokenizer_data)
        self.runtime.compression_stats = dict(self.runtime.compression_stats or {})
        self.runtime.compression_stats["tokenizer_validation"] = tokenizer_validation

        for fname in sorted(self.tokenizer_data.keys()):
            data = self.tokenizer_data[fname]
            tok_key = f"tokenizer/{fname}" if "/" not in fname else fname
            self.osmo_model.auxiliary_files[tok_key] = data

        self.osmo_model.auxiliary_files["runtime.json"] = self.runtime.to_dict()
        if self.readme:
            self.osmo_model.auxiliary_files["README.md"] = self.readme

    def _serialize_legacy_payload(self) -> bytes:
        with tempfile.TemporaryDirectory() as tmpdir:
            serializer = OsmoSerializer(use_compression=True)
            tmp_zip = serializer.save(self.osmo_model, Path(tmpdir) / "model")
            return tmp_zip.read_bytes()

    def _serialize_v2_payload(
        self,
        *,
        zstd_levels=ADAPTIVE_ZSTD_LEVELS,
    ) -> tuple[bytes, Dict[str, Any]]:
        payload, stats = pack_osmo_model_lossless_v2(
            self.osmo_model,
            zstd_levels=zstd_levels,
        )
        return payload, stats

    def _write_osmoimg_with_runtime(
        self,
        path: Path,
        *,
        payload: bytes,
        header_version: int,
    ) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "wb") as f:
            write_osmoimg_header(f, self.runtime.to_dict(), version=header_version)
            f.write(payload)

    def save(
        self,
        path: Union[str, Path],
        compression_profile: Optional[str] = None,
        compression_engine: Optional[str] = None,
        trust_sign: Optional[bool] = None,
        provenance_mode: Optional[str] = None,
        repro_mode: Optional[str] = None,
        signer_id: Optional[str] = None,
        signer_backend: Optional[str] = None,
    ) -> Path:
        """Save the OsmoImage to disk as a .osmoimg file."""
        path = Path(path)
        if path.suffix not in (OSMOIMG_EXTENSION, ".osmo"):
            path = path.with_suffix(OSMOIMG_EXTENSION)
        elif path.suffix == ".osmo":
            path = path.with_suffix(OSMOIMG_EXTENSION)

        profile = self._normalize_compression_profile(
            compression_profile or self.runtime.compression_profile or "legacy"
        )
        self.runtime.compression_profile = profile
        self.runtime.provenance_mode = self._normalize_provenance_mode(
            provenance_mode or self.runtime.provenance_mode or "private"
        )
        self.runtime.repro_mode = self._normalize_repro_mode(
            repro_mode or self.runtime.repro_mode or "strict"
        )
        existing_optimizer = (
            self.runtime.compression_stats.get("optimizer", {})
            if isinstance(self.runtime.compression_stats, dict)
            else {}
        )
        resolved_compression_engine = resolve_compression_engine_mode(
            compression_engine
            if compression_engine is not None
            else str(existing_optimizer.get("optimizer_mode", "auto"))
        )
        self.runtime.signer_id = (signer_id or self.runtime.signer_id or DEFAULT_SIGNER_ID).strip() or DEFAULT_SIGNER_ID
        if signer_backend:
            self.runtime.signer_backend = signer_backend.strip() or DEFAULT_SIGNER_BACKEND
        elif not self.runtime.signer_backend:
            self.runtime.signer_backend = get_signer_backend().backend_name
        effective_trust_sign = (
            bool(self.runtime.trust.get("enabled", True))
            if trust_sign is None
            else bool(trust_sign)
        )
        # Keep payload-side runtime trust metadata deterministic/minimal.
        # The signed trust envelope is stored in the v3 header after payload assembly.
        self.runtime.trust = {"enabled": effective_trust_sign}
        self.runtime.built_at = self._compute_built_at(self.runtime.repro_mode)
        self.runtime.provenance = self._build_provenance_record()
        self._prepare_repro_manifest()

        def _optimizer_report(selection: Any) -> Dict[str, Any]:
            report = selection.to_report_dict()
            if self.runtime.repro_mode == "strict":
                report["selection_latency_ms"] = 0.0
            return report

        if self.runtime.original_size_mb <= 0:
            total_bytes = sum(arr.nbytes for arr in self.osmo_model.parameters.values())
            self.runtime.original_size_mb = total_bytes / (1024 * 1024)
        self._prepare_auxiliary_payload()
        if self.runtime.repro_mode == "strict":
            self._normalize_manifest_for_repro()
            self._sort_model_maps_for_repro()

        if profile == "legacy":
            compression_selection = select_compression_selection(
                osmo_model=self.osmo_model,
                compression_profile=profile,
                compression_engine_mode=resolved_compression_engine,
                update_state=self.runtime.repro_mode != "strict",
            )
            payload = self._serialize_legacy_payload()
            self.runtime.payload_codec = "zip-deflate"
            self.runtime.compression_engine_version = "legacy-zip-v1"
            self.runtime.compression_stats = {
                "engine_version": self.runtime.compression_engine_version,
                "payload_codec": self.runtime.payload_codec,
                "payload_bytes": len(payload),
                "optimizer": _optimizer_report(compression_selection),
            }
        else:
            ensure_zstandard()
            compression_selection = select_compression_selection(
                osmo_model=self.osmo_model,
                compression_profile=profile,
                compression_engine_mode=resolved_compression_engine,
                update_state=self.runtime.repro_mode != "strict",
            )
            zstd_levels = (
                tuple(int(level) for level in compression_selection.zstd_levels)
                if compression_selection.zstd_levels
                else ADAPTIVE_ZSTD_LEVELS
            )
            payload, stats = self._serialize_v2_payload(zstd_levels=zstd_levels)
            self.runtime.payload_codec = ADAPTIVE_PAYLOAD_CODEC
            self.runtime.compression_engine_version = ADAPTIVE_COMPRESSION_ENGINE_VERSION
            stats = dict(stats)
            stats["engine_version"] = self.runtime.compression_engine_version
            stats["payload_codec"] = self.runtime.payload_codec
            stats["optimizer"] = _optimizer_report(compression_selection)
            self.runtime.compression_stats = stats

        if self.runtime.repro_mode == "strict":
            payload = deterministic_zip_bytes(payload)

        self.runtime.payload_sha256 = self._sha256_bytes(payload)
        self._compute_artifact_identity()
        header_version = 2 if profile == "lossless-max" else 1

        if effective_trust_sign:
            import model_osmosis

            ensure_default_signer()
            statement = {
                "artifact_version": 3,
                "payload_sha256": self.runtime.payload_sha256,
                "model_fingerprint": self.osmo_model.fingerprint,
                "payload_codec": self.runtime.payload_codec,
                "compression_profile": self.runtime.compression_profile,
                "provenance_mode": self.runtime.provenance_mode,
                "repro_mode": self.runtime.repro_mode,
                "artifact_id": self.runtime.artifact_id,
                "manifest_sha256": self.runtime.manifest_sha256,
                "tensor_index_sha256": self.runtime.tensor_index_sha256,
                "tool_version": model_osmosis.__version__,
            }
            trust_block = sign_statement_with_backend(
                statement,
                signer_id=self.runtime.signer_id or DEFAULT_SIGNER_ID,
                backend_name=self.runtime.signer_backend,
            )
            self.runtime.trust = trust_block
            self.runtime.trust_status = "trusted"
            self.runtime.signer_id = trust_block.get("signer_id")
            self.runtime.key_id = trust_block.get("key_id")
            self.runtime.signature_alg = trust_block.get("alg")
            self.runtime.signer_backend = str(
                trust_block.get("backend", self.runtime.signer_backend or DEFAULT_SIGNER_BACKEND)
            )
            self.runtime.verified_at = None
            self.runtime.transparency_status = None
            header_version = 3
        else:
            self.runtime.trust = {}
            self.runtime.trust_status = "unsigned"
            self.runtime.key_id = None
            self.runtime.signature_alg = None
            self.runtime.signer_backend = DEFAULT_SIGNER_BACKEND
            self.runtime.verified_at = None
            self.runtime.transparency_status = None

        # Write once, then recompute using actual on-disk size and rewrite header
        # so stored stats reflect the real artifact.
        self._write_osmoimg_with_runtime(
            path, payload=payload, header_version=header_version
        )
        compressed_mb = path.stat().st_size / (1024 * 1024) if path.exists() else 0.0
        self.runtime.compressed_size_mb = round(compressed_mb, 2)
        if compressed_mb > 0 and self.runtime.original_size_mb > 0:
            self.runtime.compression_ratio = round(
                self.runtime.original_size_mb / compressed_mb, 2
            )
        else:
            self.runtime.compression_ratio = 1.0

        # Store true effective ratio based on persisted artifact size.
        self.runtime.compression_stats = dict(self.runtime.compression_stats or {})
        self.runtime.compression_stats["effective_compression_ratio"] = (
            round(self.runtime.original_size_mb / compressed_mb, 4)
            if compressed_mb > 0 and self.runtime.original_size_mb > 0
            else 1.0
        )
        self.runtime.compression_stats["artifact_size_bytes"] = (
            path.stat().st_size if path.exists() else 0
        )
        self.runtime.compression_stats["header_version"] = header_version

        # Rewrite once with updated runtime metadata.
        self._write_osmoimg_with_runtime(
            path, payload=payload, header_version=header_version
        )

        # Phase 2 (non-blocking): publish trust evidence to transparency hooks.
        if effective_trust_sign and isinstance(self.runtime.trust, dict) and self.runtime.trust:
            transparency = publish_transparency_event(
                self.runtime.trust,
                payload_sha256=self.runtime.payload_sha256,
                artifact_path=str(path),
            )
            self.runtime.transparency_status = transparency.get("status")

        return path

    # ── Load ───────────────────────────────────────────────────────────

    @staticmethod
    def _load_legacy_payload(payload: bytes) -> OsmoModel:
        with tempfile.TemporaryDirectory() as tmpdir:
            zip_path = Path(tmpdir) / "model.zip"
            zip_path.write_bytes(payload)
            loader = OsmoLoader()
            return loader.load(zip_path)

    @staticmethod
    def _inspect_legacy_payload(payload: bytes) -> OsmoModel:
        with tempfile.TemporaryDirectory() as tmpdir:
            zip_path = Path(tmpdir) / "model.zip"
            zip_path.write_bytes(payload)
            loader = OsmoLoader()
            return loader.load(zip_path, load_parameters=False, load_knowledge=False)

    @staticmethod
    def _load_v2_payload(payload: bytes) -> tuple[OsmoModel, Dict[str, Any]]:
        osmo, stats = unpack_osmo_model_lossless_v2(payload)
        return osmo, stats

    @classmethod
    def load(cls, path: Union[str, Path]) -> "OsmoImage":
        """Load an OsmoImage from a .osmoimg file.

        Validates magic bytes, reads the runtime header, then
        extracts the zip payload to reconstruct the OsmoModel.

        Args:
            path: Path to .osmoimg file.

        Returns:
            OsmoImage ready for inference or export.
        """
        path = Path(path)

        payload_stats: Dict[str, Any] = {}
        payload_bytes: bytes = b""

        # Read .osmoimg format (magic header + payload)
        payload_bytes: bytes = b""
        if is_osmoimg(path):
            with open(path, "rb") as f:
                runtime_data = read_osmoimg_header(f)
                payload_bytes = f.read()

            payload_codec = runtime_data.get("payload_codec", "zip-deflate")

            if payload_codec == ADAPTIVE_PAYLOAD_CODEC:
                try:
                    osmo, payload_stats = cls._load_v2_payload(payload_bytes)
                except (PayloadIntegrityError, CompressionDependencyError):
                    raise
                except Exception as exc:
                    raise ValueError(
                        f"Failed to decode v2 .osmoimg payload: {exc}"
                    ) from exc
            else:
                osmo = cls._load_legacy_payload(payload_bytes)
        else:
            # Fallback: try loading as plain .osmo zip
            loader = OsmoLoader()
            osmo = loader.load(path)
            runtime_data = osmo.auxiliary_files.pop("runtime.json", {})

        # Parse runtime
        if isinstance(runtime_data, dict):
            runtime = OsmoRuntime.from_dict(runtime_data)
        else:
            runtime = OsmoRuntime()

        if payload_stats:
            runtime.compression_stats = dict(runtime.compression_stats or {})
            runtime.compression_stats.update(payload_stats)
        if runtime.payload_codec == "zip-deflate" and runtime.compression_profile not in cls.SUPPORTED_COMPRESSION_PROFILES:
            runtime.compression_profile = "legacy"
        if runtime.payload_codec == ADAPTIVE_PAYLOAD_CODEC and runtime.compression_profile == "legacy":
            runtime.compression_profile = "lossless-max"
        if payload_bytes:
            trust_result = verify_trust_block_with_backend(runtime.trust, payload_bytes=payload_bytes)
            runtime.trust_status = trust_result.get("trust_status", "unsigned")
            runtime.signer_id = trust_result.get("signer_id")
            runtime.key_id = trust_result.get("key_id")
            runtime.signature_alg = trust_result.get("signature_alg")
            runtime.signer_backend = trust_result.get("signer_backend", runtime.signer_backend)
            runtime.payload_sha256 = trust_result.get("payload_sha256", "")
            runtime.verified_at = trust_result.get("verified_at")
            runtime.transparency_status = trust_result.get("transparency_status")
        if not runtime.manifest_sha256:
            runtime.manifest_sha256 = cls._sha256_bytes(
                cls._canonical_json_bytes(osmo.manifest.to_dict())
            )
        if not runtime.tensor_index_sha256:
            tensor_index: List[Dict[str, Any]] = []
            for name in sorted(osmo.parameters.keys()):
                arr = np.asarray(osmo.parameters[name])
                order = "F" if arr.flags.f_contiguous and not arr.flags.c_contiguous else "C"
                raw = arr.tobytes(order=order)
                tensor_index.append(
                    {
                        "name": name,
                        "dtype": str(arr.dtype),
                        "shape": list(arr.shape),
                        "order": order,
                        "sha256": cls._sha256_bytes(raw),
                        "raw_len": len(raw),
                    }
                )
            runtime.tensor_index_sha256 = cls._sha256_bytes(
                cls._canonical_json_bytes(tensor_index)
            )
        if not runtime.artifact_id:
            artifact_descriptor = {
                "payload_sha256": runtime.payload_sha256,
                "manifest_sha256": runtime.manifest_sha256,
                "tensor_index_sha256": runtime.tensor_index_sha256,
                "model_fingerprint": osmo.fingerprint,
                "payload_codec": runtime.payload_codec,
                "compression_profile": runtime.compression_profile,
            }
            runtime.artifact_id = cls._sha256_bytes(
                cls._canonical_json_bytes(artifact_descriptor)
            )
        else:
            runtime.trust_status = runtime.trust_status or "unsigned"

        # Extract tokenizer data
        tokenizer_data = {}
        for key in list(osmo.auxiliary_files.keys()):
            if key.startswith("tokenizer/"):
                fname = key.split("/", 1)[1]
                tokenizer_data[fname] = osmo.auxiliary_files.pop(key)

        tokenizer_validation = cls.validate_tokenizer_bundle(tokenizer_data)
        runtime.compression_stats = dict(runtime.compression_stats or {})
        runtime.compression_stats["tokenizer_validation"] = tokenizer_validation
        if (
            tokenizer_validation.get("manifest_schema_version", 0) >= 2
            and tokenizer_validation.get("status") == "missing_or_corrupt"
        ):
            raise ValueError(
                "Tokenizer bundle validation failed: "
                + "; ".join(tokenizer_validation.get("errors", [])[:3])
            )

        # Extract README
        readme = osmo.auxiliary_files.pop("README.md", "")

        return cls(
            osmo_model=osmo,
            runtime=runtime,
            tokenizer_data=tokenizer_data,
            readme=readme if isinstance(readme, str) else "",
        )

    # ── Run ────────────────────────────────────────────────────────────

    def run(
        self,
        prompt: Optional[str] = None,
        input_data: Optional[np.ndarray] = None,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        include_prompt: bool = False,
        runtime_profile: str = "default",
        device: str = "auto",
        runtime_cache: bool = True,
        runtime_quantization: str = "none",
    ) -> Any:
        """Run inference on the loaded model.

        For language models: pass a text prompt.
        For other models: pass numpy input_data.

        Args:
            prompt: Text prompt (for language models).
            input_data: Numpy array input (for generic models).
            max_tokens: Max tokens to generate.
            temperature: Sampling temperature.
            include_prompt: Include original prompt in decoded output for language
                models. Default is False (return only newly generated tokens).
            runtime_profile: Runtime profile (`default`, `cpu-lowram`, `gpu-fast`).
            device: Device preference (`auto`, `cpu`, `cuda`).
            runtime_cache: Whether to reuse in-process cached runtime objects.
            runtime_quantization: Runtime-only quantization (`none`, `int8`).

        Returns:
            Model output (text for LMs, numpy array for others).
        """
        if max_tokens is None:
            max_tokens = self.runtime.default_max_tokens
        if temperature is None:
            temperature = self.runtime.default_temperature

        if prompt is not None:
            return self._run_language_model(
                prompt,
                max_tokens,
                temperature,
                include_prompt=include_prompt,
                runtime_profile=runtime_profile,
                device=device,
                runtime_cache=runtime_cache,
                runtime_quantization=runtime_quantization,
            )
        elif input_data is not None:
            return self._run_generic(
                input_data,
                runtime_profile=runtime_profile,
                device=device,
                runtime_cache=runtime_cache,
                runtime_quantization=runtime_quantization,
            )
        else:
            raise ValueError("Must provide either 'prompt' or 'input_data'")

    def clear_runtime_cache(self) -> None:
        """Clear in-memory runtime cache for this image instance."""
        with self._runtime_cache_lock:
            self._runtime_cache.clear()
        try:
            import torch

            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except Exception:
            pass

    def warm_runtime(
        self,
        *,
        runtime_profile: str = "default",
        device: str = "auto",
        runtime_quantization: str = "none",
        runtime_cache: bool = True,
        warmup_prompt: Optional[str] = None,
        warmup_input_data: Optional[np.ndarray] = None,
        warmup_runs: int = 1,
    ) -> Dict[str, Any]:
        """Eagerly materialize and optionally warm runtime inference objects."""
        bundle = self._load_for_inference(
            runtime_profile=runtime_profile,
            device=device,
            runtime_cache=runtime_cache,
            runtime_quantization=runtime_quantization,
        )
        runs = max(0, int(warmup_runs))
        for _ in range(runs):
            if warmup_prompt is not None:
                self._run_language_model(
                    warmup_prompt,
                    max_tokens=1,
                    temperature=0.0,
                    include_prompt=False,
                    runtime_profile=runtime_profile,
                    device=device,
                    runtime_cache=runtime_cache,
                    runtime_quantization=runtime_quantization,
                )
            elif warmup_input_data is not None:
                self._run_generic(
                    warmup_input_data,
                    runtime_profile=runtime_profile,
                    device=device,
                    runtime_cache=runtime_cache,
                    runtime_quantization=runtime_quantization,
                )
        return {
            "runtime_profile": bundle.get("runtime_profile", "default"),
            "resolved_device": bundle.get("resolved_device", "cpu"),
            "runtime_quantization": bundle.get("runtime_quantization", "none"),
            "cache_key": bundle.get("cache_key", ""),
        }

    def _runtime_cache_key(
        self,
        *,
        runtime_profile: str,
        resolved_device: str,
        runtime_quantization: str,
    ) -> str:
        parts = [
            runtime_profile,
            resolved_device,
            runtime_quantization,
            self.runtime.preferred_framework or "pytorch",
            self.runtime.model_type or "generic",
            self.runtime.hf_model_id or "",
        ]
        return "|".join(parts)

    def _apply_runtime_optimizations(
        self,
        model: Any,
        *,
        runtime_profile: str,
        resolved_device: str,
        runtime_quantization: str,
    ) -> tuple[Any, str]:
        import torch

        effective_quantization = runtime_quantization

        # Runtime quantization is CPU-only in this release.
        if resolved_device != "cpu" and runtime_quantization != "none":
            print(
                "Warning: runtime_quantization=int8 is only supported on CPU; "
                f"falling back to none on device={resolved_device}."
            )
            effective_quantization = "none"

        if resolved_device == "cpu" and runtime_quantization == "int8":
            try:
                quantize_dynamic = getattr(torch.ao.quantization, "quantize_dynamic", None)
                if quantize_dynamic is None:
                    quantize_dynamic = torch.quantization.quantize_dynamic
                model = quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
                effective_quantization = "int8"
            except Exception as exc:
                print(
                    "Warning: runtime int8 quantization failed; falling back to none. "
                    f"Reason: {exc}"
                )
                effective_quantization = "none"

        if hasattr(model, "to"):
            try:
                model = model.to(resolved_device)
            except Exception:
                if resolved_device != "cpu":
                    model = model.to("cpu")
                    resolved_device = "cpu"

        if resolved_device == "cuda" and runtime_profile == "gpu-fast" and hasattr(model, "half"):
            try:
                model = model.half()
            except Exception:
                pass

        if hasattr(model, "eval"):
            model.eval()

        return model, effective_quantization

    def _rebuild_runtime_bundle(
        self,
        *,
        runtime_profile: str,
        resolved_device: str,
        runtime_quantization: str,
        cache_key: str,
    ) -> Dict[str, Any]:
        # Reconstruct model from OsmoModel
        from model_osmosis.format.converters import UniversalConverter

        converter = UniversalConverter()
        tokenizer = None

        if self.runtime.hf_model_id:
            # Use HF architecture
            from model_osmosis.format.converters import HuggingFaceConverter

            hf_conv = HuggingFaceConverter()
            model = hf_conv.from_osmo(
                self.osmo_model,
                model_id=self.runtime.hf_model_id,
                strict_load=True,
                allow_lossy_cast=False,
            )
            # Make sure it's loaded as CausalLM if requested
            if self.runtime.model_type == "causal_lm" and not hasattr(model, "generate"):
                from transformers import AutoModelForCausalLM

                try:
                    with tempfile.TemporaryDirectory() as tmpdir:
                        tmppath = Path(tmpdir)
                        model.save_pretrained(tmppath)
                        model = AutoModelForCausalLM.from_pretrained(
                            tmppath, trust_remote_code=True
                        )
                except Exception as e:
                    print(f"Warning: Could not convert base model to CausalLM: {e}")
        else:
            model = converter.from_osmo(
                self.osmo_model,
                target_framework=self.runtime.preferred_framework,
                strict_load=True,
                allow_lossy_cast=False,
            )

        # Reconstruct tokenizer
        if self.tokenizer_data:
            if self._has_real_tokenizer_assets(self.tokenizer_data):
                with tempfile.TemporaryDirectory() as tmpdir:
                    for fname, data in self.tokenizer_data.items():
                        fpath = Path(tmpdir) / fname
                        self._write_tokenizer_asset(fpath, data)

                    from transformers import AutoTokenizer

                    tokenizer = AutoTokenizer.from_pretrained(tmpdir)

        if tokenizer is None and self.runtime.hf_model_id:
            from transformers import AutoTokenizer

            tokenizer = AutoTokenizer.from_pretrained(self.runtime.hf_model_id)

        model, effective_quantization = self._apply_runtime_optimizations(
            model,
            runtime_profile=runtime_profile,
            resolved_device=resolved_device,
            runtime_quantization=runtime_quantization,
        )

        return {
            "cache_key": cache_key,
            "model": model,
            "tokenizer": tokenizer,
            "runtime_profile": runtime_profile,
            "resolved_device": resolved_device,
            "runtime_quantization": effective_quantization,
        }

    def _run_language_model(
        self,
        prompt: str,
        max_tokens: int,
        temperature: float,
        *,
        include_prompt: bool,
        runtime_profile: str,
        device: str,
        runtime_cache: bool,
        runtime_quantization: str,
    ) -> str:
        """Run inference for a language model."""
        bundle = self._load_for_inference(
            runtime_profile=runtime_profile,
            device=device,
            runtime_cache=runtime_cache,
            runtime_quantization=runtime_quantization,
        )
        model = bundle["model"]
        tokenizer = bundle["tokenizer"]
        resolved_device = bundle.get("resolved_device", "cpu")
        if tokenizer is None:
            raise RuntimeError(
                "Tokenizer is unavailable for this artifact runtime. "
                "Use --input for generic inference or rebuild with tokenizer metadata."
            )

        inputs = tokenizer(prompt, return_tensors="pt")
        import torch
        if resolved_device == "cuda":
            inputs = {k: v.to("cuda") for k, v in inputs.items()}
        with torch.inference_mode():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                do_sample=temperature > 0,
                temperature=temperature if temperature > 0 else None,
                pad_token_id=tokenizer.eos_token_id,
            )

        if hasattr(outputs, "detach"):
            outputs = outputs.detach().cpu()
        output_ids = outputs[0]
        if not include_prompt:
            input_ids = inputs.get("input_ids")
            is_encoder_decoder = bool(
                getattr(getattr(model, "config", None), "is_encoder_decoder", False)
            )
            if input_ids is not None and not is_encoder_decoder:
                prompt_len = int(input_ids.shape[-1])
                output_ids = output_ids[prompt_len:]
        result = tokenizer.decode(output_ids, skip_special_tokens=True)
        return result

    def _run_generic(
        self,
        input_data: np.ndarray,
        *,
        runtime_profile: str,
        device: str,
        runtime_cache: bool,
        runtime_quantization: str,
    ) -> np.ndarray:
        """Run inference for a generic model."""
        bundle = self._load_for_inference(
            runtime_profile=runtime_profile,
            device=device,
            runtime_cache=runtime_cache,
            runtime_quantization=runtime_quantization,
        )
        model = bundle["model"]
        resolved_device = bundle.get("resolved_device", "cpu")

        import torch
        if hasattr(model, "__call__"):
            x = torch.from_numpy(input_data).float()
            if resolved_device == "cuda":
                x = x.to("cuda")
            with torch.inference_mode():
                out = model(x)
                if isinstance(out, torch.Tensor):
                    return out.detach().cpu().numpy()
                return out
        return input_data

    def _load_for_inference(
        self,
        *,
        runtime_profile: str = "default",
        device: str = "auto",
        runtime_cache: bool = True,
        runtime_quantization: str = "none",
    ) -> Dict[str, Any]:
        """Load model/tokenizer for inference with optional runtime caching."""
        normalized_profile = self._normalize_runtime_profile(runtime_profile)
        normalized_device = self._normalize_runtime_device(device)
        normalized_quantization = self._normalize_runtime_quantization(runtime_quantization)
        resolved_device = self._resolve_runtime_device(normalized_device)
        cache_key = self._runtime_cache_key(
            runtime_profile=normalized_profile,
            resolved_device=resolved_device,
            runtime_quantization=normalized_quantization,
        )

        if runtime_cache:
            with self._runtime_cache_lock:
                cached = self._runtime_cache.get(cache_key)
                if cached is not None:
                    return cached

        bundle = self._rebuild_runtime_bundle(
            runtime_profile=normalized_profile,
            resolved_device=resolved_device,
            runtime_quantization=normalized_quantization,
            cache_key=cache_key,
        )
        if runtime_cache:
            with self._runtime_cache_lock:
                self._runtime_cache[cache_key] = bundle
        return bundle

    # ── Inspect ────────────────────────────────────────────────────────

    @classmethod
    def inspect(cls, path: Union[str, Path]) -> dict:
        """Inspect an .osmoimg image without loading weights.

        For .osmoimg files: reads only the magic header for instant
        metadata access (no need to decompress the full payload).

        Args:
            path: Path to .osmoimg file.

        Returns:
            Dict with image metadata.
        """
        path = Path(path)

        payload_stats: Dict[str, Any] = {}
        payload_bytes: bytes = b""

        # Fast path: read just the header from .osmoimg
        if is_osmoimg(path):
            with open(path, "rb") as f:
                runtime_data = read_osmoimg_header(f)
                payload_bytes = f.read()

            payload_codec = runtime_data.get("payload_codec", "zip-deflate")

            if payload_codec == ADAPTIVE_PAYLOAD_CODEC:
                osmo, payload_stats = inspect_osmo_model_lossless_v2(payload_bytes)
            else:
                osmo = cls._inspect_legacy_payload(payload_bytes)
        else:
            loader = OsmoLoader()
            osmo = loader.load(path, load_parameters=False, load_knowledge=False)
            runtime_data = osmo.auxiliary_files.get("runtime.json", {})

        if isinstance(runtime_data, dict):
            runtime = OsmoRuntime.from_dict(runtime_data)
        else:
            runtime = OsmoRuntime()
        if payload_stats:
            runtime.compression_stats = dict(runtime.compression_stats or {})
            runtime.compression_stats.update(payload_stats)
        if runtime.payload_codec == ADAPTIVE_PAYLOAD_CODEC and runtime.compression_profile == "legacy":
            runtime.compression_profile = "lossless-max"
        if payload_bytes:
            trust_result = verify_trust_block_with_backend(runtime.trust, payload_bytes=payload_bytes)
            runtime.trust_status = trust_result.get("trust_status", "unsigned")
            runtime.signer_id = trust_result.get("signer_id")
            runtime.key_id = trust_result.get("key_id")
            runtime.signature_alg = trust_result.get("signature_alg")
            runtime.signer_backend = trust_result.get("signer_backend", runtime.signer_backend)
            runtime.payload_sha256 = trust_result.get("payload_sha256", "")
            runtime.verified_at = trust_result.get("verified_at")
            runtime.transparency_status = trust_result.get("transparency_status")

        manifest = osmo.manifest
        arch = osmo.architecture

        # File size
        file_size_mb = path.stat().st_size / (1024 * 1024) if path.exists() else 0
        effective_ratio = (
            round(runtime.original_size_mb / file_size_mb, 4)
            if runtime.original_size_mb > 0 and file_size_mb > 0
            else 1.0
        )

        return {
            "name": manifest.name,
            "description": manifest.description,
            "version": manifest.version,
            "modality": manifest.modality.value if hasattr(manifest.modality, 'value') else str(manifest.modality),
            "framework": manifest.framework_origin.value if hasattr(manifest.framework_origin, 'value') else str(manifest.framework_origin),
            "parameter_count": manifest.parameter_count,
            "architecture": {
                "type": arch.architecture_type,
                "hidden_size": arch.hidden_size,
                "num_layers": arch.num_layers,
                "num_attention_heads": arch.num_attention_heads,
                "vocab_size": arch.vocab_size,
            },
            "runtime": {
                "source_framework": runtime.source_framework,
                "model_type": runtime.model_type,
                "min_ram_gb": runtime.min_ram_gb,
                "supports_cpu": runtime.supports_cpu,
                "is_encrypted": runtime.is_encrypted,
                "preserve_quality": runtime.preserve_quality,
                "quantization": runtime.quantization,
                "hf_model_id": runtime.hf_model_id,
                "payload_codec": runtime.payload_codec,
                "compression_profile": runtime.compression_profile,
                "compression_engine_version": runtime.compression_engine_version,
                "compression_stats": runtime.compression_stats,
                "trust_status": runtime.trust_status,
                "signer_id": runtime.signer_id,
                "key_id": runtime.key_id,
                "signature_alg": runtime.signature_alg,
                "signer_backend": runtime.signer_backend,
                "payload_sha256": runtime.payload_sha256,
                "artifact_id": runtime.artifact_id,
                "manifest_sha256": runtime.manifest_sha256,
                "tensor_index_sha256": runtime.tensor_index_sha256,
                "verified_at": runtime.verified_at,
                "transparency_status": runtime.transparency_status,
                "provenance_mode": runtime.provenance_mode,
                "repro_mode": runtime.repro_mode,
            },
            "image": {
                "file_size_mb": round(file_size_mb, 2),
                "original_size_mb": runtime.original_size_mb,
                "compression_ratio": runtime.compression_ratio,
                "compression_ratio_stored": runtime.compression_ratio,
                "effective_compression_ratio": effective_ratio,
                "built_at": runtime.built_at,
                "image_version": runtime.image_version,
            },
            "capabilities": {
                "tasks": manifest.capabilities.tasks if hasattr(manifest, 'capabilities') else [],
                "languages": manifest.capabilities.languages if hasattr(manifest, 'capabilities') else [],
            },
            "tags": manifest.tags,
            "license": manifest.license,
        }

    @classmethod
    def verify(
        cls,
        path: Union[str, Path],
        *,
        require_trust: bool = False,
    ) -> Dict[str, Any]:
        """Verify trust metadata/signature for a .osmoimg artifact."""
        path = Path(path)
        if not is_osmoimg(path):
            result = {
                "trust_status": "unsigned",
                "reason": "Not an .osmoimg artifact",
                "verified_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
            }
            if require_trust:
                raise ValueError("Artifact is not trusted: not an .osmoimg file")
            return result

        with open(path, "rb") as f:
            runtime_data = read_osmoimg_header(f)
            payload = f.read()

        runtime = OsmoRuntime.from_dict(runtime_data if isinstance(runtime_data, dict) else {})
        if not isinstance(runtime_data, dict):
            raise ValueError("Invalid .osmoimg header metadata; expected JSON object")
        header_version = int(runtime_data.get("__osmoimg_version", 1))
        if header_version >= 2 and runtime_data.get("payload_codec") not in {"zip-deflate", ADAPTIVE_PAYLOAD_CODEC}:
            raise ValueError(
                f"Invalid payload_codec for v{header_version}: {runtime_data.get('payload_codec')}"
            )
        if header_version >= 3:
            if not isinstance(runtime_data.get("trust"), dict):
                raise ValueError("Invalid v3 trust block: missing or malformed trust metadata")
            if not str(runtime_data.get("payload_sha256", "")).strip():
                raise ValueError("Invalid v3 header: missing payload_sha256")

        trust_result = verify_trust_block_with_backend(
            runtime.trust,
            payload_bytes=payload,
            require_trust=require_trust,
        )
        trust_result["header_version"] = int(runtime_data.get("__osmoimg_version", 1))
        trust_result["payload_codec"] = runtime_data.get("payload_codec", "zip-deflate")
        trust_result["provenance_mode"] = runtime_data.get("provenance_mode", "private")
        trust_result["repro_mode"] = runtime_data.get("repro_mode", "strict")
        trust_result["artifact_id"] = runtime_data.get("artifact_id", "")
        trust_result["manifest_sha256"] = runtime_data.get("manifest_sha256", "")
        trust_result["tensor_index_sha256"] = runtime_data.get("tensor_index_sha256", "")
        trust_result["signer_id"] = trust_result.get("signer_id") or runtime_data.get("signer_id")
        trust_result["key_id"] = trust_result.get("key_id") or runtime_data.get("key_id")
        trust_result["signer_backend"] = (
            trust_result.get("signer_backend")
            or runtime_data.get("signer_backend")
            or DEFAULT_SIGNER_BACKEND
        )

        if require_trust and trust_result.get("trust_status") != "trusted":
            raise ValueError(
                f"Trusted artifact required but status is {trust_result.get('trust_status')}: "
                f"{trust_result.get('reason', 'untrusted artifact')}"
            )
        return trust_result

    # ── Export ──────────────────────────────────────────────────────────

    def export(
        self,
        target_format: str = "pytorch",
        output_path: Optional[Union[str, Path]] = None,
        strict_load: bool = True,
        allow_lossy_cast: bool = False,
    ) -> Any:
        """Export the OsmoImage to any ML framework.

        Args:
            target_format: One of "pytorch", "tensorflow", "jax",
                "onnx", "huggingface".
            output_path: Where to save the exported model.
                If None, returns the model object in memory.

        Returns:
            The model in the target framework format.
        """
        from model_osmosis.format.converters import UniversalConverter

        converter = UniversalConverter()
        model = converter.from_osmo(
            self.osmo_model,
            target_framework=target_format,
            strict_load=strict_load,
            allow_lossy_cast=allow_lossy_cast,
        )

        if output_path:
            output_path = Path(output_path)
            output_path.parent.mkdir(parents=True, exist_ok=True)

            if target_format == "huggingface" and hasattr(model, "save_pretrained"):
                output_path.mkdir(parents=True, exist_ok=True)
                model.save_pretrained(output_path)
                # Also save tokenizer if available
                if self.tokenizer_data:
                    for fname, data in self.tokenizer_data.items():
                        fpath = output_path / fname
                        self._write_tokenizer_asset(fpath, data)
            elif target_format == "pytorch":
                import torch
                if hasattr(model, "state_dict"):
                    torch.save(model.state_dict(), output_path)
                else:
                    torch.save(model, output_path)

        return model

    # ── Info / Summary ─────────────────────────────────────────────────

    @property
    def name(self) -> str:
        return self.osmo_model.manifest.name

    @property
    def size_mb(self) -> float:
        return self.runtime.original_size_mb

    @property
    def is_encrypted(self) -> bool:
        return self.runtime.is_encrypted

    @property
    def parameter_count(self) -> int:
        return self.osmo_model.manifest.parameter_count

    def summary(self) -> str:
        """Human-readable summary."""
        m = self.osmo_model.manifest
        r = self.runtime
        effective_ratio = (
            r.compression_stats.get("effective_compression_ratio")
            if isinstance(r.compression_stats, dict)
            else None
        )
        lines = [
            f"OsmoImage: {m.name}",
            f"  Parameters: {m.parameter_count:,}",
            f"  Framework:  {r.source_framework}",
            f"  Encrypted:  {'✓' if r.is_encrypted else '✗'}",
            f"  Size:       {r.original_size_mb:.1f} MB"
            + (f" → {r.compressed_size_mb:.1f} MB ({r.compression_ratio:.1f}x)"
               if r.compressed_size_mb > 0 else ""),
            f"  Compression Profile: {r.compression_profile}",
            f"  Payload Codec:       {r.payload_codec}",
            f"  Compression Engine:  {r.compression_engine_version}",
            f"  Trust Status:        {r.trust_status}",
            f"  CPU:        {'✓' if r.supports_cpu else '✗'}",
            f"  RAM:        {r.min_ram_gb:.1f} GB minimum",
        ]
        if effective_ratio is not None:
            lines.append(f"  Effective Ratio:     {float(effective_ratio):.2f}x")
        if r.artifact_id:
            lines.append(f"  Artifact ID:         {r.artifact_id}")
        if r.manifest_sha256:
            lines.append(f"  Manifest SHA256:     {r.manifest_sha256}")
        if r.tensor_index_sha256:
            lines.append(f"  Tensor Index SHA256: {r.tensor_index_sha256}")
        if r.hf_model_id:
            lines.append(f"  HF Model:   {r.hf_model_id}")
        if r.signer_id:
            lines.append(f"  Signer:            {r.signer_id} ({r.key_id or 'unknown'})")
            lines.append(f"  Signer Backend:    {r.signer_backend}")
        if r.transparency_status:
            lines.append(f"  Transparency:      {r.transparency_status}")
        return "\n".join(lines)

    # ── Helpers ────────────────────────────────────────────────────────

    @staticmethod
    def _generate_readme(osmo: OsmoModel, runtime: OsmoRuntime) -> str:
        """Auto-generate a model card README."""
        m = osmo.manifest
        lines = [
            f"# {m.name or 'OsmoImage Model'}",
            "",
            f"> Built with Model Osmosis v0.8.3",
            "",
            "## Overview",
            "",
            m.description or "A portable ML model packaged as an OsmoImage.",
            "",
            "## Specifications",
            "",
            f"| Property | Value |",
            f"|----------|-------|",
            f"| Parameters | {m.parameter_count:,} |",
            f"| Framework | {runtime.source_framework} |",
            f"| Encrypted | {'Yes' if runtime.is_encrypted else 'No'} |",
            f"| Quality Preserved | {'Yes' if runtime.preserve_quality else 'No'} |",
            f"| Min RAM | {runtime.min_ram_gb:.1f} GB |",
            f"| CPU Support | {'Yes' if runtime.supports_cpu else 'No'} |",
            "",
            "## Usage",
            "",
            "```bash",
            f"# Run inference",
            f'osmosis run {m.name or "model"}.osmoimg --prompt "Your prompt here"',
            "",
            f"# Export to PyTorch",
            f'osmosis export {m.name or "model"}.osmoimg --format pytorch -o model.pt',
            "",
            f"# View info",
            f'osmosis inspect {m.name or "model"}.osmoimg',
            "```",
            "",
            "## Python API",
            "",
            "```python",
            "from model_osmosis.format.image import OsmoImage",
            "",
            f'image = OsmoImage.load("{m.name or "model"}.osmoimg")',
            f'output = image.run("Your prompt here")',
            "```",
        ]
        if m.license != "unknown":
            lines.extend(["", f"## License", "", f"{m.license}"])

        return "\n".join(lines)

    def __repr__(self) -> str:
        return (
            f"OsmoImage(name={self.name!r}, "
            f"params={self.parameter_count:,}, "
            f"encrypted={self.is_encrypted})"
        )
