"""
Model Osmosis CLI — Docker-style interface for model operations.

Usage:
    osmosis make       --mode merge --models MODEL_A MODEL_B --output-type osmo
    osmosis protect    MODEL -o model.osmoimg
    osmosis check      model.osmoimg --self-check-level smoke
    osmosis launch     model.osmoimg --action run --prompt "Hello"
    osmosis expand     hf:repo/model --target-params x1.5
    osmosis expand-status --json
    osmosis expand-rollback --to model.r0001
    osmosis artifact   track model.osmoimg
    osmosis build      MODEL -o image.osmoimg
    osmosis run        image.osmoimg --prompt "Hello world"
    osmosis inspect    image.osmoimg
    osmosis export     image.osmoimg --format pytorch -o model.pt
    osmosis register                  (register .osmoimg with OS)
    osmosis merge      MODEL1 MODEL2 [MODEL3...] -o OUTPUT -s STRATEGY
    osmosis obfuscate  MODEL -o OUTPUT
    osmosis compare    MODEL1 MODEL2
    osmosis benchmark  MODEL1 MODEL2 [MODEL3...]
    osmosis info       MODEL
    osmosis version
    osmosis images     [DIRECTORY]
    osmosis pull       user/Model -o model.osmoimg
    osmosis push       model.osmoimg user/repo
    osmosis convert    model.pt --to tensorflow -o model.tf
    osmosis diff       a.osmoimg b.osmoimg
    osmosis validate   model.osmoimg
    osmosis verify     model.osmoimg --require-trust
    osmosis quantize   model.osmoimg --bits 8
    osmosis pipeline  --mode merge --output-type osmoimg --models MODEL_A MODEL_B
    osmosis trust keygen --name default
    osmosis demo
    osmosis doctor
    osmosis transcode  model.osmoimg -o model.v2.osmoimg --compression-profile lossless-max
    osmosis benchmark-runtime model.osmoimg --profile i5_2gb --json
    osmosis serve      model.osmoimg --port 8080
    osmosis history    model.osmoimg

Examples:
    osmosis make --mode weight-merge --models hf:model/a hf:model/b --output-type osmoimg -o merged.osmoimg
    osmosis protect HuggingFaceTB/SmolLM2-135M-Instruct -o trusted.osmoimg
    osmosis check trusted.osmoimg --self-check-level parity --json
    osmosis launch trusted.osmoimg --action serve --port 8080
    osmosis expand hf:HuggingFaceTB/SmolLM2-135M-Instruct --data-mode auto
    osmosis artifact track trusted.osmoimg
    osmosis artifact verify --all
    osmosis build user/ModelName -o model.osmoimg
    osmosis demo
    osmosis run model.osmoimg --prompt "def fibonacci(n):"
    osmosis pull user/ModelName -o model.osmoimg
    osmosis pipeline --mode absorb --output-type pytorch --source hf:source/model --target pt:target.pt
    osmosis transcode model.osmoimg --compression-profile lossless-max
    osmosis benchmark-runtime model.osmoimg --profile i5_2gb --json
    osmosis serve model.osmoimg --port 8080
"""

from __future__ import annotations

import argparse
import sys
import os
import platform
import time
import subprocess
import importlib.util
import json
import shutil
import urllib.error
import urllib.request
from datetime import datetime, timezone
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np


DEMO_DEFAULT_MODEL_ID = "HuggingFaceTB/SmolLM2-135M-Instruct"
DEMO_DEFAULT_OUTPUT = "models/quick_demo_smollm2_135m_instruct.osmoimg"
DEMO_DEFAULT_PROMPT = "Write one short friendly sentence introducing yourself."
RUNTIME_BENCHMARK_DEFAULT_PROMPT = "Hello from osmosis runtime benchmark."
RUNTIME_BENCHMARK_BASELINES = Path("benchmarks/runtime_baselines.json")
DEFAULT_RUNTIME_BENCHMARK_BASELINES: Dict[str, Any] = {
    "schema_version": 2,
    "profiles": {
        "i5_2gb": {
            "max_cold_start_seconds": 40.0,
            "max_peak_rss_mb": 2000.0,
            "max_first_token_latency_seconds": 20.0,
            "max_warm_first_token_latency_seconds": 8.0,
            "min_tokens_per_second": 0.5,
            "max_p95_latency_seconds": 35.0,
        },
        "ci_standard": {
            "max_cold_start_seconds": 35.0,
            "max_peak_rss_mb": 2600.0,
            "max_first_token_latency_seconds": 18.0,
            "max_warm_first_token_latency_seconds": 7.0,
            "min_tokens_per_second": 0.4,
            "max_p95_latency_seconds": 30.0,
        },
        "ci_gpu_standard": {
            "max_cold_start_seconds": 20.0,
            "max_peak_rss_mb": 12000.0,
            "max_first_token_latency_seconds": 8.0,
            "max_warm_first_token_latency_seconds": 2.0,
            "min_tokens_per_second": 2.0,
            "max_p95_latency_seconds": 8.0,
        },
    },
}
ARTIFACT_LOCKFILE_DEFAULT = Path(".model_osmosis") / "artifacts.lock.json"
PIPELINE_SCHEMES = {"auto", "hf", "pt", "onnx", "tf", "jax"}
EXPAND_SCHEMES = {
    "auto",
    "hf",
    "pt",
    "onnx",
    "tf",
    "jax",
    "vision",
    "audio",
    "diffusion",
    "multimodal",
    "lora",
    "mlx",
    "osmo",
    "osmoimg",
}
PIPELINE_MODES = {"merge", "absorb", "stack", "weight-merge"}
MAKE_MODES = {"merge", "weight-merge", "absorb", "stack", "cross-merge"}
PIPELINE_OUTPUT_TYPES = {
    "osmoimg",
    "osmo",
    "pytorch",
    "hf",
    "onnx",
    "tensorflow",
    "keras",
    "jax",
    "flax",
    "safetensors",
    "torchscript",
}
TRUST_PROFILES = {"standard", "production"}
PRODUCTION_POLICY_BASELINE: Dict[str, Any] = {
    "min_header_version": 3,
    "required_trust_status": "trusted",
}
RELEASE_REQUIRED_WORKFLOWS = [
    "clean-machine-validation",
    "trusted-osmoimg",
    "runtime-benchmark",
    "cross-framework-reliability",
]
STRATEGY_ENGINE_MODES = {"auto", "off"}
COMPRESSION_ENGINE_MODES = {"auto", "off"}
EXPAND_TASKS = {"auto", "text", "vision", "audio", "code", "classification", "regression"}
EXPAND_EXPORT_FORMATS = {"osmoimg", "pytorch", "hf", "onnx", "tensorflow", "jax"}
EXPAND_FALLBACK_MODES = {"off", "distill"}


# ─── Colors ────────────────────────────────────────────────────────────────

class Colors:
    RESET  = "\033[0m"
    BOLD   = "\033[1m"
    DIM    = "\033[2m"
    GREEN  = "\033[32m"
    BLUE   = "\033[34m"
    CYAN   = "\033[36m"
    YELLOW = "\033[33m"
    RED    = "\033[31m"
    MAGENTA= "\033[35m"

    @classmethod
    def disable(cls):
        for attr in ["RESET", "BOLD", "DIM", "GREEN", "BLUE", "CYAN",
                      "YELLOW", "RED", "MAGENTA"]:
            setattr(cls, attr, "")


# Disable colors if not a terminal
if not sys.stdout.isatty():
    Colors.disable()


def _logo():
    return f"""{Colors.CYAN}{Colors.BOLD}
  ╔══════════════════════════════════════╗
  ║       🧬 Model Osmosis v0.8.3       ║
  ║   Universal Model Merge & Protect   ║
  ╚══════════════════════════════════════╝{Colors.RESET}"""


def _print_step(icon: str, msg: str):
    print(f"  {Colors.GREEN}{icon}{Colors.RESET} {msg}")


def _print_error(msg: str):
    print(f"  {Colors.RED}✗{Colors.RESET} {msg}", file=sys.stderr)


def _print_info(key: str, value: str):
    print(f"  {Colors.DIM}{key:20s}{Colors.RESET} {value}")


def _reliability_fallbacks(allow_partial_load: bool, allow_lossy_cast: bool) -> str:
    fallbacks: List[str] = []
    if allow_partial_load:
        fallbacks.append("partial-load")
    if allow_lossy_cast:
        fallbacks.append("lossy-cast")
    return ", ".join(fallbacks) if fallbacks else "none"


def _normalize_mode_arg(value: Any, allowed: Sequence[str], default: str) -> str:
    if not isinstance(value, str):
        return default
    normalized = value.strip().lower()
    return normalized if normalized in set(allowed) else default


def _resolve_strategy_engine_mode(requested: Optional[str]) -> str:
    from model_osmosis.merging.smart_merge_engine import resolve_strategy_engine_mode

    return resolve_strategy_engine_mode(requested)


def _resolve_compression_engine_mode(requested: Optional[str]) -> str:
    from model_osmosis.optimizer.compression_engine import resolve_compression_engine_mode

    return resolve_compression_engine_mode(requested)


def _truthy_env(value: str) -> bool:
    return str(value).strip().lower() in {"1", "true", "yes", "on"}


def _trust_required(flag: bool) -> bool:
    return bool(flag) or _truthy_env(os.environ.get("OSMOSIS_REQUIRE_TRUST", "0"))


def _resolve_trust_profile(profile: Optional[str]) -> str:
    if not isinstance(profile, str):
        profile = None
    requested = (profile or os.environ.get("OSMOSIS_TRUST_PROFILE", "standard")).strip().lower()
    if requested not in TRUST_PROFILES:
        raise ValueError(
            f"Unsupported trust profile '{requested}'. "
            f"Expected one of: {sorted(TRUST_PROFILES)}"
        )
    return requested


def _effective_require_trust(require_trust_flag: bool, trust_profile: str) -> bool:
    if trust_profile == "production":
        return True
    return _trust_required(require_trust_flag)


def _policy_with_trust_profile(policy: Dict[str, Any], trust_profile: str) -> Dict[str, Any]:
    merged = dict(policy or {})
    if trust_profile != "production":
        return merged
    baseline = dict(PRODUCTION_POLICY_BASELINE)
    current_min = int(merged.get("min_header_version", 0) or 0)
    merged["min_header_version"] = max(int(baseline["min_header_version"]), current_min)
    merged["required_trust_status"] = str(baseline["required_trust_status"])
    return merged


def _next_command_line(command: str) -> str:
    return f"Next command: {command.strip()}"


def _error_with_next_command(message: str, command: str) -> RuntimeError:
    return RuntimeError(f"{message}\n{_next_command_line(command)}")


def _infer_next_command(args: Any, error_message: str) -> str:
    cmd_name = str(getattr(args, "command", "") or "").strip()
    msg = str(error_message or "")
    lower = msg.lower()

    marker = "Install with:"
    if marker in msg:
        install_cmd = msg.split(marker, 1)[1].strip().splitlines()[0].strip()
        if install_cmd:
            return install_cmd

    image = str(getattr(args, "image", "<image.osmoimg>"))
    if "blocked by policy" in lower or "policy file" in lower:
        if cmd_name in {"run", "serve", "export"}:
            return (
                f"osmosis {cmd_name} {image} --trust-profile production "
                "--policy-file <policy.json>"
            )
        return "osmosis verify <image.osmoimg> --policy-file <policy.json>"
    if "untrusted artifact" in lower or "trusted artifact required" in lower:
        return f"osmosis verify {image} --require-trust"

    if cmd_name:
        return f"osmosis {cmd_name} --help"
    return "osmosis --help"


def _warn_insecure_output(command_name: str) -> None:
    _print_step(
        "⚠",
        f"{command_name} is creating an unsigned artifact (--no-trust-sign). "
        "This is insecure for production trust workflows.",
    )


def _auto_register_enabled() -> bool:
    return _truthy_env(os.environ.get("OSMOSIS_AUTO_REGISTER", "1"))


def _packaged_osmoimg_icon_path() -> Optional[str]:
    icon = Path(__file__).resolve().parent / "assets" / "osmoimg_icon.ico"
    if icon.exists():
        return str(icon)
    return None


def _read_winreg_default(hive: Any, path: str) -> str:
    import winreg

    with winreg.OpenKey(hive, path) as key:
        value, _ = winreg.QueryValueEx(key, "")
    return str(value or "").strip()


def _is_windows_osmoimg_registered() -> bool:
    if platform.system() != "Windows":
        return False
    try:
        import winreg
        from model_osmosis.format.file_format import OSMOIMG_EXTENSION, OSMOIMG_PROG_ID

        prog_id = _read_winreg_default(
            winreg.HKEY_CURRENT_USER,
            rf"Software\Classes\{OSMOIMG_EXTENSION}",
        )
        if prog_id != OSMOIMG_PROG_ID:
            return False

        icon = _read_winreg_default(
            winreg.HKEY_CURRENT_USER,
            rf"Software\Classes\{OSMOIMG_PROG_ID}\DefaultIcon",
        )
        return bool(icon)
    except Exception:
        return False


def _ensure_windows_auto_register(command: Optional[str]) -> None:
    # pip/wheel install hooks are unreliable across installers; perform a
    # silent first-run registration on Windows instead.
    if platform.system() != "Windows":
        return
    if not _auto_register_enabled():
        return
    if command in {None, "register"}:
        return
    if _is_windows_osmoimg_registered():
        return
    try:
        from model_osmosis.format.file_format import register_file_format

        register_file_format(verbose=False, icon_path=_packaged_osmoimg_icon_path())
    except Exception:
        return


def _utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _load_policy_file(path: Optional[str]) -> Dict[str, Any]:
    if not path:
        return {}
    p = Path(path)
    if not p.exists():
        raise FileNotFoundError(f"Policy file not found: {p}")
    try:
        data = json.loads(p.read_text(encoding="utf-8"))
    except Exception as exc:
        raise ValueError(f"Invalid policy file JSON: {p}") from exc
    if not isinstance(data, dict):
        raise ValueError("Policy file must contain a JSON object")
    return data


def _evaluate_policy(verify_result: Dict[str, Any], policy: Dict[str, Any]) -> Dict[str, Any]:
    if not policy:
        return {"checked": False, "passed": True, "violations": []}

    violations: List[str] = []
    header_version = int(verify_result.get("header_version", 1) or 1)
    trust_status = str(verify_result.get("trust_status", "unsigned"))
    signer_id = str(verify_result.get("signer_id") or "")
    key_id = str(verify_result.get("key_id") or "")
    provenance_mode = str(verify_result.get("provenance_mode") or "")
    repro_mode = str(verify_result.get("repro_mode") or "")
    transparency_status = str(verify_result.get("transparency_status") or "")

    min_header_version = policy.get("min_header_version")
    if min_header_version is not None and header_version < int(min_header_version):
        violations.append(
            f"header_version={header_version} is below min_header_version={int(min_header_version)}"
        )

    required_trust_status = policy.get("required_trust_status")
    if required_trust_status and trust_status != str(required_trust_status):
        violations.append(
            f"trust_status={trust_status} does not match required_trust_status={required_trust_status}"
        )

    allowed_signer_ids = policy.get("allowed_signer_ids") or []
    if allowed_signer_ids and signer_id not in [str(x) for x in allowed_signer_ids]:
        violations.append(f"signer_id={signer_id or 'none'} is not in allowed_signer_ids")

    allowed_key_ids = policy.get("allowed_key_ids") or []
    if allowed_key_ids and key_id not in [str(x) for x in allowed_key_ids]:
        violations.append(f"key_id={key_id or 'none'} is not in allowed_key_ids")

    required_provenance_mode = policy.get("required_provenance_mode")
    if required_provenance_mode and provenance_mode != str(required_provenance_mode):
        violations.append(
            f"provenance_mode={provenance_mode or 'unknown'} does not match required_provenance_mode={required_provenance_mode}"
        )

    required_repro_mode = policy.get("required_repro_mode")
    if required_repro_mode and repro_mode != str(required_repro_mode):
        violations.append(
            f"repro_mode={repro_mode or 'unknown'} does not match required_repro_mode={required_repro_mode}"
        )

    if bool(policy.get("require_transparency", False)):
        if transparency_status not in {"local_found", "published", "local_only"}:
            violations.append(
                f"transparency_status={transparency_status or 'none'} does not satisfy require_transparency=true"
            )

    return {
        "checked": True,
        "passed": len(violations) == 0,
        "violations": violations,
    }


def _emit_audit_event(
    *,
    action: str,
    image_path: str,
    verify_result: Dict[str, Any],
    policy_result: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    from model_osmosis.format.trust import append_audit_event

    return append_audit_event(
        action=action,
        image_path=image_path,
        trust_result=verify_result,
        policy_result=policy_result,
    )


def _has_non_empty_allowlist(policy: Dict[str, Any]) -> bool:
    signer_ids = [str(v).strip() for v in (policy.get("allowed_signer_ids") or [])]
    key_ids = [str(v).strip() for v in (policy.get("allowed_key_ids") or [])]
    return any(signer_ids) or any(key_ids)


def _check_trust_for_action(
    image_path: str,
    *,
    action: str,
    require_trust: bool,
    policy_file: Optional[str] = None,
    trust_profile: str = "standard",
) -> dict:
    from model_osmosis.format.image import OsmoImage

    normalized_profile = _resolve_trust_profile(trust_profile)
    effective_require_trust = _effective_require_trust(require_trust, normalized_profile)
    base_cmd = f"osmosis {action} {image_path}"
    production_policy_cmd = (
        f"{base_cmd} --trust-profile production --policy-file <policy.json>"
    )

    raw_policy = _load_policy_file(policy_file)
    if normalized_profile == "production" and not policy_file:
        raise _error_with_next_command(
            "Production trust profile requires a policy file with signer allowlists.",
            production_policy_cmd,
        )

    verify_result = OsmoImage.verify(image_path, require_trust=False)
    policy = _policy_with_trust_profile(
        raw_policy,
        normalized_profile,
    )
    if normalized_profile == "production" and not _has_non_empty_allowlist(policy):
        raise _error_with_next_command(
            "Production trust profile requires non-empty allowed_signer_ids or allowed_key_ids in policy.",
            production_policy_cmd,
        )

    policy_result = _evaluate_policy(verify_result, policy)
    verify_result["policy_result"] = policy_result
    verify_result["trust_profile"] = normalized_profile
    _emit_audit_event(
        action=action,
        image_path=image_path,
        verify_result=verify_result,
        policy_result=policy_result,
    )
    if policy_result.get("checked") and not policy_result.get("passed"):
        remedy_cmd = (
            f"{base_cmd} --trust-profile {normalized_profile} --policy-file {policy_file}"
            if policy_file
            else production_policy_cmd
        )
        raise _error_with_next_command(
            f"{action} blocked by policy. Violations: "
            + "; ".join(policy_result.get("violations", [])),
            remedy_cmd,
        )

    status = verify_result.get("trust_status", "unsigned")
    if status != "trusted":
        msg = (
            f"{action} on untrusted artifact ({status}). "
            f"Reason: {verify_result.get('reason', 'unknown')}. "
            f"Run `osmosis verify {image_path}`."
        )
        if effective_require_trust:
            raise _error_with_next_command(
                msg,
                f"osmosis verify {image_path} --require-trust",
            )
        _print_step("⚠", msg)
    return verify_result


def _run_repro_check(image_path: str, verify_result: Dict[str, Any]) -> Dict[str, Any]:
    from model_osmosis.format.image import OsmoImage
    import hashlib
    import json

    status = str(verify_result.get("trust_status", "unsigned"))
    if status == "tampered":
        return {
            "checked": True,
            "passed": False,
            "mode": "identity-validation",
            "reason": "Cannot run repro check on tampered artifact",
            "expected_payload_sha256": verify_result.get("payload_sha256", ""),
            "rebuilt_payload_sha256": "",
        }

    source = OsmoImage.load(image_path)

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

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

    manifest_sha = _sha256_bytes(_canon(source.osmo_model.manifest.to_dict()))
    tensor_index: List[Dict[str, Any]] = []
    for name in sorted(source.osmo_model.parameters.keys()):
        arr = source.osmo_model.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": _sha256_bytes(raw),
                "raw_len": len(raw),
            }
        )
    tensor_index_sha = _sha256_bytes(_canon(tensor_index))
    expected_payload_sha = str(verify_result.get("payload_sha256", "") or "")
    artifact_descriptor = {
        "payload_sha256": expected_payload_sha,
        "manifest_sha256": manifest_sha,
        "tensor_index_sha256": tensor_index_sha,
        "model_fingerprint": source.osmo_model.fingerprint,
        "payload_codec": source.runtime.payload_codec,
        "compression_profile": source.runtime.compression_profile,
    }
    artifact_id = _sha256_bytes(_canon(artifact_descriptor))

    expected_manifest = str(getattr(source.runtime, "manifest_sha256", "") or "")
    expected_tensor_index = str(getattr(source.runtime, "tensor_index_sha256", "") or "")
    expected_artifact_id = str(getattr(source.runtime, "artifact_id", "") or "")

    checks = {
        "manifest_sha256_match": bool(expected_manifest and expected_manifest == manifest_sha),
        "tensor_index_sha256_match": bool(expected_tensor_index and expected_tensor_index == tensor_index_sha),
        "artifact_id_match": bool(expected_artifact_id and expected_artifact_id == artifact_id),
    }
    passed = all(checks.values())
    return {
        "checked": True,
        "passed": passed,
        "mode": "identity-validation",
        "reason": "" if passed else "Artifact identity fields do not match recomputed values",
        "expected_payload_sha256": expected_payload_sha,
        "manifest_sha256": manifest_sha,
        "tensor_index_sha256": tensor_index_sha,
        "artifact_id": artifact_id,
        "checks": checks,
    }


# ─── MERGE ─────────────────────────────────────────────────────────────────

def cmd_merge(args):
    """Merge multiple models into one."""
    print(_logo())
    print(f"\n  {Colors.BOLD}MERGE{Colors.RESET}\n")

    models = args.models
    output = args.output or "merged_model"
    strategy = args.strategy
    weights = args.weights

    _print_info("Models:", ", ".join(str(m) for m in models))
    _print_info("Strategy:", strategy)
    _print_info("Strategy engine:", _resolve_strategy_engine_mode(getattr(args, "strategy_engine", "auto")))
    _print_info("Cross opt stage:", str(getattr(args, "cross_opt_stage", "auto")))
    _print_info("Output:", output)
    if weights:
        _print_info("Weights:", ", ".join(f"{w:.2f}" for w in weights))
    print()

    from model_osmosis import OsmosisPipeline, OsmosisConfig
    from model_osmosis.utils.determinism import set_global_seed

    set_global_seed(args.seed)
    config = OsmosisConfig(verbose=True, seed=args.seed)
    pipeline = OsmosisPipeline(config)
    loaded = [_load_model(m) for m in models]
    sample_data = None
    if args.cross and args.samples:
        sample_data = [s.strip() for s in args.samples.split("|") if s.strip()]

    merge_result = _run_merge_backend(
        pipeline,
        loaded_models=loaded,
        strategy=strategy,
        weights=weights,
        sample_data=sample_data,
        source_refs=models,
        force_cross=bool(args.cross),
        epochs=args.epochs or 3,
        seed=args.seed,
        strategy_engine_mode=getattr(args, "strategy_engine", "auto"),
        cross_opt_stage=getattr(args, "cross_opt_stage", "auto"),
        return_metadata=True,
    )
    if isinstance(merge_result, tuple) and len(merge_result) == 2 and isinstance(merge_result[1], dict):
        result = merge_result[0]
    else:
        result = merge_result

    if isinstance(result, tuple):
        model, tokenizer = result
        os.makedirs(output, exist_ok=True)
        model.save_pretrained(output)
        if tokenizer is not None and hasattr(tokenizer, "save_pretrained"):
            tokenizer.save_pretrained(output)
    else:
        _save_model(result, output)

    # Optionally obfuscate
    if args.obfuscate:
        _print_step("🔒", "Obfuscating merged model...")
        from model_osmosis.obfuscation import ObfuscationConfig
        obf_config = ObfuscationConfig(preserve_quality=True)
        if isinstance(result, tuple):
            result = pipeline.obfuscate(result[0], obfuscation_config=obf_config)
        else:
            result = pipeline.obfuscate(result, obfuscation_config=obf_config)
        _save_model(result, output)
        _print_step("✓", "Weights obfuscated (quality preserved)")

    print()
    _print_step("✓", f"Saved to {Colors.BOLD}{output}{Colors.RESET}")
    print()


# ─── OBFUSCATE ─────────────────────────────────────────────────────────────

def cmd_obfuscate(args):
    """Obfuscate model weights to prevent tracing."""
    print(_logo())
    print(f"\n  {Colors.BOLD}OBFUSCATE{Colors.RESET}\n")

    model_path = args.model
    output = args.output or f"{model_path}_obfuscated"

    _print_info("Input:", model_path)
    _print_info("Output:", output)
    _print_info("Quality:", "preserved" if args.preserve_quality else "standard")
    print()

    from model_osmosis import OsmosisPipeline, OsmosisConfig
    from model_osmosis.obfuscation import ObfuscationConfig

    config = OsmosisConfig(verbose=True)
    pipeline = OsmosisPipeline(config)

    obf_config = ObfuscationConfig(
        preserve_quality=args.preserve_quality,
        seed=args.seed,
        num_rounds=args.rounds,
    )

    model = _load_model(model_path)
    _print_step("🔒", "Obfuscating weights...")

    result = pipeline.obfuscate(model, obfuscation_config=obf_config)

    _save_model(result, output)
    _print_step("✓", f"Saved to {Colors.BOLD}{output}{Colors.RESET}")
    print()


# ─── COMPARE ──────────────────────────────────────────────────────────────

def cmd_compare(args):
    """Compare two models' outputs."""
    print(_logo())
    print(f"\n  {Colors.BOLD}COMPARE{Colors.RESET}\n")

    _print_info("Model A:", args.model_a)
    _print_info("Model B:", args.model_b)
    print()

    import numpy as np
    from model_osmosis import OsmosisPipeline, OsmosisConfig

    config = OsmosisConfig(verbose=False)
    pipeline = OsmosisPipeline(config)

    model_a = _load_model(args.model_a)
    model_b = _load_model(args.model_b)

    sample = np.random.rand(args.samples, 10).astype(np.float32)
    result = pipeline.compare(model_a, model_b, sample)

    # Pretty print results
    cos = result["cosine_similarity"]
    mse = result["mse"]
    corr = result["correlation"]
    maxe = result["max_error"]

    cos_color = Colors.GREEN if cos > 0.99 else Colors.YELLOW if cos > 0.9 else Colors.RED
    mse_color = Colors.GREEN if mse < 1e-6 else Colors.YELLOW if mse < 1e-3 else Colors.RED

    print(f"  {Colors.BOLD}Results:{Colors.RESET}")
    print(f"  ┌──────────────────────────────────┐")
    print(f"  │ Cosine Similarity  {cos_color}{cos:>12.6f}{Colors.RESET}  │")
    print(f"  │ MSE                {mse_color}{mse:>12.2e}{Colors.RESET}  │")
    print(f"  │ Correlation        {corr:>12.6f}  │")
    print(f"  │ Max Error          {maxe:>12.2e}  │")
    print(f"  │ Samples            {result['num_samples']:>12d}  │")
    print(f"  └──────────────────────────────────┘")

    if cos > 0.9999:
        print(f"\n  {Colors.GREEN}✓ Models are functionally identical{Colors.RESET}")
    elif cos > 0.99:
        print(f"\n  {Colors.YELLOW}~ Models are very similar{Colors.RESET}")
    else:
        print(f"\n  {Colors.RED}✗ Models produce different outputs{Colors.RESET}")
    print()


# ─── BENCHMARK ────────────────────────────────────────────────────────────

def cmd_benchmark(args):
    """Benchmark merge strategies on given models."""
    print(_logo())
    print(f"\n  {Colors.BOLD}BENCHMARK{Colors.RESET}\n")

    models = args.models
    _print_info("Models:", ", ".join(str(m) for m in models))
    print()

    import numpy as np
    from model_osmosis import OsmosisPipeline, OsmosisConfig

    config = OsmosisConfig(verbose=False)
    pipeline = OsmosisPipeline(config)

    loaded = [_load_model(m) for m in models]
    eval_data = np.random.rand(50, 10).astype(np.float32)

    _print_step("⏱", "Running benchmark...")
    print()

    results = pipeline.benchmark(loaded, eval_data)

    print(f"  {Colors.BOLD}{'Rank':<6}{'Strategy':<18}{'Score':<14}{'Weights'}{Colors.RESET}")
    print(f"  {'─'*55}")
    for i, r in enumerate(results):
        rank = f"#{i+1}"
        color = Colors.GREEN if i == 0 else Colors.DIM
        weights_str = ", ".join(f"{w:.2f}" for w in r["weights"])
        print(f"  {color}{rank:<6}{r['strategy']:<18}{r['score']:<14.6f}[{weights_str}]{Colors.RESET}")

    if results:
        best = results[0]
        print(f"\n  {Colors.GREEN}✓ Best: {Colors.BOLD}{best['strategy']}{Colors.RESET}"
              f" (score: {best['score']:.6f})")
    print()


# ─── INFO ──────────────────────────────────────────────────────────────────

def cmd_info(args):
    """Show model information."""
    print(_logo())
    print(f"\n  {Colors.BOLD}INFO{Colors.RESET}\n")

    model_path = args.model
    _print_info("Model:", model_path)
    print()

    model = _load_model(model_path)

    if hasattr(model, "parameters"):
        # PyTorch model
        import torch
        total_params = sum(p.numel() for p in model.parameters())
        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)

        print(f"  {Colors.BOLD}Architecture:{Colors.RESET}")
        _print_info("Type:", type(model).__name__)
        _print_info("Parameters:", f"{total_params:,}")
        _print_info("Trainable:", f"{trainable:,}")
        _print_info("Size (est):", f"{total_params * 4 / 1024**2:.1f} MB (fp32)")

        print(f"\n  {Colors.BOLD}Layers:{Colors.RESET}")
        for name, param in list(model.named_parameters())[:15]:
            print(f"    {Colors.DIM}{name:40s}{Colors.RESET} {list(param.shape)}")
        if sum(1 for _ in model.named_parameters()) > 15:
            remaining = sum(1 for _ in model.named_parameters()) - 15
            print(f"    {Colors.DIM}... and {remaining} more{Colors.RESET}")
    elif hasattr(model, "config"):
        # HuggingFace model
        print(f"  {Colors.BOLD}Config:{Colors.RESET}")
        config = model.config
        for key in ["model_type", "hidden_size", "num_hidden_layers",
                     "num_attention_heads", "vocab_size"]:
            if hasattr(config, key):
                _print_info(key + ":", str(getattr(config, key)))
    else:
        _print_info("Type:", type(model).__name__)

    print()


# ─── BUILD ─────────────────────────────────────────────────────────────────

def cmd_build(args):
    """Build an OsmoImage from any model."""
    print(_logo())
    print(f"\n  {Colors.BOLD}BUILD{Colors.RESET}\n")

    signer_backend = getattr(args, "signer_backend", "local-keyring")
    compression_engine = _normalize_mode_arg(
        getattr(args, "compression_engine", "auto"),
        COMPRESSION_ENGINE_MODES,
        "auto",
    )
    _print_info("Source:", args.model)
    _print_info("Output:", args.output)
    _print_info("Compression profile:", args.compression_profile)
    _print_info("Compression engine:", _resolve_compression_engine_mode(compression_engine))
    _print_info("Trust sign:", "yes" if args.trust_sign else "no")
    _print_info("Signer backend:", signer_backend)
    _print_info("Provenance mode:", args.provenance)
    _print_info("Repro mode:", args.repro_mode)
    _print_info("Encrypt:", "✓ always (lossless)")
    _print_info("Compress:", "✓ always (lossless)")
    print()
    if not args.trust_sign:
        _warn_insecure_output("build")
        print()

    from model_osmosis.format.image import OsmoImage

    _print_step("📦", f"Building image from {args.model}...")
    image = OsmoImage.build(
        model=args.model,
        name=args.name or "",
        compression_profile=args.compression_profile,
        compression_engine=compression_engine,
        trust_sign=args.trust_sign,
        provenance_mode=args.provenance,
        repro_mode=args.repro_mode,
        signer_id=args.signer_id,
        signer_backend=signer_backend,
    )

    _print_step("💾", f"Saving to {args.output}...")
    result_path = image.save(
        args.output,
        compression_profile=args.compression_profile,
        compression_engine=compression_engine,
        trust_sign=args.trust_sign,
        provenance_mode=args.provenance,
        repro_mode=args.repro_mode,
        signer_id=args.signer_id,
        signer_backend=signer_backend,
    )

    # Show summary
    print()
    print(f"  {Colors.BOLD}Image Summary:{Colors.RESET}")
    _print_info("Name:", image.name)
    _print_info("Parameters:", f"{image.parameter_count:,}")
    _print_info("Size:", f"{image.runtime.original_size_mb:.1f} MB")
    _print_info("Payload codec:", image.runtime.payload_codec)
    _print_info("Compression ratio:", f"{image.runtime.compression_ratio:.2f}x")
    _print_info("Encrypted:", "✓ yes" if image.is_encrypted else "✗ no")
    _print_info("Trust status:", image.runtime.trust_status)
    _print_info("Signer:", image.runtime.signer_id or "none")
    _print_info("Signer backend:", image.runtime.signer_backend or "none")
    _print_info("Transparency:", image.runtime.transparency_status or "n/a")
    print()
    _print_step("✓", f"Image saved to {Colors.BOLD}{result_path}{Colors.RESET}")
    print()


# ─── RUN ───────────────────────────────────────────────────────────────────

def cmd_run(args):
    """Run inference from an OsmoImage."""
    print(_logo())
    print(f"\n  {Colors.BOLD}RUN{Colors.RESET}\n")

    _print_info("Image:", args.image)
    if args.prompt:
        _print_info("Prompt:", args.prompt[:60] + ("..." if len(args.prompt) > 60 else ""))
    _print_info("Runtime profile:", args.runtime_profile)
    _print_info("Device:", args.device)
    _print_info("Runtime quantization:", args.runtime_quantization)
    _print_info("Runtime cache:", "on" if args.runtime_cache else "off")
    _print_info("Include prompt:", "yes" if args.include_prompt else "no")
    require_trust_flag = bool(getattr(args, "require_trust", False))
    trust_profile = _resolve_trust_profile(getattr(args, "trust_profile", "standard"))
    require_trust = _effective_require_trust(require_trust_flag, trust_profile)
    _print_info("Trust profile:", trust_profile)
    _print_info("Require trust:", "yes" if require_trust else "no")
    _print_info("Policy file:", getattr(args, "policy_file", None) or "none")
    print()

    from model_osmosis.format.image import OsmoImage
    import numpy as np

    _check_trust_for_action(
        args.image,
        action="run",
        require_trust=require_trust,
        policy_file=getattr(args, "policy_file", None),
        trust_profile=trust_profile,
    )

    _print_step("📥", "Loading image...")
    image = OsmoImage.load(args.image)
    print(f"  {Colors.DIM}{image.summary()}{Colors.RESET}")
    print()

    if args.prompt:
        _print_step("🧠", "Running inference...")
        output = image.run(
            prompt=args.prompt,
            max_tokens=args.max_tokens,
            temperature=args.temperature,
            include_prompt=args.include_prompt,
            runtime_profile=args.runtime_profile,
            device=args.device,
            runtime_cache=args.runtime_cache,
            runtime_quantization=args.runtime_quantization,
        )
        print()
        print(f"  {Colors.BOLD}Output:{Colors.RESET}")
        print(f"  {output}")
    elif args.input:
        _print_step("🧠", f"Running inference on {args.input}...")
        data = np.load(args.input)
        output = image.run(
            input_data=data,
            runtime_profile=args.runtime_profile,
            device=args.device,
            runtime_cache=args.runtime_cache,
            runtime_quantization=args.runtime_quantization,
        )
        if args.output_file:
            np.save(args.output_file, output)
            _print_step("💾", f"Saved to {args.output_file}")
        else:
            print(f"  Output shape: {output.shape}")
    else:
        _print_error("Must provide --prompt or --input")
        return

    print()


# ─── INSPECT ───────────────────────────────────────────────────────────────

def cmd_inspect(args):
    """Inspect an OsmoImage without loading weights."""
    print(_logo())
    print(f"\n  {Colors.BOLD}INSPECT{Colors.RESET}\n")

    from model_osmosis.format.image import OsmoImage

    info = OsmoImage.inspect(args.image)

    # Model
    print(f"  {Colors.BOLD}Model{Colors.RESET}")
    _print_info("Name:", info["name"] or "(unnamed)")
    _print_info("Description:", info.get("description", "") or "(none)")
    _print_info("Parameters:", f"{info['parameter_count']:,}")
    _print_info("License:", info.get("license", "unknown"))
    print()

    # Architecture
    print(f"  {Colors.BOLD}Architecture{Colors.RESET}")
    arch = info["architecture"]
    _print_info("Type:", arch.get("type", "unknown"))
    _print_info("Hidden size:", str(arch.get("hidden_size", "?")))
    _print_info("Layers:", str(arch.get("num_layers", "?")))
    _print_info("Attention heads:", str(arch.get("num_attention_heads", "?")))
    _print_info("Vocab size:", str(arch.get("vocab_size", "?")))
    print()

    # Runtime
    print(f"  {Colors.BOLD}Runtime{Colors.RESET}")
    rt = info["runtime"]
    _print_info("Framework:", rt["source_framework"])
    _print_info("Model type:", rt["model_type"])
    _print_info("Min RAM:", f"{rt['min_ram_gb']:.1f} GB")
    _print_info("CPU support:", "✓" if rt["supports_cpu"] else "✗")
    encrypted_color = Colors.GREEN if rt["is_encrypted"] else Colors.YELLOW
    _print_info("Encrypted:", f"{encrypted_color}{'✓ yes' if rt['is_encrypted'] else '✗ no'}{Colors.RESET}")
    _print_info("Compression profile:", rt.get("compression_profile", "legacy"))
    _print_info("Payload codec:", rt.get("payload_codec", "zip-deflate"))
    _print_info("Compression engine:", rt.get("compression_engine_version", "legacy-zip-v1"))
    optimizer_meta = (rt.get("compression_stats", {}) or {}).get("optimizer", {})
    if isinstance(optimizer_meta, dict) and optimizer_meta:
        _print_info("Optimizer mode:", str(optimizer_meta.get("optimizer_mode", "auto")))
        _print_info("Optimizer profile:", str(optimizer_meta.get("profile_id", "n/a")))
        _print_info("Optimizer pack:", str(optimizer_meta.get("pack_version", "n/a")))
        _print_info("Optimizer selector:", str(optimizer_meta.get("selector_version", "n/a")))
        _print_info("Canary status:", str(optimizer_meta.get("canary_status", "shadow")))
        _print_info("Rollback reason:", str(optimizer_meta.get("rollback_reason", "")) or "none")
        _print_info(
            "Selection latency:",
            f"{float(optimizer_meta.get('selection_latency_ms', 0.0) or 0.0):.4f} ms",
        )
    _print_info("Trust status:", rt.get("trust_status", "unsigned"))
    _print_info("Signer ID:", rt.get("signer_id") or "none")
    _print_info("Key ID:", rt.get("key_id") or "none")
    _print_info("Signature alg:", rt.get("signature_alg") or "none")
    _print_info("Signer backend:", rt.get("signer_backend") or "local-keyring")
    _print_info("Payload SHA256:", rt.get("payload_sha256") or "n/a")
    _print_info("Artifact ID:", rt.get("artifact_id") or "n/a")
    _print_info("Manifest SHA256:", rt.get("manifest_sha256") or "n/a")
    _print_info("Tensor index SHA256:", rt.get("tensor_index_sha256") or "n/a")
    _print_info("Verified at:", rt.get("verified_at") or "n/a")
    _print_info("Transparency:", rt.get("transparency_status") or "n/a")
    _print_info("Provenance mode:", rt.get("provenance_mode", "private"))
    _print_info("Repro mode:", rt.get("repro_mode", "strict"))
    if rt.get("hf_model_id"):
        _print_info("HF Model:", rt["hf_model_id"])
    print()

    # Image
    print(f"  {Colors.BOLD}Image{Colors.RESET}")
    img = info["image"]
    _print_info("File size:", f"{img['file_size_mb']:.1f} MB")
    _print_info("Original size:", f"{img['original_size_mb']:.1f} MB")
    _print_info("Compression (stored):", f"{img.get('compression_ratio_stored', img['compression_ratio']):.2f}x")
    _print_info("Compression (effective):", f"{img.get('effective_compression_ratio', img['compression_ratio']):.2f}x")
    _print_info("Built at:", img.get("built_at", "unknown"))
    print()


# ─── EXPORT ────────────────────────────────────────────────────────────────

def cmd_export(args):
    """Export an OsmoImage to any framework."""
    print(_logo())
    print(f"\n  {Colors.BOLD}EXPORT{Colors.RESET}\n")

    _print_info("Image:", args.image)
    _print_info("Format:", args.format)
    _print_info("Output:", args.output or "(in memory)")
    require_trust_flag = bool(getattr(args, "require_trust", False))
    trust_profile = _resolve_trust_profile(getattr(args, "trust_profile", "standard"))
    require_trust = _effective_require_trust(require_trust_flag, trust_profile)
    _print_info("Trust profile:", trust_profile)
    _print_info("Require trust:", "yes" if require_trust else "no")
    _print_info("Policy file:", getattr(args, "policy_file", None) or "none")
    print()

    from model_osmosis.format.image import OsmoImage

    _check_trust_for_action(
        args.image,
        action="export",
        require_trust=require_trust,
        policy_file=getattr(args, "policy_file", None),
        trust_profile=trust_profile,
    )

    _print_step("📥", "Loading image...")
    image = OsmoImage.load(args.image)

    _print_step("🔄", f"Exporting to {args.format}...")
    model = image.export(
        target_format=args.format,
        output_path=args.output,
    )

    if args.output:
        _print_step("✓", f"Exported to {Colors.BOLD}{args.output}{Colors.RESET}")
    else:
        _print_step("✓", f"Model loaded in memory ({type(model).__name__})")
    print()


# ─── REGISTER ─────────────────────────────────────────────────────────────

def cmd_register(args):
    """Register the .osmoimg file format with the OS."""
    print(_logo())
    print(f"\n  {Colors.BOLD}REGISTER{Colors.RESET}\n")

    from model_osmosis.format.file_format import (
        register_file_format,
        OSMOIMG_EXTENSION,
        OSMOIMG_MIME_TYPE,
        OSMOIMG_PROG_ID,
    )

    _print_info("Extension:", OSMOIMG_EXTENSION)
    _print_info("MIME Type:", OSMOIMG_MIME_TYPE)
    _print_info("ProgID:", OSMOIMG_PROG_ID)
    _print_info("Icon:", args.icon or "(auto)")
    print()

    if args.unregister:
        from model_osmosis.format.file_format import unregister_file_format
        _print_step("🗑", "Unregistering .osmoimg file format...")
        success = unregister_file_format(verbose=True)
    else:
        _print_step("📝", "Registering .osmoimg file format...")
        success = register_file_format(verbose=True, icon_path=args.icon)

    if success:
        if not args.unregister:
            print()
            print(f"  {Colors.BOLD}What this does:{Colors.RESET}")
            print(f"    • .osmoimg files are now recognized by your OS")
            print(f"    • Double-click to run with: osmosis run <file>")
            print(f"    • Right-click for: Inspect, Export options")
    print()


# ─── IMAGES ───────────────────────────────────────────────────────────────

def cmd_images(args):
    """List all .osmoimg files in a directory."""
    print(_logo())
    print(f"\n  {Colors.BOLD}IMAGES{Colors.RESET}\n")

    search_dir = Path(args.directory)
    if not search_dir.is_dir():
        _print_error(f"Directory not found: {search_dir}")
        return

    from model_osmosis.format.image import OsmoImage
    from model_osmosis.format.file_format import OSMOIMG_EXTENSION

    files = sorted(search_dir.rglob(f"*{OSMOIMG_EXTENSION}"))
    if not files:
        print(f"  {Colors.DIM}No .osmoimg files found in {search_dir}{Colors.RESET}")
        print()
        return

    # Table header
    print(f"  {Colors.BOLD}{'NAME':<20}{'SIZE':>10}{'PARAMS':>14}{'ENCRYPTED':>12}{'BUILT':<12}{Colors.RESET}")
    print(f"  {'─'*68}")

    for f in files:
        try:
            info = OsmoImage.inspect(f)
            name = info["name"] or f.stem
            size_mb = info["image"]["file_size_mb"]
            params = info["parameter_count"]
            encrypted = "✓" if info["runtime"]["is_encrypted"] else "✗"
            built = info["image"].get("built_at", "unknown")[:10]

            if size_mb >= 1024:
                size_str = f"{size_mb/1024:.1f} GB"
            else:
                size_str = f"{size_mb:.1f} MB"

            print(f"  {name:<20}{size_str:>10}{params:>14,}{encrypted:>12}  {built}")
        except Exception as e:
            print(f"  {f.name:<20}{Colors.RED}(error: {e}){Colors.RESET}")

    print(f"\n  {Colors.DIM}{len(files)} image(s) found{Colors.RESET}")
    print()


# ─── PULL ─────────────────────────────────────────────────────────────────

def cmd_pull(args):
    """Download a HuggingFace model and save as .osmoimg."""
    print(_logo())
    print(f"\n  {Colors.BOLD}PULL{Colors.RESET}\n")

    model_id = args.model
    signer_backend = getattr(args, "signer_backend", "local-keyring")
    compression_engine = _normalize_mode_arg(
        getattr(args, "compression_engine", "auto"),
        COMPRESSION_ENGINE_MODES,
        "auto",
    )
    output = args.output or f"{model_id.split('/')[-1]}.osmoimg"

    _print_info("Model:", model_id)
    _print_info("Output:", output)
    _print_info("Compression profile:", args.compression_profile)
    _print_info("Compression engine:", _resolve_compression_engine_mode(compression_engine))
    _print_info("Trust sign:", "yes" if args.trust_sign else "no")
    _print_info("Signer backend:", signer_backend)
    _print_info("Provenance mode:", args.provenance)
    _print_info("Repro mode:", args.repro_mode)
    print()
    if not args.trust_sign:
        _warn_insecure_output("pull")
        print()

    from model_osmosis.format.image import OsmoImage

    _print_step("📥", f"Pulling {model_id} from HuggingFace...")
    image = OsmoImage.build(
        model=model_id,
        name=model_id.split("/")[-1],
        compression_profile=args.compression_profile,
        compression_engine=compression_engine,
        trust_sign=args.trust_sign,
        provenance_mode=args.provenance,
        repro_mode=args.repro_mode,
        signer_id=args.signer_id,
        signer_backend=signer_backend,
    )

    _print_step("💾", f"Saving to {output}...")
    result_path = image.save(
        output,
        compression_profile=args.compression_profile,
        compression_engine=compression_engine,
        trust_sign=args.trust_sign,
        provenance_mode=args.provenance,
        repro_mode=args.repro_mode,
        signer_id=args.signer_id,
        signer_backend=signer_backend,
    )

    print()
    print(f"  {Colors.BOLD}Image Summary:{Colors.RESET}")
    _print_info("Name:", image.name)
    _print_info("Parameters:", f"{image.parameter_count:,}")
    _print_info("Size:", f"{image.runtime.original_size_mb:.1f} MB")
    _print_info("Payload codec:", image.runtime.payload_codec)
    _print_info("Compression ratio:", f"{image.runtime.compression_ratio:.2f}x")
    _print_info("Encrypted:", "✓ yes")
    _print_info("Trust status:", image.runtime.trust_status)
    _print_info("Signer:", image.runtime.signer_id or "none")
    _print_info("Signer backend:", image.runtime.signer_backend or "none")
    _print_info("Transparency:", image.runtime.transparency_status or "n/a")
    print()
    _print_step("✓", f"Saved to {Colors.BOLD}{result_path}{Colors.RESET}")
    print()


# ─── DEMO ──────────────────────────────────────────────────────────────────

def _is_module_available(module_name: str) -> bool:
    return importlib.util.find_spec(module_name) is not None


def _require_demo_dependencies(compression_profile: str) -> None:
    missing = []
    for module_name in ("torch", "transformers", "huggingface_hub", "cryptography"):
        if not _is_module_available(module_name):
            missing.append(module_name)

    if compression_profile == "lossless-max" and not _is_module_available("zstandard"):
        missing.append("zstandard")

    if missing:
        missing_str = ", ".join(sorted(set(missing)))
        raise RuntimeError(
            f"Missing required dependencies for `osmosis demo`: {missing_str}. "
            "Install with: pip install \"osmosis[huggingface,compression]\""
        )


def _validate_demo_inspect_contract(info: dict) -> None:
    runtime = info.get("runtime", {})
    image = info.get("image", {})
    required_runtime = [
        "source_framework",
        "model_type",
        "min_ram_gb",
        "supports_cpu",
        "is_encrypted",
        "payload_codec",
        "compression_profile",
        "trust_status",
    ]
    missing_runtime = [k for k in required_runtime if k not in runtime]
    if missing_runtime:
        raise RuntimeError(
            "Demo inspect contract failed; missing runtime fields: "
            + ", ".join(missing_runtime)
        )

    if image.get("file_size_mb", 0) <= 0:
        raise RuntimeError("Demo inspect contract failed; invalid file size")


def _check_hf_network(timeout_sec: float = 5.0) -> None:
    try:
        req = urllib.request.Request(
            "https://huggingface.co",
            method="GET",
            headers={"User-Agent": "model-osmosis-demo-preflight"},
        )
        with urllib.request.urlopen(req, timeout=timeout_sec) as resp:  # nosec B310
            if int(getattr(resp, "status", 200)) >= 400:
                raise RuntimeError(f"status={getattr(resp, 'status', 0)}")
    except Exception as exc:
        raise RuntimeError(
            "Network preflight failed for HuggingFace Hub. "
            "Check internet access or run behind a configured proxy."
        ) from exc


def _demo_api_smoke(
    image,
    *,
    prompt: str,
    max_tokens: int,
    temperature: float,
    timeout_sec: float = 8.0,
) -> Dict[str, Any]:
    from http.server import BaseHTTPRequestHandler, HTTPServer
    import threading

    class _DemoHandler(BaseHTTPRequestHandler):
        def do_GET(self):
            if self.path == "/health":
                payload = json.dumps({"status": "ok", "model": image.name}).encode("utf-8")
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(payload)
                return
            self.send_response(404)
            self.end_headers()

        def do_POST(self):
            if self.path != "/predict":
                self.send_response(404)
                self.end_headers()
                return
            content_len = int(self.headers.get("Content-Length", 0))
            body = self.rfile.read(content_len)
            request = json.loads(body.decode("utf-8") or "{}")
            result = image.run(
                prompt=str(request.get("prompt", "")),
                max_tokens=int(request.get("max_tokens", max_tokens)),
                temperature=float(request.get("temperature", temperature)),
            )
            payload = json.dumps({"output": result}).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(payload)

        def log_message(self, fmt, *args):
            return

    server = HTTPServer(("127.0.0.1", 0), _DemoHandler)
    port = int(server.server_port)
    base_url = f"http://127.0.0.1:{port}"
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()

    try:
        with urllib.request.urlopen(f"{base_url}/health", timeout=timeout_sec) as resp:  # nosec B310
            if int(getattr(resp, "status", 200)) != 200:
                raise RuntimeError("API health endpoint failed")

        req = urllib.request.Request(
            f"{base_url}/predict",
            method="POST",
            headers={"Content-Type": "application/json"},
            data=json.dumps(
                {
                    "prompt": prompt,
                    "max_tokens": int(max_tokens),
                    "temperature": float(temperature),
                }
            ).encode("utf-8"),
        )
        with urllib.request.urlopen(req, timeout=timeout_sec) as resp:  # nosec B310
            body = resp.read().decode("utf-8")
        data = json.loads(body or "{}")
        output = str(data.get("output", ""))
        if not output.strip():
            raise RuntimeError("API smoke output was empty")
        return {
            "api_url": base_url,
            "output": output,
            "health": "ok",
        }
    finally:
        server.shutdown()
        server.server_close()
        thread.join(timeout=timeout_sec)


def cmd_demo(args):
    """Run one-command first-time demo flow and print PASS/FAIL summary."""
    print(_logo())
    print(f"\n  {Colors.BOLD}DEMO{Colors.RESET}\n")

    start = time.perf_counter()
    output_path = Path(args.output)
    if output_path.suffix != ".osmoimg":
        output_path = output_path.with_suffix(".osmoimg")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    _print_info("Model:", args.model_id)
    _print_info("Output:", str(output_path))
    _print_info("Compression profile:", args.compression_profile)
    _print_info("Timeout:", f"{args.timeout_minutes} minute(s)")
    print()

    _print_step("🩺", "Running preflight checks...")
    _require_demo_dependencies(args.compression_profile)
    _check_hf_network(timeout_sec=5.0)

    from model_osmosis.format.image import OsmoImage

    try:
        _print_step("📦", "Building demo image...")
        image = OsmoImage.build(
            model=args.model_id,
            name=output_path.stem,
            compression_profile=args.compression_profile,
        )
        _print_step("💾", f"Saving image to {output_path}...")
        saved_path = image.save(output_path, compression_profile=args.compression_profile)
    except Exception as exc:
        msg = str(exc)
        if any(token in msg.lower() for token in ("connection", "timeout", "download", "http", "network")):
            raise RuntimeError(
                f"Failed downloading/loading demo model `{args.model_id}`. "
                f"Retry with: osmosis demo --model-id {args.model_id}"
            ) from exc
        if "transformers required" in msg.lower():
            raise RuntimeError(
                "Missing transformers dependency. "
                "Install with: pip install \"osmosis[huggingface,compression]\""
            ) from exc
        raise

    _print_step("🔍", "Inspecting generated image...")
    info = OsmoImage.inspect(saved_path)
    _validate_demo_inspect_contract(info)

    _print_step("🧠", "Running smoke prompt...")
    loaded = OsmoImage.load(saved_path)
    smoke_output = loaded.run(
        prompt=DEMO_DEFAULT_PROMPT,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
    )
    if not str(smoke_output).strip():
        raise RuntimeError("Smoke prompt failed: model returned empty output")
    _print_step("🌐", "Running API smoke...")
    api_smoke = _demo_api_smoke(
        loaded,
        prompt=DEMO_DEFAULT_PROMPT,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
    )

    elapsed = time.perf_counter() - start
    timeout_seconds = float(args.timeout_minutes) * 60.0
    if elapsed > timeout_seconds:
        raise RuntimeError(
            f"Demo exceeded timeout ({elapsed:.1f}s > {timeout_seconds:.1f}s)"
        )

    runtime = info["runtime"]
    payload_codec = runtime.get("payload_codec", "unknown")
    profile = runtime.get("compression_profile", "unknown")
    min_ram = float(runtime.get("min_ram_gb", 0.0))
    trust_status = runtime.get("trust_status", "unsigned")

    print()
    print(f"  {Colors.BOLD}Quick Demo Summary{Colors.RESET}")
    _print_info("Status:", f"{Colors.GREEN}PASS{Colors.RESET}")
    _print_info("Output path:", str(saved_path))
    _print_info("Min RAM:", f"{min_ram:.2f} GB")
    _print_info("Payload codec:", payload_codec)
    _print_info("Compression profile:", profile)
    _print_info("Trust status:", trust_status)
    _print_info("API URL:", api_smoke.get("api_url", "n/a"))
    _print_info("Elapsed:", f"{elapsed:.1f}s")
    _print_info("Smoke output:", str(smoke_output)[:120].replace("\n", " "))
    print()


# ─── DOCTOR ────────────────────────────────────────────────────────────────

def _get_available_ram_gb() -> float:
    try:
        import psutil

        return float(psutil.virtual_memory().available) / (1024 ** 3)
    except Exception:
        return 0.0


def cmd_doctor(args):
    """Run first-time environment diagnostics for demo/build flows."""
    print(_logo())
    print(f"\n  {Colors.BOLD}DOCTOR{Colors.RESET}\n")

    compression_profile = getattr(args, "compression_profile", "lossless-max")
    _print_info("Compression profile:", compression_profile)
    print()

    checks: List[Tuple[str, bool, str]] = []

    core_modules = ["torch", "transformers", "huggingface_hub", "cryptography"]
    if compression_profile == "lossless-max":
        core_modules.append("zstandard")

    missing = [m for m in core_modules if not _is_module_available(m)]
    checks.append(
        (
            "Dependencies",
            len(missing) == 0,
            "ok" if not missing else (
                "Missing: " + ", ".join(missing)
                + '. Fix: pip install "osmosis[huggingface,compression]"'
            ),
        )
    )

    try:
        _check_hf_network(timeout_sec=4.0)
        checks.append(("Network", True, "huggingface.co reachable"))
    except Exception:
        checks.append(
            (
                "Network",
                False,
                "Cannot reach huggingface.co. Fix: verify internet/proxy settings.",
            )
        )

    free_gb = shutil.disk_usage(Path.cwd()).free / (1024 ** 3)
    checks.append(
        (
            "Disk",
            free_gb >= 2.0,
            f"{free_gb:.1f} GB free"
            + ("" if free_gb >= 2.0 else " (need ~2 GB+). Fix: free disk space."),
        )
    )

    avail_ram = _get_available_ram_gb()
    checks.append(
        (
            "RAM",
            avail_ram >= 2.0 if avail_ram > 0 else True,
            f"{avail_ram:.1f} GB available" if avail_ram > 0 else "unavailable (install psutil for precise check)",
        )
    )

    ok_count = 0
    for name, passed, details in checks:
        icon = "✓" if passed else "✗"
        color = Colors.GREEN if passed else Colors.RED
        print(f"  {color}{icon}{Colors.RESET} {name:<14} {details}")
        if passed:
            ok_count += 1

    print()
    if ok_count == len(checks):
        _print_info("Doctor status:", f"{Colors.GREEN}PASS{Colors.RESET}")
    else:
        _print_info("Doctor status:", f"{Colors.RED}FAIL{Colors.RESET}")
        raise RuntimeError("Environment checks failed. See remediation above.")


# ─── RUNTIME BENCHMARK ────────────────────────────────────────────────────

def _current_rss_mb() -> float:
    try:
        import psutil

        return float(psutil.Process(os.getpid()).memory_info().rss) / (1024 ** 2)
    except Exception:
        return 0.0


def _load_runtime_baselines(path: Path) -> Dict[str, Any]:
    if not path.exists():
        return dict(DEFAULT_RUNTIME_BENCHMARK_BASELINES)
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception as exc:
        raise ValueError(f"Invalid runtime baseline JSON: {path}") from exc
    if not isinstance(data, dict):
        raise ValueError("Runtime baseline file must be a JSON object")
    return data


def _resolve_runtime_profile_limits(
    baseline_data: Dict[str, Any],
    profile: str,
) -> Optional[Dict[str, float]]:
    # v2 schema: {"schema_version": 2, "profiles": {...}}
    profiles = baseline_data.get("profiles")
    if isinstance(profiles, dict):
        limits = profiles.get(profile)
        return limits if isinstance(limits, dict) else None
    # Legacy schema: {"profile_name": {...}}
    limits = baseline_data.get(profile)
    return limits if isinstance(limits, dict) else None


def _median(values: List[float]) -> float:
    if not values:
        return 0.0
    ordered = sorted(float(v) for v in values)
    mid = len(ordered) // 2
    if len(ordered) % 2:
        return ordered[mid]
    return (ordered[mid - 1] + ordered[mid]) / 2.0


def _is_non_empty_runtime_output(output: Any) -> bool:
    if output is None:
        return False
    if isinstance(output, str):
        return bool(output.strip())
    if isinstance(output, (list, tuple, dict, set)):
        return len(output) > 0
    try:
        import numpy as np

        if isinstance(output, np.ndarray):
            return bool(output.size > 0)
    except Exception:
        pass
    return True


def _infer_benchmark_input(image: Any):
    import numpy as np

    # Best-effort generic fallback for non-text artifacts.
    for _, arr in sorted(image.osmo_model.parameters.items()):
        if getattr(arr, "ndim", 0) == 2 and arr.shape[1] > 0:
            return np.zeros((1, int(arr.shape[1])), dtype=np.float32)
        if getattr(arr, "ndim", 0) == 4 and arr.shape[1] > 0:
            side = int(max(8, arr.shape[2], arr.shape[3]))
            return np.zeros((1, int(arr.shape[1]), side, side), dtype=np.float32)
        if getattr(arr, "ndim", 0) == 3 and arr.shape[1] > 0:
            width = int(max(8, arr.shape[2]))
            return np.zeros((1, int(arr.shape[1]), width), dtype=np.float32)
    return np.zeros((1, 8), dtype=np.float32)


def cmd_benchmark_runtime(args):
    """Benchmark runtime KPIs for a built .osmoimg artifact."""
    print(_logo())
    print(f"\n  {Colors.BOLD}BENCHMARK-RUNTIME{Colors.RESET}\n")

    from model_osmosis.format.image import OsmoImage

    image_path = Path(args.image)
    _print_info("Image:", str(image_path))
    _print_info("Profile:", args.profile)
    _print_info("Iterations:", str(args.iterations))
    _print_info("Warmup runs:", str(args.warmup_runs))
    _print_info("Runtime profile:", args.runtime_profile)
    _print_info("Device:", args.device)
    _print_info("Runtime quantization:", args.runtime_quantization)
    _print_info("Runtime cache:", "on" if args.runtime_cache else "off")
    _print_info("Baseline file:", str(args.baseline_file))
    print()

    baseline_data = _load_runtime_baselines(Path(args.baseline_file))
    profile_limits = _resolve_runtime_profile_limits(baseline_data, args.profile)
    if not isinstance(profile_limits, dict):
        raise ValueError(f"Profile '{args.profile}' not found in baseline file")

    cold_start_begin = time.perf_counter()
    image = OsmoImage.load(image_path)
    cold_start_seconds = time.perf_counter() - cold_start_begin

    peak_rss_mb = _current_rss_mb()

    execution_mode = "text"
    benchmark_input = None

    run_kwargs = {
        "runtime_profile": args.runtime_profile,
        "device": args.device,
        "runtime_cache": args.runtime_cache,
        "runtime_quantization": args.runtime_quantization,
    }

    first_token_begin = time.perf_counter()
    try:
        first_token_output = image.run(
            prompt=args.prompt,
            max_tokens=1,
            temperature=args.temperature,
            **run_kwargs,
        )
    except Exception:
        execution_mode = "generic"
        benchmark_input = _infer_benchmark_input(image)
        first_token_output = image.run(input_data=benchmark_input, **run_kwargs)
    first_token_latency_seconds = time.perf_counter() - first_token_begin
    peak_rss_mb = max(peak_rss_mb, _current_rss_mb())
    if not _is_non_empty_runtime_output(first_token_output):
        raise RuntimeError("Runtime benchmark failed: empty first-token output")

    if args.warmup_runs > 0:
        if execution_mode == "text":
            image.warm_runtime(
                runtime_profile=args.runtime_profile,
                device=args.device,
                runtime_quantization=args.runtime_quantization,
                runtime_cache=True,
                warmup_prompt=args.prompt,
                warmup_runs=int(args.warmup_runs),
            )
        else:
            image.warm_runtime(
                runtime_profile=args.runtime_profile,
                device=args.device,
                runtime_quantization=args.runtime_quantization,
                runtime_cache=True,
                warmup_input_data=benchmark_input,
                warmup_runs=int(args.warmup_runs),
            )
        peak_rss_mb = max(peak_rss_mb, _current_rss_mb())

    warm_first_latencies: List[float] = []
    for _ in range(int(args.iterations)):
        t0 = time.perf_counter()
        if execution_mode == "text":
            out = image.run(
                prompt=args.prompt,
                max_tokens=1,
                temperature=args.temperature,
                **run_kwargs,
            )
        else:
            out = image.run(input_data=benchmark_input, **run_kwargs)
        latency = time.perf_counter() - t0
        if not _is_non_empty_runtime_output(out):
            raise RuntimeError("Runtime benchmark failed: empty warm first-token output")
        warm_first_latencies.append(latency)
        peak_rss_mb = max(peak_rss_mb, _current_rss_mb())
    warm_first_token_latency_seconds = _median(warm_first_latencies)

    full_latencies: List[float] = []
    full_tps: List[float] = []
    for _ in range(int(args.iterations)):
        t0 = time.perf_counter()
        if execution_mode == "text":
            full_output = image.run(
                prompt=args.prompt,
                max_tokens=args.max_tokens,
                temperature=args.temperature,
                **run_kwargs,
            )
        else:
            full_output = image.run(input_data=benchmark_input, **run_kwargs)
        elapsed = time.perf_counter() - t0
        if not _is_non_empty_runtime_output(full_output):
            raise RuntimeError("Runtime benchmark failed: empty generation output")
        full_latencies.append(elapsed)
        full_tps.append(float(args.max_tokens) / max(elapsed, 1e-6))
        peak_rss_mb = max(peak_rss_mb, _current_rss_mb())

    tokens_per_second = _median(full_tps)

    if full_latencies:
        sorted_lat = sorted(full_latencies)
        p95_index = max(0, min(len(sorted_lat) - 1, int(round(0.95 * len(sorted_lat) + 0.5)) - 1))
        p95_latency_seconds = float(sorted_lat[p95_index])
    else:
        p95_latency_seconds = 0.0

    metrics = {
        "profile": args.profile,
        "image_path": str(image_path),
        "measured_at": _utc_now(),
        "runtime_profile": args.runtime_profile,
        "device": args.device,
        "runtime_quantization": args.runtime_quantization,
        "runtime_cache": bool(args.runtime_cache),
        "warmup_runs": int(args.warmup_runs),
        "cold_start_seconds": round(cold_start_seconds, 4),
        "peak_rss_mb": round(peak_rss_mb, 2),
        "first_token_latency_seconds": round(first_token_latency_seconds, 4),
        "warm_first_token_latency_seconds": round(warm_first_token_latency_seconds, 4),
        "tokens_per_second": round(tokens_per_second, 4),
        "p95_latency_seconds": round(p95_latency_seconds, 4),
        "execution_mode": execution_mode,
    }

    regressions: List[str] = []
    max_cold = profile_limits.get("max_cold_start_seconds")
    if max_cold is not None and metrics["cold_start_seconds"] > float(max_cold):
        regressions.append(
            f"cold_start_seconds {metrics['cold_start_seconds']} > {float(max_cold)}"
        )
    max_rss = profile_limits.get("max_peak_rss_mb")
    if max_rss is not None and metrics["peak_rss_mb"] > float(max_rss):
        regressions.append(f"peak_rss_mb {metrics['peak_rss_mb']} > {float(max_rss)}")
    max_first = profile_limits.get("max_first_token_latency_seconds")
    if max_first is not None and metrics["first_token_latency_seconds"] > float(max_first):
        regressions.append(
            f"first_token_latency_seconds {metrics['first_token_latency_seconds']} > {float(max_first)}"
        )
    max_warm_first = profile_limits.get("max_warm_first_token_latency_seconds")
    if (
        max_warm_first is not None
        and metrics["warm_first_token_latency_seconds"] > float(max_warm_first)
    ):
        regressions.append(
            "warm_first_token_latency_seconds "
            f"{metrics['warm_first_token_latency_seconds']} > {float(max_warm_first)}"
        )
    min_tps = profile_limits.get("min_tokens_per_second")
    if min_tps is not None and metrics["tokens_per_second"] < float(min_tps):
        regressions.append(
            f"tokens_per_second {metrics['tokens_per_second']} < {float(min_tps)}"
        )
    max_p95 = profile_limits.get("max_p95_latency_seconds")
    if max_p95 is not None and metrics["p95_latency_seconds"] > float(max_p95):
        regressions.append(
            f"p95_latency_seconds {metrics['p95_latency_seconds']} > {float(max_p95)}"
        )

    metrics["baseline_profile"] = profile_limits
    metrics["regressions"] = regressions
    metrics["passed"] = len(regressions) == 0

    if args.json:
        print(json.dumps(metrics, indent=2, sort_keys=True))
    else:
        _print_info("Execution mode:", execution_mode)
        _print_info("Cold start:", f"{metrics['cold_start_seconds']:.3f}s")
        _print_info("Peak RSS:", f"{metrics['peak_rss_mb']:.1f} MB")
        _print_info("First token:", f"{metrics['first_token_latency_seconds']:.3f}s")
        _print_info("Warm first token:", f"{metrics['warm_first_token_latency_seconds']:.3f}s")
        _print_info("Tokens/sec:", f"{metrics['tokens_per_second']:.2f}")
        _print_info("p95 latency:", f"{metrics['p95_latency_seconds']:.3f}s")
        _print_info("Status:", f"{Colors.GREEN}PASS{Colors.RESET}" if metrics["passed"] else f"{Colors.RED}FAIL{Colors.RESET}")
        if regressions:
            _print_info("Regressions:", "; ".join(regressions))
    if not metrics["passed"]:
        raise RuntimeError("Runtime benchmark regression detected against baseline profile")


# ─── TRANSCODE ────────────────────────────────────────────────────────────

def _normalize_tokenizer_validation_status(validation: Dict[str, Any]) -> str:
    status = str((validation or {}).get("status", "missing_or_corrupt")).strip().lower()
    if status not in {"real_assets", "placeholder_only", "missing_or_corrupt"}:
        return "missing_or_corrupt"
    return status


def _migration_remediation(status: str, result_path: Path, post_trust_status: str) -> str:
    if status == "missing_or_corrupt":
        return (
            "Rebuild with real tokenizer assets and re-sign, e.g.: "
            f"osmosis build <model> -o {result_path} --compression-profile lossless-max --trust-sign"
        )
    if status == "placeholder_only":
        return (
            "Rebuild from source with real tokenizer files, then upgrade trust, e.g.: "
            f"osmosis transcode {result_path} --upgrade-trust --compression-profile lossless-max"
        )
    if post_trust_status != "trusted":
        return (
            "Artifact is signed but not locally trusted. Import signer key and verify again: "
            f"osmosis trust import-signer --name <signer> --pubkey <pubkey-file> && osmosis verify {result_path}"
        )
    return ""


def cmd_transcode(args):
    """Recompress an existing .osmoimg using a selected compression profile."""
    print(_logo())
    print(f"\n  {Colors.BOLD}TRANSCODE{Colors.RESET}\n")

    source_path = Path(args.image)
    default_output = source_path.with_name(f"{source_path.stem}.v2{source_path.suffix}")
    output = Path(args.output) if args.output else default_output
    compression_engine = _normalize_mode_arg(
        getattr(args, "compression_engine", "auto"),
        COMPRESSION_ENGINE_MODES,
        "auto",
    )

    _print_info("Input:", str(source_path))
    _print_info("Output:", str(output))
    _print_info("Compression profile:", args.compression_profile)
    _print_info("Compression engine:", _resolve_compression_engine_mode(compression_engine))
    _print_info("Trust sign:", "yes" if args.trust_sign else "no")
    _print_info("Upgrade trust:", "yes" if args.upgrade_trust else "no")
    _print_info("Signer backend:", args.signer_backend)
    migration_report_path = getattr(args, "migration_report", None)
    _print_info("Migration report:", migration_report_path or "none")
    print()
    if not args.trust_sign and not args.upgrade_trust:
        _warn_insecure_output("transcode")
        print()

    from model_osmosis.format.image import OsmoImage

    _print_step("📥", "Loading source image...")
    try:
        image = OsmoImage.load(source_path)
    except Exception as exc:
        raise _error_with_next_command(
            f"Failed to load source image for transcode: {exc}",
            f"osmosis verify {source_path}",
        ) from exc
    source_tokenizer_validation = (
        image.runtime.compression_stats.get("tokenizer_validation", {})
        if isinstance(image.runtime.compression_stats, dict)
        else {}
    )
    if not isinstance(source_tokenizer_validation, dict) or not source_tokenizer_validation:
        source_tokenizer_validation = OsmoImage.validate_tokenizer_bundle(image.tokenizer_data)
    if (
        int(source_tokenizer_validation.get("manifest_schema_version", 0) or 0) >= 2
        and _normalize_tokenizer_validation_status(source_tokenizer_validation) == "missing_or_corrupt"
    ):
        raise _error_with_next_command(
            "Cannot transcode artifact: tokenizer manifest validation failed for source artifact.",
            f"osmosis verify {source_path}",
        )

    pre_verify = OsmoImage.verify(source_path, require_trust=False)
    pre_trust_status = str(pre_verify.get("trust_status", "unsigned"))

    _print_step("🔄", f"Transcoding with profile {args.compression_profile}...")
    trust_sign = bool(args.trust_sign or args.upgrade_trust)
    result_path = image.save(
        output,
        compression_profile=args.compression_profile,
        compression_engine=compression_engine,
        trust_sign=trust_sign,
        signer_id=args.signer_id,
        signer_backend=args.signer_backend,
    )
    post_verify = OsmoImage.verify(result_path, require_trust=False)
    post_trust_status = str(post_verify.get("trust_status", "unsigned"))
    try:
        loaded_result = OsmoImage.load(result_path)
    except Exception as exc:
        raise _error_with_next_command(
            f"Failed to load transcoded output for validation: {exc}",
            f"osmosis verify {result_path}",
        ) from exc
    tokenizer_validation = (
        loaded_result.runtime.compression_stats.get("tokenizer_validation", {})
        if isinstance(loaded_result.runtime.compression_stats, dict)
        else {}
    )
    if not isinstance(tokenizer_validation, dict) or not tokenizer_validation:
        tokenizer_validation = OsmoImage.validate_tokenizer_bundle(loaded_result.tokenizer_data)
    if (
        int(tokenizer_validation.get("manifest_schema_version", 0) or 0) >= 2
        and _normalize_tokenizer_validation_status(tokenizer_validation) == "missing_or_corrupt"
    ):
        raise _error_with_next_command(
            "Transcode output failed tokenizer validation; artifact was not upgraded safely.",
            f"osmosis verify {result_path}",
        )
    tokenizer_status = _normalize_tokenizer_validation_status(tokenizer_validation)
    remediation = _migration_remediation(tokenizer_status, Path(result_path), post_trust_status)
    migration_report = {
        "schema_version": 1,
        "input_image": str(source_path),
        "output_image": str(result_path),
        "compression_profile": args.compression_profile,
        "upgrade_trust": bool(args.upgrade_trust),
        "pre_trust_status": pre_trust_status,
        "post_trust_status": post_trust_status,
        "tokenizer_validation_status": tokenizer_status,
        "tokenizer_validation": tokenizer_validation,
        "recommended_remediation": remediation,
        "generated_at": _utc_now(),
    }
    if migration_report_path:
        report_path = Path(migration_report_path)
        report_path.parent.mkdir(parents=True, exist_ok=True)
        report_path.write_text(json.dumps(migration_report, indent=2, sort_keys=True), encoding="utf-8")

    print()
    _print_info("Payload codec:", image.runtime.payload_codec)
    _print_info("Trust status (pre):", pre_trust_status)
    _print_info("Trust status (post):", post_trust_status)
    _print_info("Tokenizer validation:", tokenizer_status)
    _print_info("Stored ratio:", f"{image.runtime.compression_ratio:.2f}x")
    effective_ratio = image.runtime.compression_stats.get("effective_compression_ratio", 1.0)
    _print_info("Effective ratio:", f"{float(effective_ratio):.2f}x")
    if remediation:
        _print_info("Remediation:", remediation)
    if migration_report_path:
        _print_info("Migration report:", str(Path(migration_report_path)))
    _print_step("✓", f"Saved to {Colors.BOLD}{result_path}{Colors.RESET}")
    print()


# ─── PUSH ─────────────────────────────────────────────────────────────────

def cmd_push(args):
    """Upload an .osmoimg file to HuggingFace Hub."""
    print(_logo())
    print(f"\n  {Colors.BOLD}PUSH{Colors.RESET}\n")

    image_path = args.image
    repo_id = args.repo

    _print_info("Image:", image_path)
    _print_info("Repo:", repo_id)
    _print_info("Private:", "yes" if args.private else "no")
    print()

    try:
        from huggingface_hub import HfApi
    except ImportError:
        _print_error("huggingface_hub required: pip install huggingface_hub")
        return

    from model_osmosis.format.image import OsmoImage

    _print_step("📥", "Reading image metadata...")
    image = OsmoImage.load(image_path)

    api = HfApi()

    _print_step("📤", f"Creating repo {repo_id}...")
    api.create_repo(repo_id, private=args.private, exist_ok=True)

    _print_step("📤", f"Uploading {image_path}...")
    api.upload_file(
        path_or_fileobj=image_path,
        path_in_repo=Path(image_path).name,
        repo_id=repo_id,
    )

    if image.readme:
        _print_step("📝", "Uploading model card...")
        api.upload_file(
            path_or_fileobj=image.readme.encode("utf-8"),
            path_in_repo="README.md",
            repo_id=repo_id,
        )

    print()
    _print_step("✓", f"Pushed to {Colors.BOLD}{repo_id}{Colors.RESET}")
    print()


# ─── CONVERT ──────────────────────────────────────────────────────────────

def cmd_convert(args):
    """Convert a model between frameworks."""
    print(_logo())
    print(f"\n  {Colors.BOLD}CONVERT{Colors.RESET}\n")

    _print_info("Input:", args.model)
    _print_info("Target:", args.to)
    _print_info("Output:", args.output or "(in memory)")
    _print_info("Seed:", str(args.seed))
    _print_info("Cross opt stage:", str(getattr(args, "cross_opt_stage", "auto")))
    _print_info("Strict load:", "no" if args.allow_partial_load else "yes")
    _print_info("Allow lossy cast:", "yes" if args.allow_lossy_cast else "no")
    print()

    from model_osmosis.format.converters import UniversalConverter
    from model_osmosis.format.image import OsmoImage
    from model_osmosis.format.file_format import is_osmoimg
    from model_osmosis.utils.determinism import set_global_seed

    set_global_seed(args.seed)
    input_path = Path(args.model)
    converter = UniversalConverter()

    if input_path.suffix == ".osmoimg" or (input_path.exists() and is_osmoimg(input_path)):
        _print_step("📥", "Loading .osmoimg...")
        image = OsmoImage.load(input_path)
        _print_step("🔄", f"Exporting to {args.to}...")
        model = image.export(
            target_format=args.to,
            output_path=args.output,
            strict_load=not args.allow_partial_load,
            allow_lossy_cast=args.allow_lossy_cast,
        )
        report = {
            "missing_keys": [],
            "unexpected_keys": [],
            "lossy_casts": [],
            "dtype_transports": [],
        }
    else:
        _print_step("📥", f"Loading {args.model}...")
        model = _load_model(args.model)
        _print_step("🔄", f"Converting to {args.to}...")
        model = converter.convert(
            model,
            target_framework=args.to,
            strict_load=not args.allow_partial_load,
            allow_lossy_cast=args.allow_lossy_cast,
        )
        report = converter.last_report

        if args.output:
            _save_model(model, args.output)

    if args.output:
        _print_step("✓", f"Saved to {Colors.BOLD}{args.output}{Colors.RESET}")
    else:
        _print_step("✓", f"Converted in memory ({type(model).__name__})")
    reliability_mode = (
        "degraded" if args.allow_partial_load or args.allow_lossy_cast else "strict"
    )
    _print_info("Reliability mode:", reliability_mode)
    _print_info("Missing keys:", str(len(report.get("missing_keys", []))))
    _print_info("Unexpected keys:", str(len(report.get("unexpected_keys", []))))
    _print_info("Lossy casts:", str(len(report.get("lossy_casts", []))))
    _print_info("Dtype transports:", str(len(report.get("dtype_transports", []))))
    _print_info(
        "Fallbacks:",
        _reliability_fallbacks(args.allow_partial_load, args.allow_lossy_cast),
    )
    print()


# ─── DIFF ─────────────────────────────────────────────────────────────────

def cmd_diff(args):
    """Compare metadata of two .osmoimg files."""
    print(_logo())
    print(f"\n  {Colors.BOLD}DIFF{Colors.RESET}\n")

    _print_info("Image A:", args.image_a)
    _print_info("Image B:", args.image_b)
    print()

    from model_osmosis.format.image import OsmoImage

    info_a = OsmoImage.inspect(args.image_a)
    info_b = OsmoImage.inspect(args.image_b)

    def _compare_fields(section, keys, data_a, data_b):
        print(f"  {Colors.BOLD}{section}{Colors.RESET}")
        print(f"  {'Field':<22}{'Image A':>20}{'Image B':>20}  {'':>4}")
        print(f"  {'─'*70}")
        for key in keys:
            val_a = data_a.get(key, "?")
            val_b = data_b.get(key, "?")
            if isinstance(val_a, float):
                val_a = f"{val_a:.2f}"
            if isinstance(val_b, float):
                val_b = f"{val_b:.2f}"
            val_a, val_b = str(val_a), str(val_b)
            marker = f"{Colors.GREEN}={Colors.RESET}" if val_a == val_b else f"{Colors.YELLOW}≠{Colors.RESET}"
            print(f"  {key:<22}{val_a:>20}{val_b:>20}  {marker}")
        print()

    _compare_fields("Model", ["name", "parameter_count", "framework", "version", "license"],
                    info_a, info_b)
    _compare_fields("Architecture",
                    list(set(list(info_a.get("architecture", {}).keys()) +
                             list(info_b.get("architecture", {}).keys()))),
                    info_a.get("architecture", {}), info_b.get("architecture", {}))
    _compare_fields("Runtime",
                    ["source_framework", "model_type", "is_encrypted", "min_ram_gb"],
                    info_a.get("runtime", {}), info_b.get("runtime", {}))
    _compare_fields("Image",
                    ["file_size_mb", "original_size_mb", "compression_ratio", "built_at"],
                    info_a.get("image", {}), info_b.get("image", {}))


# ─── VALIDATE ─────────────────────────────────────────────────────────────

def cmd_validate(args):
    """Verify the integrity of an .osmoimg file."""
    print(_logo())
    print(f"\n  {Colors.BOLD}VALIDATE{Colors.RESET}\n")

    _print_info("Image:", args.image)
    print()

    import io
    import zipfile

    image_path = Path(args.image)
    all_ok = True

    # 1. File exists
    if not image_path.exists():
        _print_step("✗", f"{Colors.RED}File not found{Colors.RESET}")
        return
    _print_step("✓", "File exists")

    # 2. Magic bytes
    from model_osmosis.format.file_format import OSMOIMG_MAGIC, read_osmoimg_header
    try:
        with open(image_path, "rb") as f:
            magic = f.read(8)
        if magic == OSMOIMG_MAGIC:
            _print_step("✓", "Magic bytes valid")
        else:
            _print_step("✗", f"{Colors.RED}Invalid magic bytes{Colors.RESET}")
            all_ok = False
            print()
            return
    except Exception as e:
        _print_step("✗", f"{Colors.RED}Cannot read file: {e}{Colors.RESET}")
        return

    # 3. Header valid
    try:
        with open(image_path, "rb") as f:
            runtime_data = read_osmoimg_header(f)
            remaining = f.read()
        _print_step("✓", "Header valid")
    except Exception as e:
        _print_step("✗", f"{Colors.RED}Invalid header: {e}{Colors.RESET}")
        all_ok = False
        print()
        return

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

    # 4. Payload valid
    try:
        if payload_codec == "zip-deflate":
            if zipfile.is_zipfile(io.BytesIO(remaining)):
                _print_step("✓", "Zip payload valid")
            else:
                _print_step("✗", f"{Colors.RED}Invalid zip payload{Colors.RESET}")
                all_ok = False
        else:
            from model_osmosis.format.adaptive_compression import inspect_osmo_model_lossless_v2

            inspect_osmo_model_lossless_v2(remaining)
            _print_step("✓", f"V2 payload valid ({payload_codec})")
    except Exception as e:
        _print_step("✗", f"{Colors.RED}Payload check failed: {e}{Colors.RESET}")
        all_ok = False

    # 5. Manifest loadable
    try:
        from model_osmosis.format.image import OsmoImage
        info = OsmoImage.inspect(image_path)
        _print_step("✓", f"Manifest valid (name: {info['name'] or '(unnamed)'})")
    except Exception as e:
        _print_step("✗", f"{Colors.RED}Cannot read manifest: {e}{Colors.RESET}")
        all_ok = False

    # 6. Parameters loadable
    try:
        from model_osmosis.format.image import OsmoImage
        image = OsmoImage.load(image_path)
        param_count = image.parameter_count
        _print_step("✓", f"Parameters loadable ({param_count:,} params)")
    except Exception as e:
        _print_step("✗", f"{Colors.RED}Cannot load parameters: {e}{Colors.RESET}")
        all_ok = False

    # 7. Trust verification
    try:
        from model_osmosis.format.image import OsmoImage

        verify_result = OsmoImage.verify(image_path)
        trust_status = verify_result.get("trust_status", "unsigned")
        if trust_status == "trusted":
            _print_step("✓", "Trust signature valid (trusted signer)")
        elif trust_status == "signed_untrusted":
            _print_step("⚠", "Signature valid but signer is not trusted locally")
        elif trust_status == "unsigned":
            _print_step("⚠", "Artifact unsigned (v1/v2 compatibility path)")
        else:
            _print_step("✗", f"{Colors.RED}Trust verification failed ({trust_status}){Colors.RESET}")
            all_ok = False
    except Exception as e:
        _print_step("✗", f"{Colors.RED}Trust verification error: {e}{Colors.RESET}")
        all_ok = False

    print()
    if all_ok:
        print(f"  {Colors.GREEN}{Colors.BOLD}✓ Image is valid{Colors.RESET}")
    else:
        print(f"  {Colors.RED}{Colors.BOLD}✗ Image has errors{Colors.RESET}")
    print()


def cmd_verify(args):
    """Verify trusted-signature status for an .osmoimg artifact."""
    print(_logo())
    print(f"\n  {Colors.BOLD}VERIFY{Colors.RESET}\n")
    _print_info("Image:", args.image)
    _print_info("Require trust:", "yes" if args.require_trust else "no")
    _print_info("Repro check:", "yes" if getattr(args, "repro_check", False) else "no")
    _print_info("Policy file:", getattr(args, "policy_file", None) or "none")
    print()

    from model_osmosis.format.image import OsmoImage

    result = OsmoImage.verify(args.image, require_trust=False)
    policy = _load_policy_file(getattr(args, "policy_file", None))
    policy_result = _evaluate_policy(result, policy)
    result["policy_result"] = policy_result
    audit_result = _emit_audit_event(
        action="verify",
        image_path=args.image,
        verify_result=result,
        policy_result=policy_result,
    )
    if getattr(args, "repro_check", False):
        result["repro_check"] = _run_repro_check(args.image, result)
    status = result.get("trust_status", "unsigned")

    if args.json:
        print(json.dumps(result, indent=2, sort_keys=True))
    else:
        color = Colors.GREEN if status == "trusted" else Colors.YELLOW
        if status == "tampered":
            color = Colors.RED
        _print_info("Trust status:", f"{color}{status}{Colors.RESET}")
        _print_info("Signer ID:", str(result.get("signer_id") or "none"))
        _print_info("Key ID:", str(result.get("key_id") or "none"))
        _print_info("Algorithm:", str(result.get("signature_alg") or "none"))
        _print_info("Signer backend:", str(result.get("signer_backend") or "local-keyring"))
        _print_info("Payload SHA256:", str(result.get("payload_sha256") or "n/a"))
        _print_info("Artifact ID:", str(result.get("artifact_id") or "n/a"))
        _print_info("Manifest SHA256:", str(result.get("manifest_sha256") or "n/a"))
        _print_info("Tensor index SHA256:", str(result.get("tensor_index_sha256") or "n/a"))
        _print_info("Verified at:", str(result.get("verified_at") or "n/a"))
        _print_info("Transparency:", str(result.get("transparency_status") or "n/a"))
        _print_info("Reason:", str(result.get("reason") or ""))
        _print_info("Header version:", str(result.get("header_version", "?")))
        _print_info("Payload codec:", str(result.get("payload_codec", "unknown")))
        _print_info("Policy:", "PASS" if policy_result.get("passed", True) else "FAIL")
        if policy_result.get("checked") and not policy_result.get("passed"):
            _print_info("Policy violations:", "; ".join(policy_result.get("violations", [])))
        repro_result = result.get("repro_check")
        if isinstance(repro_result, dict):
            _print_info("Repro check:", "PASS" if repro_result.get("passed") else "FAIL")
            if not repro_result.get("passed"):
                _print_info("Repro reason:", str(repro_result.get("reason") or "unknown"))
        _print_info("Audit logged:", "yes" if audit_result.get("logged") else "no")
        if status == "unsigned":
            print()
            print(
                "  Upgrade path: transcode and sign with trusted metadata, e.g.\n"
                f"    osmosis transcode {args.image} --compression-profile lossless-max --upgrade-trust\n"
                f"    osmosis build <model> -o trusted.osmoimg --trust-sign"
            )

    if policy_result.get("checked") and not policy_result.get("passed"):
        raise RuntimeError(
            "Verification blocked by policy: "
            + "; ".join(policy_result.get("violations", []))
        )
    repro_result = result.get("repro_check")
    if getattr(args, "repro_check", False) and isinstance(repro_result, dict):
        if not repro_result.get("passed", False):
            raise RuntimeError(
                "Reproducibility check failed: "
                + str(repro_result.get("reason") or "payload mismatch")
            )
    if args.require_trust and status != "trusted":
        raise RuntimeError(f"Trusted artifact required; status={status}")
    print()


def cmd_trust_keygen(args):
    """Generate a local signer keypair and trust it."""
    print(_logo())
    print(f"\n  {Colors.BOLD}TRUST KEYGEN{Colors.RESET}\n")
    from model_osmosis.format.trust import generate_signer

    result = generate_signer(args.name, overwrite=args.overwrite)
    _print_info("Signer ID:", result.get("signer_id", ""))
    _print_info("Key ID:", result.get("key_id", ""))
    _print_info("Private key:", result.get("private_key_path", ""))
    _print_info("Public key:", result.get("public_key_path", ""))
    _print_info("Trusted key:", result.get("trusted_key_path", ""))
    _print_info("Created:", result.get("created", "false"))
    print()


def cmd_trust_list_keys(args):
    """List local private keys and trusted signer keys."""
    print(_logo())
    print(f"\n  {Colors.BOLD}TRUST LIST-KEYS{Colors.RESET}\n")
    from model_osmosis.format.trust import list_keys

    data = list_keys()
    private_keys = data.get("private_keys", [])
    trusted = data.get("trusted_signers", [])

    print(f"  {Colors.BOLD}Private Keys{Colors.RESET}")
    if not private_keys:
        print("  (none)")
    for item in private_keys:
        print(
            f"  - {item.get('signer_id')} "
            f"[{item.get('key_id')}] -> {item.get('private_key_path')}"
        )
    print()
    print(f"  {Colors.BOLD}Trusted Signers{Colors.RESET}")
    if not trusted:
        print("  (none)")
    for item in trusted:
        print(
            f"  - {item.get('signer_id')} "
            f"[{item.get('key_id')}] -> {item.get('trusted_key_path')}"
        )
    print()


def cmd_trust_import_signer(args):
    """Import a trusted signer public key."""
    print(_logo())
    print(f"\n  {Colors.BOLD}TRUST IMPORT-SIGNER{Colors.RESET}\n")
    from model_osmosis.format.trust import import_trusted_signer

    result = import_trusted_signer(args.name, args.pubkey)
    _print_info("Signer ID:", result.get("signer_id", ""))
    _print_info("Key ID:", result.get("key_id", ""))
    _print_info("Trusted key:", result.get("trusted_key_path", ""))
    print()


def cmd_trust_revoke_signer(args):
    """Remove a signer from local trust store."""
    print(_logo())
    print(f"\n  {Colors.BOLD}TRUST REVOKE-SIGNER{Colors.RESET}\n")
    from model_osmosis.format.trust import revoke_trusted_signer

    removed = revoke_trusted_signer(args.name)
    if removed:
        _print_step("✓", f"Revoked signer '{args.name}'")
    else:
        _print_step("⚠", f"Signer '{args.name}' was not in trusted store")
    print()


# ─── ARTIFACT TRACKING ─────────────────────────────────────────────────────

def cmd_artifact_track(args):
    """Track one or more .osmoimg artifacts in a deterministic lock file."""
    print(_logo())
    print(f"\n  {Colors.BOLD}ARTIFACT TRACK{Colors.RESET}\n")
    from model_osmosis.services.artifact_tracking import track_artifacts

    try:
        result = track_artifacts(
            images=list(getattr(args, "images", []) or []),
            lock_file=getattr(args, "lock_file", None),
            default_lock_file=ARTIFACT_LOCKFILE_DEFAULT,
        )
    except RuntimeError as exc:
        if "No artifacts were provided to track" in str(exc):
            raise _error_with_next_command(
                str(exc),
                "osmosis artifact track <image.osmoimg>",
            ) from exc
        raise

    tracked_paths = list(result.get("tracked_artifacts", []))
    _print_info("Tracked artifacts:", str(result.get("tracked_count", len(tracked_paths))))
    _print_info("Lock file:", str(result.get("lock_file", getattr(args, "lock_file", ""))))
    for tracked in sorted(str(x) for x in tracked_paths):
        print(f"  - {tracked}")
    print()


def cmd_artifact_verify(args):
    """Verify tracked artifacts against lock metadata and fail on drift/tamper."""
    print(_logo())
    print(f"\n  {Colors.BOLD}ARTIFACT VERIFY{Colors.RESET}\n")

    from model_osmosis.services.artifact_tracking import verify_artifacts

    verify_all = bool(getattr(args, "all", False))
    requested_image = str(getattr(args, "image", "") or "").strip()
    try:
        payload = verify_artifacts(
            image=requested_image or None,
            verify_all=verify_all,
            lock_file=getattr(args, "lock_file", None),
            default_lock_file=ARTIFACT_LOCKFILE_DEFAULT,
        )
    except RuntimeError as exc:
        message = str(exc)
        if message.startswith("No tracked artifacts found in lock file:"):
            raise _error_with_next_command(
                message,
                "osmosis artifact track <image.osmoimg>",
            ) from exc
        if message == "Specify a tracked IMAGE or use --all.":
            raise _error_with_next_command(
                message,
                "osmosis artifact verify --all",
            ) from exc
        if message.startswith("Artifact is not tracked:"):
            missing = message.split(":", 1)[1].strip()
            raise _error_with_next_command(
                message,
                f"osmosis artifact track {missing}",
            ) from exc
        raise

    failed = [item for item in (payload.get("results", []) or []) if not bool(item.get("passed", False))]

    if bool(getattr(args, "json", False)):
        print(json.dumps(payload, indent=2, sort_keys=True))
    else:
        _print_info(
            "Status:",
            f"{Colors.GREEN}PASS{Colors.RESET}" if not failed else f"{Colors.RED}FAIL{Colors.RESET}",
        )
        _print_info("Lock file:", str(payload.get("lock_file", getattr(args, "lock_file", ""))))
        _print_info("Checked:", str(payload.get("checked_artifacts", 0)))
        _print_info("Failed:", str(payload.get("failed_artifacts", len(failed))))
        for item in payload["results"]:
            marker = "✓" if item.get("passed") else "✗"
            detail = str(item.get("reason", "") or "ok")
            print(f"  {marker} {item.get('path')} ({detail})")
        print()

    if failed:
        remediation_target = requested_image or str(failed[0].get("path", "<image.osmoimg>"))
        raise _error_with_next_command(
            f"artifact verify failed for {len(failed)} artifact(s).",
            f"osmosis artifact track {remediation_target}",
        )


def cmd_artifact_hook_install(args):
    """Install git hooks to enforce artifact verification on commit/push."""
    print(_logo())
    print(f"\n  {Colors.BOLD}ARTIFACT HOOK INSTALL{Colors.RESET}\n")
    from model_osmosis.services.artifact_tracking import install_artifact_hooks

    try:
        result = install_artifact_hooks(
            lock_file=getattr(args, "lock_file", None),
            force=bool(getattr(args, "force", False)),
            default_lock_file=ARTIFACT_LOCKFILE_DEFAULT,
        )
    except RuntimeError as exc:
        message = str(exc)
        if "not managed by ModelOsmosis" in message:
            raise _error_with_next_command(
                message,
                "osmosis artifact hook install --force",
            ) from exc
        raise

    _print_info("Hooks dir:", str(result.get("hooks_dir", "")))
    _print_info("Verify command:", str(result.get("verify_command", "osmosis artifact verify --all")))
    _print_info("Installed:", ", ".join(str(x) for x in (result.get("installed_hooks", []) or [])))
    print()


# ─── QUANTIZE ─────────────────────────────────────────────────────────────

def cmd_quantize(args):
    """Quantize model weights in an .osmoimg."""
    print(_logo())
    print(f"\n  {Colors.BOLD}QUANTIZE{Colors.RESET}\n")

    _print_info("Input:", args.image)
    _print_info("Bits:", str(args.bits))
    _print_info("Output:", args.output or args.image)
    print()

    import numpy as np
    from model_osmosis.format.image import OsmoImage

    _print_step("📥", "Loading image...")
    image = OsmoImage.load(args.image)

    _print_step("🔧", f"Quantizing to {args.bits}-bit...")

    original_size = 0
    quantized_size = 0

    for name, param in image.osmo_model.parameters.items():
        original_size += param.nbytes
        if args.bits == 16:
            image.osmo_model.parameters[name] = param.astype(np.float16)
        elif args.bits == 8:
            scale = np.abs(param).max() / 127.0 if np.abs(param).max() > 0 else 1.0
            quantized = np.clip(np.round(param / scale), -128, 127).astype(np.int8)
            image.osmo_model.parameters[name] = quantized.astype(np.float32) * scale
        elif args.bits == 4:
            scale = np.abs(param).max() / 7.0 if np.abs(param).max() > 0 else 1.0
            quantized = np.clip(np.round(param / scale), -8, 7).astype(np.int8)
            image.osmo_model.parameters[name] = quantized.astype(np.float32) * scale
        quantized_size += image.osmo_model.parameters[name].nbytes

    image.runtime.quantization = f"int{args.bits}" if args.bits in (4, 8) else f"fp{args.bits}"

    output = args.output or args.image
    _print_step("💾", f"Saving to {output}...")
    result_path = image.save(output)

    print()
    _print_info("Original:", f"{original_size / 1024**2:.1f} MB")
    _print_info("Quantized:", f"{quantized_size / 1024**2:.1f} MB")
    if original_size > 0:
        _print_info("Reduction:", f"{(1 - quantized_size/original_size)*100:.1f}%")
    print()
    _print_step("✓", f"Saved to {Colors.BOLD}{result_path}{Colors.RESET}")
    print()


# ─── SERVE ────────────────────────────────────────────────────────────────

def cmd_serve(args):
    """Start an HTTP inference server."""
    print(_logo())
    print(f"\n  {Colors.BOLD}SERVE{Colors.RESET}\n")

    _print_info("Image:", args.image)
    _print_info("Host:", args.host)
    _print_info("Port:", str(args.port))
    _print_info("Runtime profile:", args.runtime_profile)
    _print_info("Device:", args.device)
    _print_info("Runtime quantization:", args.runtime_quantization)
    _print_info("Runtime cache:", "on" if args.runtime_cache else "off")
    _print_info("Warmup:", "on" if args.warmup else "off")
    require_trust_flag = bool(getattr(args, "require_trust", False))
    trust_profile = _resolve_trust_profile(getattr(args, "trust_profile", "standard"))
    require_trust = _effective_require_trust(require_trust_flag, trust_profile)
    _print_info("Trust profile:", trust_profile)
    _print_info("Require trust:", "yes" if require_trust else "no")
    _print_info("Policy file:", getattr(args, "policy_file", None) or "none")
    print()

    from model_osmosis.format.image import OsmoImage
    import json as json_mod

    _check_trust_for_action(
        args.image,
        action="serve",
        require_trust=require_trust,
        policy_file=getattr(args, "policy_file", None),
        trust_profile=trust_profile,
    )

    _print_step("📥", "Loading image...")
    image = OsmoImage.load(args.image)
    print(f"  {Colors.DIM}{image.summary()}{Colors.RESET}")
    print()

    warmup_prompt = None
    warmup_input = None
    warmup_runs = 1 if args.warmup else 0
    if args.warmup:
        if (image.runtime.model_type or "").lower() in {"causal_lm", "seq2seq", "language_model"}:
            warmup_prompt = RUNTIME_BENCHMARK_DEFAULT_PROMPT
        else:
            warmup_input = _infer_benchmark_input(image)
    _print_step("🔥", "Warming runtime...")
    image.warm_runtime(
        runtime_profile=args.runtime_profile,
        device=args.device,
        runtime_quantization=args.runtime_quantization,
        runtime_cache=args.runtime_cache,
        warmup_prompt=warmup_prompt,
        warmup_input_data=warmup_input,
        warmup_runs=warmup_runs,
    )

    from http.server import HTTPServer, BaseHTTPRequestHandler

    class PredictHandler(BaseHTTPRequestHandler):
        def do_POST(self):
            if self.path == "/predict":
                content_len = int(self.headers.get("Content-Length", 0))
                body = self.rfile.read(content_len)
                try:
                    data = json_mod.loads(body)
                    prompt = data.get("prompt", "")
                    max_tokens = data.get("max_tokens", 256)
                    temperature = data.get("temperature", 0.7)

                    output = image.run(
                        prompt=prompt,
                        max_tokens=max_tokens,
                        temperature=temperature,
                        runtime_profile=args.runtime_profile,
                        device=args.device,
                        runtime_cache=args.runtime_cache,
                        runtime_quantization=args.runtime_quantization,
                    )

                    response = json_mod.dumps({"output": output})
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.end_headers()
                    self.wfile.write(response.encode())
                except Exception as e:
                    error = json_mod.dumps({"error": str(e)})
                    self.send_response(500)
                    self.send_header("Content-Type", "application/json")
                    self.end_headers()
                    self.wfile.write(error.encode())
            else:
                self.send_response(404)
                self.end_headers()

        def do_GET(self):
            if self.path == "/health":
                response = json_mod.dumps({"status": "ok", "model": image.name})
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(response.encode())
            elif self.path == "/info":
                response = json_mod.dumps({
                    "name": image.name,
                    "parameters": image.parameter_count,
                    "encrypted": image.is_encrypted,
                    "framework": image.runtime.source_framework,
                })
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.end_headers()
                self.wfile.write(response.encode())
            else:
                self.send_response(404)
                self.end_headers()

        def log_message(self, format, *log_args):
            print(f"  {Colors.DIM}{self.address_string()} - {format % log_args}{Colors.RESET}")

    _print_step("🌐", f"Starting server on {args.host}:{args.port}")
    print(f"\n  {Colors.BOLD}Endpoints:{Colors.RESET}")
    print(f"    POST /predict  — Run inference (JSON: {{\"prompt\": \"...\", \"max_tokens\": 256}})")
    print(f"    GET  /health   — Health check")
    print(f"    GET  /info     — Model info")
    print(f"\n  {Colors.DIM}Press Ctrl+C to stop{Colors.RESET}\n")

    server = HTTPServer((args.host, args.port), PredictHandler)
    server.serve_forever()


# ─── HISTORY ──────────────────────────────────────────────────────────────

def cmd_history(args):
    """Show build/export history from .osmoimg metadata."""
    print(_logo())
    print(f"\n  {Colors.BOLD}HISTORY{Colors.RESET}\n")

    from model_osmosis.format.image import OsmoImage

    info = OsmoImage.inspect(args.image)

    _print_info("Image:", args.image)
    print()

    print(f"  {Colors.BOLD}Build History{Colors.RESET}")
    print(f"  {'─'*50}")

    rt = info.get("runtime", {})
    img = info.get("image", {})

    _print_info("Built at:", img.get("built_at", "unknown"))
    _print_info("Source framework:", rt.get("source_framework", "unknown"))
    _print_info("Model type:", rt.get("model_type", "unknown"))
    _print_info("Image version:", img.get("image_version", "unknown"))
    _print_info("Encrypted:", "✓ yes" if rt.get("is_encrypted") else "✗ no")
    _print_info("Quality preserved:", "✓ yes" if rt.get("preserve_quality") else "✗ no")
    _print_info("Quantization:", rt.get("quantization") or "none")
    _print_info("Compression profile:", rt.get("compression_profile", "legacy"))
    _print_info("Payload codec:", rt.get("payload_codec", "zip-deflate"))
    _print_info("Compression engine:", rt.get("compression_engine_version", "legacy-zip-v1"))
    _print_info("Trust status:", rt.get("trust_status", "unsigned"))
    _print_info("Signer ID:", rt.get("signer_id") or "none")
    _print_info("Key ID:", rt.get("key_id") or "none")
    _print_info("Signature alg:", rt.get("signature_alg") or "none")
    _print_info("Signer backend:", rt.get("signer_backend") or "local-keyring")
    _print_info("Payload SHA256:", rt.get("payload_sha256") or "n/a")
    _print_info("Artifact ID:", rt.get("artifact_id") or "n/a")
    _print_info("Manifest SHA256:", rt.get("manifest_sha256") or "n/a")
    _print_info("Tensor index SHA256:", rt.get("tensor_index_sha256") or "n/a")
    _print_info("Verified at:", rt.get("verified_at") or "n/a")
    _print_info("Transparency:", rt.get("transparency_status") or "n/a")
    _print_info("Provenance mode:", rt.get("provenance_mode", "private"))
    _print_info("Repro mode:", rt.get("repro_mode", "strict"))
    if rt.get("hf_model_id"):
        _print_info("HF Model ID:", rt["hf_model_id"])

    print()
    print(f"  {Colors.BOLD}Size{Colors.RESET}")
    print(f"  {'─'*50}")
    _print_info("File size:", f"{img.get('file_size_mb', 0):.1f} MB")
    _print_info("Original size:", f"{img.get('original_size_mb', 0):.1f} MB")
    stored_ratio = img.get("compression_ratio_stored", img.get("compression_ratio", 1))
    effective_ratio = img.get("effective_compression_ratio", img.get("compression_ratio", 1))
    _print_info("Compression (stored):", f"{stored_ratio:.2f}x")
    _print_info("Compression (effective):", f"{effective_ratio:.2f}x")
    optimizer_meta = (rt.get("compression_stats", {}) or {}).get("optimizer", {})
    if isinstance(optimizer_meta, dict) and optimizer_meta:
        _print_info("Optimizer mode:", str(optimizer_meta.get("optimizer_mode", "auto")))
        _print_info("Optimizer profile:", str(optimizer_meta.get("profile_id", "n/a")))
        _print_info("Optimizer pack:", str(optimizer_meta.get("pack_version", "n/a")))
        _print_info("Optimizer selector:", str(optimizer_meta.get("selector_version", "n/a")))
        _print_info("Canary status:", str(optimizer_meta.get("canary_status", "shadow")))
        _print_info("Rollback reason:", str(optimizer_meta.get("rollback_reason", "")) or "none")

    print()
    print(f"  {Colors.BOLD}Model{Colors.RESET}")
    print(f"  {'─'*50}")
    _print_info("Name:", info.get("name") or "(unnamed)")
    _print_info("Parameters:", f"{info.get('parameter_count', 0):,}")
    _print_info("License:", info.get("license", "unknown"))
    print()


# ─── OPTIMIZE-STATUS ──────────────────────────────────────────────────────

def cmd_optimize_status(args):
    """Show local self-optimizer status/state."""
    print(_logo())
    print(f"\n  {Colors.BOLD}OPTIMIZE-STATUS{Colors.RESET}\n")

    from model_osmosis.optimizer.self_optimize import get_optimizer_status

    payload = get_optimizer_status()
    if args.json:
        print(json.dumps(payload, indent=2, sort_keys=True))
        return

    _print_info("Optimizer mode:", str(payload.get("optimizer_mode", "auto")))
    _print_info("State path:", str(payload.get("state_path", "n/a")))
    _print_info("Updated at:", str(payload.get("updated_at", "")) or "n/a")
    rollback = payload.get("rollback_marker", {}) or {}
    if isinstance(rollback, dict) and rollback.get("subsystem"):
        _print_info(
            "Last rollback:",
            f"{rollback.get('subsystem')}:{rollback.get('profile_id')} ({rollback.get('reason')})",
        )
    else:
        _print_info("Last rollback:", "none")
    print()
    for name, subsystem in (payload.get("subsystems", {}) or {}).items():
        if not isinstance(subsystem, dict):
            continue
        print(f"  {Colors.BOLD}{name}{Colors.RESET}")
        _print_info("Active profile:", str(subsystem.get("active_profile_id", "")) or "n/a")
        _print_info("Promoted profile:", str(subsystem.get("promoted_profile_id", "")) or "n/a")
        _print_info("Canary status:", str(subsystem.get("canary_status", "shadow")))
        _print_info("Rollback reason:", str(subsystem.get("rollback_reason", "")) or "none")
        _print_info("Pack version:", str(subsystem.get("pack_version", "")) or "n/a")
        _print_info("Selector version:", str(subsystem.get("selector_version", "")) or "n/a")
        _print_info("Last latency ms:", f"{float(subsystem.get('selection_latency_ms', 0.0) or 0.0):.4f}")
        _print_info("Runs:", str(int(subsystem.get("runs", 0) or 0)))
        print()


# ─── OPTIMIZE-RESET ───────────────────────────────────────────────────────

def cmd_optimize_reset(args):
    """Reset local self-optimizer state only."""
    print(_logo())
    print(f"\n  {Colors.BOLD}OPTIMIZE-RESET{Colors.RESET}\n")

    from model_osmosis.optimizer.self_optimize import reset_optimizer_state

    payload = reset_optimizer_state()
    _print_step("✓", "Local optimizer state reset.")
    if args.json:
        print(json.dumps(payload, indent=2, sort_keys=True))
    else:
        _print_info("State path:", str(payload.get("state_path", "n/a")))
        _print_info("Optimizer mode:", str(payload.get("optimizer_mode", "auto")))
    print()


# ─── RELEASE-CHECK ────────────────────────────────────────────────────────

def _default_release_local_gates() -> List[Dict[str, Any]]:
    py = sys.executable
    return [
        {
            "name": "pipeline-local",
            "command": [
                py,
                "-m",
                "pytest",
                "-q",
                "tests/test_cli_new_commands.py",
                "-k",
                "pipeline_parser_merge_shape or pipeline_parser_runtime_flags",
            ],
            "timeout_seconds": 600,
        },
        {
            "name": "trust-local",
            "command": [
                py,
                "-m",
                "pytest",
                "-q",
                "tests/test_trusted_osmoimg.py",
                "-k",
                "signed_artifact_is_trusted",
            ],
            "timeout_seconds": 600,
        },
        {
            "name": "runtime-local",
            "command": [
                py,
                "-m",
                "pytest",
                "-q",
                "tests/test_cli_new_commands.py",
                "-k",
                "benchmark_runtime_outputs_json_schema",
            ],
            "timeout_seconds": 600,
        },
        {
            "name": "crossfw-core-local",
            "command": [
                py,
                "-m",
                "pytest",
                "-q",
                "tests/test_cross_framework_reliability.py",
                "-m",
                "crossfw_core",
                "--maxfail=1",
            ],
            "timeout_seconds": 1800,
        },
    ]


def _run_release_local_gate(gate: Dict[str, Any]) -> Dict[str, Any]:
    started = time.perf_counter()
    cmd = list(gate.get("command", []))
    timeout_seconds = int(gate.get("timeout_seconds", 600))
    try:
        proc = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout_seconds,
        )
        elapsed = time.perf_counter() - started
        stdout_tail = (proc.stdout or "").strip().splitlines()[-20:]
        stderr_tail = (proc.stderr or "").strip().splitlines()[-20:]
        return {
            "name": gate.get("name", "unknown"),
            "passed": proc.returncode == 0,
            "returncode": int(proc.returncode),
            "elapsed_seconds": round(float(elapsed), 3),
            "command": cmd,
            "stdout_tail": stdout_tail,
            "stderr_tail": stderr_tail,
            "reason": "" if proc.returncode == 0 else f"command_failed:{proc.returncode}",
        }
    except subprocess.TimeoutExpired:
        elapsed = time.perf_counter() - started
        return {
            "name": gate.get("name", "unknown"),
            "passed": False,
            "returncode": None,
            "elapsed_seconds": round(float(elapsed), 3),
            "command": cmd,
            "stdout_tail": [],
            "stderr_tail": [],
            "reason": f"timeout>{timeout_seconds}s",
        }
    except Exception as exc:
        elapsed = time.perf_counter() - started
        return {
            "name": gate.get("name", "unknown"),
            "passed": False,
            "returncode": None,
            "elapsed_seconds": round(float(elapsed), 3),
            "command": cmd,
            "stdout_tail": [],
            "stderr_tail": [],
            "reason": str(exc),
        }


def _current_git_sha() -> str:
    try:
        proc = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            capture_output=True,
            text=True,
            timeout=10,
        )
    except Exception:
        return ""
    if proc.returncode != 0:
        return ""
    return (proc.stdout or "").strip()


def _collect_release_ci_status(current_sha: str) -> Dict[str, Any]:
    if not current_sha:
        return {
            "available": False,
            "reason": "Unable to determine current git commit SHA for CI recency checks",
            "workflows": {},
            "current_sha": "",
        }
    if shutil.which("gh") is None:
        return {
            "available": False,
            "reason": "`gh` CLI not found on PATH",
            "workflows": {},
            "current_sha": current_sha,
        }
    cmd = [
        "gh",
        "run",
        "list",
        "--limit",
        "100",
        "--json",
        "workflowName,status,conclusion,createdAt,updatedAt,url,headSha",
    ]
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=20)
    except Exception as exc:
        return {
            "available": False,
            "reason": f"Failed to execute gh CLI: {exc}",
            "workflows": {},
            "current_sha": current_sha,
        }
    if proc.returncode != 0:
        return {
            "available": False,
            "reason": (proc.stderr or proc.stdout or "").strip() or "gh command failed",
            "workflows": {},
            "current_sha": current_sha,
        }
    try:
        runs = json.loads(proc.stdout or "[]")
    except Exception:
        return {
            "available": False,
            "reason": "Unable to parse gh run list output",
            "workflows": {},
            "current_sha": current_sha,
        }

    latest: Dict[str, Dict[str, Any]] = {}
    for wf_name in RELEASE_REQUIRED_WORKFLOWS:
        normalized = wf_name.strip().lower()
        for run in runs:
            run_name = str(run.get("workflowName", "")).strip().lower()
            run_sha = str(run.get("headSha", "")).strip().lower()
            if run_name == normalized and run_sha == current_sha.lower():
                latest[wf_name] = {
                    "status": str(run.get("status", "unknown")),
                    "conclusion": str(run.get("conclusion", "")),
                    "url": str(run.get("url", "")),
                    "updated_at": str(run.get("updatedAt", "")),
                    "head_sha": str(run.get("headSha", "")),
                }
                break
        if wf_name not in latest:
            latest[wf_name] = {
                "status": "missing",
                "conclusion": "",
                "url": "",
                "updated_at": "",
                "head_sha": "",
            }

    passed = all(
        item.get("status") == "completed"
        and item.get("conclusion") == "success"
        and str(item.get("head_sha", "")).lower() == current_sha.lower()
        for item in latest.values()
    )
    return {
        "available": True,
        "reason": "",
        "passed": passed,
        "workflows": latest,
        "current_sha": current_sha,
    }


def cmd_release_check(args):
    """Run aggregated release-ready checks with strict CI status validation."""
    print(_logo())
    print(f"\n  {Colors.BOLD}RELEASE-CHECK{Colors.RESET}\n")

    if args.local_only:
        args.with_ci = False
        args.strict_ci = False
    elif args.strict_ci and not args.with_ci:
        args.with_ci = True

    _print_info("Local only:", "yes" if args.local_only else "no")
    _print_info("With CI:", "yes" if args.with_ci else "no")
    _print_info("Strict CI:", "yes" if args.strict_ci else "no")
    _print_info("JSON:", "yes" if args.json else "no")
    print()

    started = time.perf_counter()
    local_results = [_run_release_local_gate(g) for g in _default_release_local_gates()]
    failed_gates = [r["name"] for r in local_results if not r.get("passed")]

    ci_result: Dict[str, Any] = {
        "checked": False,
        "available": False,
        "passed": True,
        "reason": "",
        "workflows": {},
        "current_sha": "",
    }
    if args.with_ci and not args.local_only:
        ci_raw = _collect_release_ci_status(_current_git_sha())
        ci_result = {
            "checked": True,
            "available": bool(ci_raw.get("available", False)),
            "passed": bool(ci_raw.get("passed", False)),
            "reason": str(ci_raw.get("reason", "")),
            "workflows": ci_raw.get("workflows", {}),
            "current_sha": str(ci_raw.get("current_sha", "")),
        }
        if not ci_result["available"] and args.strict_ci:
            failed_gates.append("ci-status-unavailable")
        elif ci_result["available"] and not ci_result["passed"]:
            failed_gates.append("ci-release-workflows")

    failed_gates = sorted(set(failed_gates))
    passed = len(failed_gates) == 0
    payload = {
        "schema_version": 1,
        "checked_at": _utc_now(),
        "elapsed_seconds": round(float(time.perf_counter() - started), 3),
        "status": "PASS" if passed else "FAIL",
        "failed_gates": failed_gates,
        "local_gates": local_results,
        "ci": ci_result,
    }

    if args.json:
        print(json.dumps(payload, indent=2, sort_keys=True))
    else:
        _print_info("Status:", f"{Colors.GREEN}PASS{Colors.RESET}" if passed else f"{Colors.RED}FAIL{Colors.RESET}")
        _print_info("Failed gates:", ", ".join(payload["failed_gates"]) if payload["failed_gates"] else "none")
        for gate in payload["local_gates"]:
            marker = "✓" if gate.get("passed") else "✗"
            _print_info(
                f"Gate {gate['name']}:",
                f"{marker} {gate.get('reason') or 'ok'} ({gate.get('elapsed_seconds', 0):.2f}s)",
            )
        if ci_result.get("checked"):
            if ci_result.get("available"):
                _print_info(
                    "CI status:",
                    "PASS" if ci_result.get("passed") else "FAIL",
                )
            else:
                _print_info("CI status:", f"unavailable ({ci_result.get('reason')})")
        print()

    if not passed:
        raise RuntimeError("release-check failed: " + ", ".join(sorted(failed_gates)))


# ─── VERSION ──────────────────────────────────────────────────────────────

def cmd_version(args):
    """Print version."""
    import model_osmosis
    print(f"Model Osmosis v{model_osmosis.__version__}")


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


@dataclass
class LoadedModelRef:
    raw: str
    scheme: str
    value: str
    model: Any


@dataclass
class ExpandModelRef:
    raw: str
    scheme: str
    value: str
    model: Any
    route: str = "unknown"
    runtime_source: str = "unknown"


def _parse_model_ref(raw: str) -> Tuple[str, str]:
    """Parse `scheme:value` model references for unified pipeline flow."""
    if ":" in raw:
        prefix, rest = raw.split(":", 1)
        if prefix in PIPELINE_SCHEMES and prefix != "auto" and rest:
            return prefix, rest
    return "auto", raw


def _parse_expand_ref(raw: str) -> Tuple[str, str]:
    """Parse `scheme:value` model references for universal expand flow."""
    if ":" in raw:
        prefix, rest = raw.split(":", 1)
        if prefix in EXPAND_SCHEMES and prefix != "auto" and rest:
            return prefix, rest
    return "auto", raw


def _auto_detect_model_scheme(value: str) -> str:
    path = Path(value)
    suffix = path.suffix.lower()
    if suffix in {".pt", ".pth", ".bin"}:
        return "pt"
    if suffix == ".onnx":
        return "onnx"
    if path.is_dir():
        if (path / "saved_model.pb").exists():
            return "tf"
        if (path / "config.json").exists():
            return "hf"
        return "pt"
    if "/" in value and not path.exists():
        return "hf"
    return "pt"


def _auto_detect_expand_scheme(value: str) -> str:
    path = Path(value)
    suffix = path.suffix.lower()
    if suffix == ".osmoimg":
        return "osmoimg"
    if suffix in {".osmo", ".zip"}:
        return "osmo"
    if suffix == ".onnx":
        return "onnx"
    if suffix in {".h5", ".keras"}:
        return "tf"
    if suffix in {".npz", ".npy"} and "jax" in path.stem.lower():
        return "jax"
    if suffix in {".pt", ".pth", ".bin", ".safetensors"}:
        return "pt"
    if path.is_dir():
        if (path / "runtime.json").exists() and (path / "manifest.json").exists():
            return "osmo"
        if (path / "saved_model.pb").exists():
            return "tf"
        if (path / "config.json").exists():
            return "hf"
        return "pt"
    if "/" in value and not path.exists():
        return "hf"
    return "pt"


def _load_hf_model(model_id: str) -> Any:
    try:
        from transformers import AutoModelForCausalLM, AutoModel
    except ImportError as exc:
        raise RuntimeError(
            "HuggingFace loading requires transformers. "
            "Install with: pip install \"osmosis[huggingface]\""
        ) from exc

    _print_step("📥", f"Loading {model_id}...")
    try:
        return AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
    except Exception:
        return AutoModel.from_pretrained(model_id, trust_remote_code=True)


def _load_jax_params(path: str) -> Any:
    import pickle
    import numpy as np

    path_obj = Path(path)
    if not path_obj.exists():
        raise FileNotFoundError(f"JAX params file not found: {path_obj}")

    suffix = path_obj.suffix.lower()
    if suffix == ".npz":
        with np.load(path_obj, allow_pickle=True) as data:
            return {k: data[k] for k in data.files}
    if suffix == ".npy":
        data = np.load(path_obj, allow_pickle=True)
        if isinstance(data, np.ndarray) and data.shape == () and data.dtype == object:
            return data.item()
        return {"params": data}
    with open(path_obj, "rb") as fh:
        return pickle.load(fh)


def _load_model_ref(raw: str) -> LoadedModelRef:
    scheme, value = _parse_model_ref(raw)
    if scheme == "auto":
        scheme = _auto_detect_model_scheme(value)

    if scheme == "hf":
        model = _load_hf_model(value)
    elif scheme == "pt":
        model = _load_model(value)
    elif scheme == "onnx":
        try:
            from model_osmosis.adapters.onnx import ONNXAdapter
        except ImportError as exc:
            raise RuntimeError(
                "ONNX loading requires onnxruntime. "
                "Install with: pip install \"osmosis[onnx]\""
            ) from exc
        adapter = ONNXAdapter(verbose=False)
        model = adapter.load(value, name=Path(value).stem)
    elif scheme == "tf":
        try:
            from model_osmosis.adapters.tensorflow import TensorFlowAdapter
        except ImportError as exc:
            raise RuntimeError(
                "TensorFlow loading requires tensorflow. "
                "Install with: pip install \"osmosis[tensorflow]\""
            ) from exc
        adapter = TensorFlowAdapter()
        model = adapter.load(value)
    elif scheme == "jax":
        model = _load_jax_params(value)
    else:
        raise ValueError(f"Unsupported model reference scheme: {scheme}")

    return LoadedModelRef(raw=raw, scheme=scheme, value=value, model=model)


def _load_expand_model_ref(raw: str) -> ExpandModelRef:
    """Load universal expand input refs including adapters and artifacts."""
    scheme, value = _parse_expand_ref(raw)
    if scheme == "auto":
        scheme = _auto_detect_expand_scheme(value)

    if scheme in {"hf", "pt", "onnx", "tf", "jax"}:
        # Reuse existing loader contract for core schemes.
        loaded = _load_model_ref(f"{scheme}:{value}" if ":" not in raw else raw)
        route = {
            "hf": "native_torch",
            "pt": "native_torch",
            "onnx": "onnx_bridge",
            "tf": "tf_bridge",
            "jax": "jax_bridge",
        }.get(scheme, "unknown")
        source = "native" if scheme in {"hf", "pt"} else "converted"
        return ExpandModelRef(
            raw=raw,
            scheme=loaded.scheme,
            value=loaded.value,
            model=loaded.model,
            route=route,
            runtime_source=source,
        )

    if scheme == "osmoimg":
        from model_osmosis.format.image import OsmoImage

        try:
            image = OsmoImage.load(value)
            model = image.export("pytorch", strict_load=True, allow_lossy_cast=False)
        except Exception as exc:
            raise _error_with_next_command(
                f"Failed to bridge .osmoimg into strict expand runtime: {type(exc).__name__}",
                f"osmosis export {value} --format pytorch -o expanded_source.pt",
            ) from exc
        return ExpandModelRef(
            raw=raw,
            scheme=scheme,
            value=value,
            model=model,
            route="artifact_bridge",
            runtime_source="artifact",
        )

    if scheme == "osmo":
        from model_osmosis.format.converters import UniversalConverter
        from model_osmosis.format.serialization import load_osmo

        try:
            osmo = load_osmo(value)
            converter = UniversalConverter()
            model = converter.from_osmo(
                osmo,
                target_framework="pytorch",
                strict_load=True,
                allow_lossy_cast=False,
            )
        except Exception as exc:
            raise _error_with_next_command(
                f"Failed to bridge .osmo into strict expand runtime: {type(exc).__name__}",
                f"osmosis convert {value} --to pytorch -o expanded_source.pt",
            ) from exc
        return ExpandModelRef(
            raw=raw,
            scheme=scheme,
            value=value,
            model=model,
            route="artifact_bridge",
            runtime_source="artifact",
        )

    if scheme in {"vision", "audio", "multimodal"}:
        from model_osmosis import get_adapter

        adapter_cls = get_adapter(scheme)
        adapter = adapter_cls(verbose=False)
        expert = adapter.load(value)
        model = expert
        if not hasattr(model, "state_dict") and hasattr(expert, "model"):
            model = getattr(expert, "model")
        return ExpandModelRef(
            raw=raw,
            scheme=scheme,
            value=value,
            model=model,
            route="adapter_torch",
            runtime_source=f"adapter:{scheme}",
        )

    if scheme == "diffusion":
        from model_osmosis.adapters.diffusion import DiffusionAdapter

        adapter = DiffusionAdapter(verbose=False)
        model = adapter.load_unet(value)
        return ExpandModelRef(
            raw=raw,
            scheme=scheme,
            value=value,
            model=model,
            route="adapter_torch",
            runtime_source="adapter:diffusion",
        )

    if scheme == "lora":
        if "::" not in value:
            raise _error_with_next_command(
                "LoRA expand input must be formatted as lora:<base_model>::<lora_adapter_path>",
                "osmosis expand lora:<base_model>::<lora_adapter_path> --fallback-mode distill",
            )
        base_model, lora_path = value.split("::", 1)
        from model_osmosis.adapters.lora import LoRAAdapter

        adapter = LoRAAdapter(verbose=False)
        expert = adapter.load_with_lora(base_model=base_model, lora_path=lora_path, merge_lora=True)
        return ExpandModelRef(
            raw=raw,
            scheme=scheme,
            value=value,
            model=expert,
            route="adapter_torch",
            runtime_source="adapter:lora",
        )

    if scheme == "mlx":
        from model_osmosis.adapters.mlx import MLXAdapter

        adapter = MLXAdapter(verbose=False)
        model = adapter.load(value)
        return ExpandModelRef(
            raw=raw,
            scheme=scheme,
            value=value,
            model=model,
            route="jax_bridge",
            runtime_source="adapter:mlx",
        )

    raise ValueError(f"Unsupported expand model reference scheme: {scheme}")


def _parse_onnx_input_shape(value: Optional[str]) -> Optional[Tuple[int, ...]]:
    if value is None:
        return None
    cleaned = value.replace("x", ",").replace(" ", "")
    parts = [p for p in cleaned.split(",") if p]
    if not parts:
        raise ValueError("Invalid --onnx-input-shape; expected e.g. 1,128")
    shape = tuple(int(p) for p in parts)
    if any(dim <= 0 for dim in shape):
        raise ValueError("All --onnx-input-shape dimensions must be positive")
    return shape


def _load_sample_data(path: Optional[str]) -> Optional[Any]:
    if not path:
        return None
    import numpy as np

    sample_path = Path(path)
    if not sample_path.exists():
        raise FileNotFoundError(f"Sample data file not found: {sample_path}")
    data = np.load(sample_path, allow_pickle=True)
    return data


def _infer_expand_feature_dim(model: Any) -> int:
    try:
        import torch.nn as nn

        if isinstance(model, nn.Module):
            for module in model.modules():
                if isinstance(module, nn.Linear):
                    return int(max(1, int(module.in_features)))
    except Exception:
        pass
    return 8


def _normalize_expand_name(raw: str) -> str:
    cleaned = "".join(ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in str(raw))
    return cleaned.strip("_") or "expanded_model"


def _default_expand_name(ref: ExpandModelRef) -> str:
    if ref.scheme == "hf":
        return _normalize_expand_name(Path(ref.value).name)
    if ref.value:
        return _normalize_expand_name(Path(ref.value).stem)
    return _normalize_expand_name(ref.raw)


def _expand_route_next_command(loaded_ref: ExpandModelRef) -> str:
    scheme = str(loaded_ref.scheme or "").strip().lower()
    value = str(loaded_ref.value or loaded_ref.raw or "").strip()
    if scheme == "onnx":
        return f"osmosis convert onnx:{value} --to pytorch -o expanded_source.pt"
    if scheme == "tf":
        return f"osmosis convert tf:{value} --to pytorch -o expanded_source.pt"
    if scheme == "jax":
        return f"osmosis convert jax:{value} --to pytorch -o expanded_source.pt"
    if scheme == "mlx":
        return f"osmosis convert mlx:{value} --to pytorch -o expanded_source.pt"
    if scheme == "osmoimg":
        return f"osmosis export {value} --format pytorch -o expanded_source.pt"
    if scheme == "osmo":
        return f"osmosis convert {value} --to pytorch -o expanded_source.pt"
    if scheme in {"vision", "audio", "diffusion", "multimodal", "lora"}:
        return f"osmosis expand {loaded_ref.raw} --fallback-mode distill"
    return "osmosis expand pt:<model.pt> --fallback-mode distill"


def _prepare_expand_runtime_model(
    loaded_ref: ExpandModelRef,
    *,
    allow_lossy_cast: bool = False,
) -> Tuple[Any, str, str]:
    """Return a torch-compatible runtime model for expansion."""
    model = loaded_ref.model
    try:
        import torch.nn as nn
    except ImportError as exc:
        raise RuntimeError("Expansion requires torch installed. Install with: pip install torch") from exc

    if isinstance(model, nn.Module):
        return model, str(loaded_ref.runtime_source or "native"), str(loaded_ref.route or "native_torch")

    from model_osmosis.format.converters import UniversalConverter

    converter = UniversalConverter()
    osmo_name = Path(loaded_ref.value).stem if loaded_ref.value else "expand_source"
    try:
        osmo = converter.to_osmo(model, name=osmo_name, allow_lossy_cast=allow_lossy_cast)
    except Exception as exc:
        raise _error_with_next_command(
            (
                f"Expand route {loaded_ref.route} cannot convert this input into strict-expand runtime: "
                f"{type(exc).__name__}"
            ),
            _expand_route_next_command(loaded_ref),
        ) from exc
    reconstructed = converter.from_osmo(
        osmo,
        target_framework="pytorch",
        strict_load=True,
        allow_lossy_cast=allow_lossy_cast,
    )
    if isinstance(reconstructed, nn.Module):
        route = str(loaded_ref.route or "unknown")
        if route == "unknown":
            route = "artifact_bridge"
        return reconstructed, "converted", route
    raise _error_with_next_command(
        "Expansion conversion produced a non-runnable runtime for strict expansion.",
        _expand_route_next_command(loaded_ref),
    )


def _expand_export_suffix(fmt: str) -> str:
    if fmt == "osmoimg":
        return ".osmoimg"
    if fmt == "pytorch":
        return ".pt"
    if fmt == "hf":
        return "_hf"
    if fmt == "onnx":
        return ".onnx"
    if fmt == "tensorflow":
        return ".npz"
    if fmt == "jax":
        return ".npz"
    return ".bin"


def _state_dict_like(model: Any) -> Optional[Dict[str, Any]]:
    if isinstance(model, dict):
        return model
    if hasattr(model, "state_dict"):
        return model.state_dict()
    return None


def _to_numpy_state_dict(model: Any) -> Optional[Dict[str, np.ndarray]]:
    state = _state_dict_like(model)
    if state is None:
        return None
    out: Dict[str, np.ndarray] = {}
    for key, value in state.items():
        arr = value
        if hasattr(arr, "detach"):
            arr = arr.detach().cpu().numpy()
        arr_np = np.asarray(arr)
        out[str(key)] = arr_np
    return out


def _models_are_directly_mergeable(models: Sequence[Any]) -> bool:
    if len(models) < 2:
        return True

    first = _state_dict_like(models[0])
    if first is None:
        return False

    first_keys = list(first.keys())
    first_shapes = {k: tuple(first[k].shape) for k in first_keys}
    for model in models[1:]:
        state = _state_dict_like(model)
        if state is None:
            return False
        if list(state.keys()) != first_keys:
            return False
        for key in first_keys:
            if tuple(state[key].shape) != first_shapes[key]:
                return False
    return True


def _run_merge_backend(
    pipeline: Any,
    *,
    loaded_models: Sequence[Any],
    strategy: str,
    weights: Optional[List[float]],
    sample_data: Optional[Any],
    source_refs: Optional[Sequence[Any]] = None,
    force_cross: bool = False,
    epochs: int = 3,
    seed: int = 42,
    strategy_engine_mode: str = "auto",
    cross_opt_stage: str = "auto",
    return_metadata: bool = False,
) -> Any:
    from model_osmosis.merging.smart_merge_engine import (
        collect_tensor_groups_from_state_dicts,
        default_selection,
        select_recipe_for_tensor_groups,
    )
    from model_osmosis.merging.cross_merge_engine import (
        select_cross_merge_recipe,
    )
    from model_osmosis.merging.strategies import get_strategy

    normalized_engine_mode = _resolve_strategy_engine_mode(strategy_engine_mode)
    requested_strategy = str(strategy or "weighted_average").strip().lower()

    selection = default_selection(
        strategy_name=requested_strategy if requested_strategy else "weighted_average",
        strategy_engine_mode=normalized_engine_mode,
        recipe_id=f"explicit::{requested_strategy or 'weighted_average'}",
        fallback_reason="explicit_strategy",
    )

    def _labels_for_selection() -> List[str]:
        labels: List[str] = []
        refs = list(source_refs) if source_refs else [f"model_{i + 1}" for i in range(len(loaded_models))]
        for idx, ref in enumerate(refs):
            if isinstance(ref, str):
                labels.append(ref)
            else:
                labels.append(f"model_{idx + 1}")
        return labels

    if force_cross:
        cross_models = list(loaded_models)
        cross_selection = select_cross_merge_recipe(
            models=cross_models,
            sample_data=sample_data,
            weights=weights,
            seed=seed,
            strategy_engine_mode=normalized_engine_mode,
            cross_opt_stage=cross_opt_stage,
            source_refs=_labels_for_selection(),
            epochs=epochs,
        )
        result = pipeline.cross_merge(
            models=cross_models,
            sample_data=sample_data,
            weights=weights,
            epochs=epochs,
            seed=seed,
            **dict(cross_selection.cross_kwargs or {}),
        )
        selection = default_selection(
            strategy_name="cross-merge",
            strategy_kwargs=dict(cross_selection.cross_kwargs or {}),
            strategy_engine_mode=normalized_engine_mode,
            recipe_id=str(cross_selection.recipe_id),
            recipe_pack_version=str(cross_selection.recipe_pack_version),
            fallback_reason=str(cross_selection.fallback_reason or "forced_cross_mode"),
            canary_status=str(cross_selection.canary_status),
            rollback_reason=str(cross_selection.rollback_reason),
        )
        if return_metadata:
            report = selection.to_report_dict()
            report["cross_opt_stage"] = str(cross_selection.cross_opt_stage)
            report["is_llm"] = bool(cross_selection.is_llm)
            return (result, report)
        return result

    if _models_are_directly_mergeable(loaded_models):
        strategy_for_merge: Any = strategy
        if requested_strategy == "auto" and normalized_engine_mode == "auto":
            numpy_states = [_to_numpy_state_dict(model) for model in loaded_models]
            if all(state is not None for state in numpy_states):
                tensor_groups = collect_tensor_groups_from_state_dicts(
                    [state for state in numpy_states if state is not None],
                    seed=seed,
                )
            else:
                tensor_groups = []
            selection = select_recipe_for_tensor_groups(
                tensor_groups=tensor_groups,
                model_count=len(loaded_models),
                seed=seed,
                strategy_engine_mode=normalized_engine_mode,
                user_weights_provided=weights is not None,
                context_label="merge",
                source_refs=_labels_for_selection(),
            )
            selected_name = selection.strategy_name
            selected_kwargs = dict(selection.strategy_kwargs or {})
            if selected_kwargs:
                if selected_name == "dare" and "seed" not in selected_kwargs:
                    selected_kwargs["seed"] = int(seed)
                strategy_for_merge = get_strategy(selected_name, **selected_kwargs)
            else:
                strategy_for_merge = selected_name
        elif requested_strategy == "auto" and normalized_engine_mode == "off":
            selection = default_selection(
                strategy_name="auto",
                strategy_engine_mode=normalized_engine_mode,
                recipe_id="legacy::auto",
                fallback_reason="strategy_engine_off",
            )

        eval_data = sample_data
        if isinstance(strategy_for_merge, str) and strategy_for_merge == "auto" and eval_data is None:
            import numpy as np
            from model_osmosis.utils.determinism import make_numpy_rng

            rng = make_numpy_rng(seed)
            eval_data = rng.random((50, 10), dtype=np.float32)
        result = pipeline.merge(
            models=list(loaded_models),
            strategy=strategy_for_merge,
            model_weights=weights,
            eval_data=eval_data,
        )
        return (result, selection.to_report_dict()) if return_metadata else result

    _print_step("🔀", "Models are not directly compatible; auto-routing to cross_merge...")
    cross_models = list(loaded_models)
    cross_sample_data = sample_data
    if cross_models and isinstance(cross_models[0], str) and not isinstance(sample_data, list):
        cross_sample_data = None
    cross_selection = select_cross_merge_recipe(
        models=cross_models,
        sample_data=cross_sample_data,
        weights=weights,
        seed=seed,
        strategy_engine_mode=normalized_engine_mode,
        cross_opt_stage=cross_opt_stage,
        source_refs=_labels_for_selection(),
        epochs=epochs,
    )
    result = pipeline.cross_merge(
        models=cross_models,
        sample_data=cross_sample_data,
        weights=weights,
        epochs=epochs,
        seed=seed,
        **dict(cross_selection.cross_kwargs or {}),
    )
    selection = default_selection(
        strategy_name="cross-merge",
        strategy_kwargs=dict(cross_selection.cross_kwargs or {}),
        strategy_engine_mode=normalized_engine_mode,
        recipe_id=str(cross_selection.recipe_id),
        recipe_pack_version=str(cross_selection.recipe_pack_version),
        fallback_reason=str(cross_selection.fallback_reason or "state_dict_incompatible_or_non_stateful_models"),
        canary_status=str(cross_selection.canary_status),
        rollback_reason=str(cross_selection.rollback_reason),
    )
    if return_metadata:
        report = selection.to_report_dict()
        report["cross_opt_stage"] = str(cross_selection.cross_opt_stage)
        report["is_llm"] = bool(cross_selection.is_llm)
        return (result, report)
    return result


def _resolve_merge_route(
    loaded_models: Sequence[Any],
    *,
    force_cross: bool = False,
) -> Tuple[str, str]:
    if force_cross:
        return "cross-merge", "forced_cross_mode"
    if _models_are_directly_mergeable(loaded_models):
        return "direct-merge", "state_dict_keys_and_shapes_match"
    return "cross-merge", "state_dict_incompatible_or_non_stateful_models"


def _run_weight_merge_backend(
    *,
    loaded_refs: Sequence[LoadedModelRef],
    strategy: str,
    weights: Optional[List[float]],
    coverage_floor: float,
    output_type: str,
    seed: int = 42,
    allow_lossy_cast: bool = False,
    strategy_engine_mode: str = "auto",
) -> Dict[str, Any]:
    from model_osmosis.format.weight_merge import merge_weights_to_osmo

    try:
        result = merge_weights_to_osmo(
            [item.model for item in loaded_refs],
            source_refs=[item.raw for item in loaded_refs],
            strategy=strategy,
            weights=weights,
            coverage_floor=coverage_floor,
            seed=seed,
            allow_lossy_cast=allow_lossy_cast,
            strategy_engine_mode=_resolve_strategy_engine_mode(strategy_engine_mode),
        )
    except RuntimeError as exc:
        message = str(exc)
        if "coverage floor not met" in message.lower():
            models_cmd = " ".join(item.raw for item in loaded_refs)
            suggested_floor = max(0.0, float(coverage_floor) - 0.2)
            raise RuntimeError(
                f"{message} "
                f"Next command: osmosis pipeline --mode weight-merge --output-type {output_type} "
                f"--models {models_cmd} --coverage-floor {suggested_floor:.2f}"
            ) from exc
        raise
    return {
        "merged_osmo": result.merged_osmo,
        "report": dict(result.report or {}),
    }


def _reconstruct_torch_module_for_weight_merge(
    *,
    merged_osmo: Any,
    loaded_refs: Sequence[LoadedModelRef],
    strict_load: bool,
    allow_lossy_cast: bool = False,
) -> Tuple[Any, Dict[str, Any]]:
    try:
        import copy
        import torch.nn as nn
    except ImportError as exc:
        raise RuntimeError("Reconstructing torch module requires torch installed.") from exc

    template = None
    for ref in loaded_refs:
        if isinstance(ref.model, nn.Module):
            template = ref.model
            break

    if template is None:
        raise RuntimeError(
            "Weight-merge export to onnx/torchscript requires at least one torch/HF input model "
            "to provide architecture template."
        )

    from model_osmosis.format.converters import UniversalConverter

    converter = UniversalConverter()
    state_dict = converter.from_osmo(
        merged_osmo,
        target_framework="pytorch",
        strict_load=False,
        allow_lossy_cast=allow_lossy_cast,
    )
    model = copy.deepcopy(template).eval()
    load_result = model.load_state_dict(state_dict, strict=False)
    missing = list(load_result.missing_keys)
    unexpected = list(load_result.unexpected_keys)

    report = dict(converter.last_report or {})
    report["missing_keys"] = missing
    report["unexpected_keys"] = unexpected
    report.setdefault("lossy_casts", [])
    report.setdefault("dtype_transports", [])
    if strict_load and (missing or unexpected):
        raise RuntimeError(
            "Weight-merge torch reconstruction is partial for this architecture template. "
            "Re-run with --allow-partial-load or export to pytorch/safetensors/osmo."
        )
    return model, report


def _pipeline_preflight_requirements(args) -> None:
    missing: List[Tuple[str, str]] = []

    def _need(module_name: str, install_hint: str) -> None:
        if not _is_module_available(module_name):
            missing.append((module_name, install_hint))

    if args.output_type in {"pytorch", "onnx", "torchscript", "safetensors"}:
        _need("torch", "pip install torch")
    if args.output_type == "hf":
        _need("transformers", "pip install \"osmosis[huggingface]\"")
    if args.output_type == "onnx":
        _need("onnxruntime", "pip install \"osmosis[onnx]\"")
    if args.output_type in {"tensorflow", "keras"}:
        _need("tensorflow", "pip install \"osmosis[tensorflow]\"")
    if args.output_type in {"jax", "flax"}:
        _need("jax", "pip install \"osmosis[jax]\"")
    if args.output_type == "safetensors":
        _need("safetensors", "pip install safetensors")
    if args.output_type == "osmoimg" and args.compression_profile == "lossless-max":
        _need("zstandard", "pip install \"osmosis[compression]\"")

    if missing:
        module, hint = missing[0]
        raise RuntimeError(
            f"Missing required dependency '{module}' for pipeline output-type={args.output_type}. "
            f"Install with: {hint}"
        )

    if args.output_type == "onnx":
        if _parse_onnx_input_shape(args.onnx_input_shape) is None:
            raise RuntimeError(
                "ONNX output requires --onnx-input-shape, e.g. --onnx-input-shape 1,128"
            )
    if args.output_type == "hf" and args.mode != "merge" and not args.hf_base_model:
        raise RuntimeError(
            "HF output for non-merge flows requires --hf-base-model when result is not "
            "already a HuggingFace model instance."
        )


def _pipeline_preflight(
    args,
    *,
    loaded_refs: Sequence[LoadedModelRef],
) -> Dict[str, Any]:
    _pipeline_preflight_requirements(args)

    route = "unknown"
    route_reason = "unknown"

    if args.mode == "merge":
        loaded_models = [item.model for item in loaded_refs]
        force_cross = getattr(args, "force_cross", False)
        if not isinstance(force_cross, bool):
            force_cross = False
        route, route_reason = _resolve_merge_route(
            loaded_models,
            force_cross=force_cross,
        )
    elif args.mode == "weight-merge":
        route = "weight-merge"
        route_reason = "exact_param_compatibility"
    elif args.mode == "absorb":
        if len(loaded_refs) != 2:
            raise RuntimeError("Absorb preflight expected exactly source and target refs.")
        _ensure_absorb_target_supported(loaded_refs[1].model, loaded_refs[1].raw)
        route = "absorb"
        route_reason = "target_absorber_supported"
    elif args.mode == "stack":
        for item in loaded_refs:
            _ensure_stack_expert(item.model, item.raw)
        route = "stack"
        route_reason = "experts_supported"
    else:
        raise ValueError(f"Unsupported mode: {args.mode}")

    resolved_inputs = [
        {"raw": item.raw, "scheme": item.scheme, "value": item.value}
        for item in sorted(loaded_refs, key=lambda x: x.raw)
    ]
    return {
        "route": route,
        "route_reason": route_reason,
        "resolved_inputs": resolved_inputs,
    }


def _resolve_pipeline_self_check_level(args) -> str:
    level = str(getattr(args, "self_check_level", "smoke") or "smoke").strip().lower()
    if not getattr(args, "self_check", True):
        return "off"
    if level not in {"off", "smoke", "parity"}:
        raise ValueError("Invalid self-check level. Expected one of: off, smoke, parity")
    return level


def _extract_parity_tensor_map(model_obj: Any) -> Optional[Dict[str, Any]]:
    if model_obj is None:
        return None

    if hasattr(model_obj, "osmo_model") and hasattr(model_obj.osmo_model, "parameters"):
        return {k: np.asarray(v) for k, v in model_obj.osmo_model.parameters.items()}

    try:
        import torch

        if isinstance(model_obj, torch.nn.Module):
            state = model_obj.state_dict()
            return {k: v.detach().cpu().numpy() for k, v in state.items()}
    except Exception:
        pass

    if hasattr(model_obj, "state_dict"):
        try:
            state = model_obj.state_dict()
            out: Dict[str, Any] = {}
            for key, value in state.items():
                if hasattr(value, "detach"):
                    out[key] = value.detach().cpu().numpy()
                else:
                    out[key] = np.asarray(value)
            return out
        except Exception:
            pass

    if isinstance(model_obj, dict):
        out = {}
        for key, value in model_obj.items():
            if hasattr(value, "detach"):
                out[str(key)] = value.detach().cpu().numpy()
            else:
                out[str(key)] = np.asarray(value)
        return out

    return None


def _parity_compare_tensor_maps(
    reference: Dict[str, Any],
    candidate: Dict[str, Any],
    *,
    rtol: float = 1e-5,
    atol: float = 1e-6,
) -> Dict[str, Any]:
    ref_keys = sorted(reference.keys())
    cand_keys = sorted(candidate.keys())
    if ref_keys != cand_keys:
        missing = sorted(set(ref_keys) - set(cand_keys))
        extra = sorted(set(cand_keys) - set(ref_keys))
        return {
            "passed": False,
            "reason": "key_mismatch",
            "missing_keys": missing,
            "unexpected_keys": extra,
            "max_error": float("inf"),
            "failing_tensors": missing[:5] + extra[:5],
        }

    max_error = 0.0
    failing: List[str] = []
    for key in ref_keys:
        ref = np.asarray(reference[key])
        cand = np.asarray(candidate[key])
        if ref.shape != cand.shape:
            failing.append(key)
            max_error = float("inf")
            continue
        if ref.dtype.kind in {"i", "u", "b"}:
            equal = np.array_equal(ref, cand)
            err = 0.0 if equal else float("inf")
        else:
            diff = np.abs(ref.astype(np.float64) - cand.astype(np.float64))
            err = float(diff.max()) if diff.size else 0.0
            equal = np.allclose(ref, cand, rtol=rtol, atol=atol, equal_nan=True)
        max_error = max(max_error, err)
        if not equal:
            failing.append(key)
            if len(failing) >= 10:
                break

    return {
        "passed": len(failing) == 0,
        "reason": "" if len(failing) == 0 else "numerical_mismatch",
        "max_error": max_error,
        "failing_tensors": failing,
        "missing_keys": [],
        "unexpected_keys": [],
    }


def _pipeline_self_check(
    *,
    output_type: str,
    output_path: Path,
    self_check_level: str = "smoke",
    reference_model: Any = None,
) -> Dict[str, Any]:
    level = str(self_check_level or "smoke").strip().lower()
    if level not in {"off", "smoke", "parity"}:
        raise RuntimeError(f"Self-check failed: invalid self_check_level={level}")
    if level == "off":
        return {
            "self_check_status": "skipped_disabled",
            "self_check_level": "off",
            "self_check_checks": [],
            "self_check_skips": [],
        }

    checks: List[str] = []
    skips: List[str] = []

    if not output_path.exists():
        raise RuntimeError(f"Self-check failed: output artifact missing at {output_path}")
    checks.append("artifact_exists")

    if output_type == "osmoimg":
        from model_osmosis.format.image import OsmoImage

        info = OsmoImage.inspect(output_path)
        verify = OsmoImage.verify(output_path)
        if verify.get("trust_status") == "tampered":
            raise RuntimeError("Self-check failed: trust verification reported tampered payload")
        file_size_mb = float(info.get("image", {}).get("file_size_mb", 0.0))
        if file_size_mb <= 0:
            raise RuntimeError("Self-check failed: inspect returned invalid file size")
        checks.extend(["inspect_ok", "verify_ok"])

        # Deterministic smoke inference for runnable artifact sanity.
        try:
            loaded = OsmoImage.load(output_path)
            model_type = str(loaded.runtime.model_type or "").strip().lower()
            if model_type in {"causal_lm", "seq2seq", "language_model"}:
                smoke_output = loaded.run(
                    prompt=RUNTIME_BENCHMARK_DEFAULT_PROMPT,
                    max_tokens=1,
                    temperature=0.0,
                    runtime_cache=False,
                )
            else:
                smoke_input = _infer_benchmark_input(loaded)
                smoke_output = loaded.run(
                    input_data=smoke_input,
                    runtime_cache=False,
                )
            if not _is_non_empty_runtime_output(smoke_output):
                raise RuntimeError("empty_output")
            checks.append("smoke_inference_ok")
        except Exception as exc:
            skips.append(f"smoke_inference_skipped:{type(exc).__name__}")
    elif output_type == "osmo":
        from model_osmosis.format.serialization import load_osmo

        osmo = load_osmo(output_path, load_parameters=False, load_knowledge=False)
        if osmo is None:
            raise RuntimeError("Self-check failed: unable to load .osmo artifact")
        checks.append("osmo_load_ok")
        skips.append("smoke_inference_not_supported:osmo")
    elif output_type == "pytorch":
        loaded_model = _load_model(str(output_path))
        checks.append("pytorch_load_ok")
        skips.append("smoke_inference_not_supported:pytorch")
    elif output_type == "hf":
        config_path = output_path / "config.json"
        if not config_path.exists():
            raise RuntimeError(
                "Self-check failed: HF export is missing config.json in output directory"
            )
        checks.append("hf_layout_ok")
        skips.append("smoke_inference_not_supported:hf")
    elif output_type == "onnx":
        _load_model(str(output_path))
        checks.append("onnx_load_ok")
        skips.append("smoke_inference_not_supported:onnx")
    elif output_type in {"tensorflow", "jax", "flax"}:
        import numpy as np

        with np.load(output_path, allow_pickle=True) as arrs:
            if len(arrs.files) == 0:
                raise RuntimeError("Self-check failed: exported NPZ contains no arrays")
        checks.append("npz_load_ok")
        skips.append(f"smoke_inference_not_supported:{output_type}")
    elif output_type == "keras":
        if not _is_module_available("tensorflow"):
            raise RuntimeError("Self-check failed: tensorflow required to load Keras artifact")
        import tensorflow as tf

        tf.keras.models.load_model(output_path)
        checks.append("keras_load_ok")
        skips.append("smoke_inference_not_supported:keras")
    elif output_type == "safetensors":
        from safetensors.numpy import load_file as load_safetensors_file

        tensors = load_safetensors_file(str(output_path))
        if not tensors:
            raise RuntimeError("Self-check failed: safetensors artifact contains no tensors")
        checks.append("safetensors_load_ok")
        skips.append("smoke_inference_not_supported:safetensors")
    elif output_type == "torchscript":
        import torch

        torch.jit.load(str(output_path))
        checks.append("torchscript_load_ok")
        skips.append("smoke_inference_not_supported:torchscript")
    else:
        raise RuntimeError(f"Self-check failed: unsupported output type {output_type}")

    if level == "parity":
        parity_supported = output_type in {"osmoimg", "pytorch"}
        if not parity_supported:
            skips.append(f"parity_not_supported:{output_type}")
        elif reference_model is None:
            skips.append("parity_reference_unavailable")
        else:
            from model_osmosis.format.image import OsmoImage

            if output_type == "osmoimg":
                candidate_model = OsmoImage.load(output_path)
            else:
                candidate_model = _load_model(str(output_path))
            ref_map = _extract_parity_tensor_map(reference_model)
            cand_map = _extract_parity_tensor_map(candidate_model)
            if ref_map is None or cand_map is None:
                skips.append("parity_tensor_map_unavailable")
            else:
                parity = _parity_compare_tensor_maps(ref_map, cand_map)
                if not parity.get("passed", False):
                    failing = [str(x) for x in (parity.get("failing_tensors", []) or [])[:3]]
                    max_error = float(parity.get("max_error", 0.0) or 0.0)
                    raise RuntimeError(
                        "Self-check parity failed: "
                        f"reason={parity.get('reason')} "
                        f"max_error={max_error:.6e} "
                        f"failing_tensors={failing}"
                    )
                checks.append("parity_ok")

    return {
        "self_check_status": "passed",
        "self_check_level": level,
        "self_check_checks": sorted(checks),
        "self_check_skips": sorted(set(skips)),
    }


def _pipeline_report_payload(
    args,
    *,
    details: Dict[str, Any],
    elapsed_seconds: float,
) -> Dict[str, Any]:
    conversion_report = details.get("conversion_report", {}) or {}
    input_refs = details.get("resolved_inputs", []) or []
    normalized_refs: List[Dict[str, Any]] = []
    for ref in input_refs:
        if isinstance(ref, dict):
            normalized_refs.append(
                {
                    "raw": str(ref.get("raw", "")),
                    "scheme": str(ref.get("scheme", "")),
                    "value": str(ref.get("value", "")),
                }
            )
        else:
            normalized_refs.append({"raw": str(ref), "scheme": "", "value": ""})
    normalized_refs = sorted(
        normalized_refs,
        key=lambda item: (item.get("raw", ""), item.get("scheme", ""), item.get("value", "")),
    )
    weight_merge_report = details.get("weight_merge_report") or {}
    normalized_weight_merge: Dict[str, Any] = {
        "compatible_tensor_count": int(weight_merge_report.get("compatible_tensor_count", 0) or 0),
        "skipped_tensor_count": int(weight_merge_report.get("skipped_tensor_count", 0) or 0),
        "merged_bytes": int(weight_merge_report.get("merged_bytes", 0) or 0),
        "min_coverage_ratio": float(weight_merge_report.get("min_coverage_ratio", 0.0) or 0.0),
        "coverage_floor": float(
            weight_merge_report.get(
                "coverage_floor",
                getattr(args, "coverage_floor", 0.0),
            ) or 0.0
        ),
        "coverage_ratio_by_bytes": {
            str(k): float(v)
            for k, v in sorted(
                (weight_merge_report.get("coverage_ratio_by_bytes", {}) or {}).items(),
                key=lambda item: str(item[0]),
            )
        },
        "skip_reason_counts": {
            "dtype_mismatch": int(
                (weight_merge_report.get("skip_reason_counts", {}) or {}).get("dtype_mismatch", 0) or 0
            ),
            "missing": int(
                (weight_merge_report.get("skip_reason_counts", {}) or {}).get("missing", 0) or 0
            ),
            "shape_mismatch": int(
                (weight_merge_report.get("skip_reason_counts", {}) or {}).get("shape_mismatch", 0) or 0
            ),
        },
    }

    report = {
        "schema_version": 1,
        "mode": args.mode,
        "output_type": args.output_type,
        "seed": int(args.seed),
        "route": details.get("route", "unknown"),
        "route_reason": details.get("route_reason", "unknown"),
        "reliability_mode": details.get("reliability_mode", "strict"),
        "fallbacks": _reliability_fallbacks(
            args.allow_partial_load,
            args.allow_lossy_cast,
        ),
        "dry_run": bool(details.get("dry_run", False)),
        "input_refs": normalized_refs,
        "output_path": str(details.get("output_path", "")),
        "protected": bool(details.get("protected", False)),
        "self_check_status": details.get("self_check_status", "skipped"),
        "self_check_level": details.get("self_check_level", "off"),
        "self_check_checks": sorted(set(details.get("self_check_checks", []) or [])),
        "self_check_skips": sorted(set(details.get("self_check_skips", []) or [])),
        "conversion_report_counts": {
            "missing_keys": len(conversion_report.get("missing_keys", [])),
            "unexpected_keys": len(conversion_report.get("unexpected_keys", [])),
            "lossy_casts": len(conversion_report.get("lossy_casts", [])),
            "dtype_transports": len(conversion_report.get("dtype_transports", [])),
        },
        "compression_profile": details.get("compression_profile", "n/a"),
        "compression_engine_mode": str(details.get("compression_engine_mode", "auto")),
        "payload_codec": details.get("payload_codec", "n/a"),
        "trust_status": details.get("trust_status", "unsigned"),
        "strategy_engine_mode": str(details.get("strategy_engine_mode", "auto")),
        "cross_opt_stage": str(details.get("cross_opt_stage", "auto")),
        "strategy_name": str(details.get("strategy_name", "")),
        "strategy_kwargs": {
            str(k): v
            for k, v in sorted((details.get("strategy_kwargs", {}) or {}).items(), key=lambda item: str(item[0]))
        },
        "optimizer_mode": str(details.get("optimizer_mode", details.get("strategy_engine_mode", "auto"))),
        "profile_id": str(details.get("profile_id", details.get("recipe_id", ""))),
        "pack_version": str(details.get("pack_version", details.get("recipe_pack_version", ""))),
        "recipe_id": str(details.get("recipe_id", "")),
        "recipe_pack_version": str(details.get("recipe_pack_version", "")),
        "selector_version": str(details.get("selector_version", "")),
        "canary_status": str(details.get("canary_status", "shadow")),
        "rollback_reason": str(details.get("rollback_reason", "")),
        "selection_latency_ms": round(float(details.get("selection_latency_ms", 0.0) or 0.0), 4),
        "fallback_reason": str(details.get("fallback_reason", "")),
        "weight_merge": normalized_weight_merge if args.mode == "weight-merge" else None,
        "elapsed_seconds": round(float(elapsed_seconds), 4),
    }
    return report


def _write_pipeline_report(report_file: str, report_payload: Dict[str, Any]) -> Path:
    path = Path(report_file)
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as fh:
        json.dump(report_payload, fh, indent=2, sort_keys=True)
    return path


def _default_pipeline_output(mode: str, output_type: str, name: Optional[str]) -> Path:
    stem = (name or f"{mode}_result").strip().replace(" ", "_")
    base = Path("models")
    if output_type == "osmoimg":
        return base / f"{stem}.osmoimg"
    if output_type == "osmo":
        return base / f"{stem}.osmo"
    if output_type == "pytorch":
        return base / f"{stem}.pt"
    if output_type == "safetensors":
        return base / f"{stem}.safetensors"
    if output_type == "torchscript":
        return base / f"{stem}.torchscript.pt"
    if output_type == "onnx":
        return base / f"{stem}.onnx"
    if output_type in {"tensorflow", "jax", "flax"}:
        return base / f"{stem}.npz"
    if output_type == "keras":
        return base / f"{stem}.keras"
    return base / f"{stem}_hf"


def _ensure_absorb_target_supported(model: Any, ref: str) -> Any:
    from model_osmosis.utils.model_utils import auto_select_absorber
    try:
        auto_select_absorber(model)
        return model
    except Exception:
        pass
    raise RuntimeError(
        "Unsupported target for absorb after auto-conversion attempt: "
        f"{ref}. Use a PyTorch/sklearn-compatible target, e.g. `pt:path/to/model.pt`."
    )


def _ensure_stack_expert(model: Any, ref: str) -> Any:
    try:
        import torch.nn as nn
    except ImportError as exc:
        raise RuntimeError("Stack mode requires torch installed.") from exc

    if isinstance(model, nn.Module):
        return model
    raise RuntimeError(
        "Unsupported stack expert after auto-conversion attempt: "
        f"{ref}. Use torch/HF/ONNX experts for stack mode."
    )


def _save_hf_output(
    model: Any,
    output_path: Path,
    *,
    hf_base_model: Optional[str],
    name: Optional[str],
    tokenizer: Optional[Any] = None,
    strict_load: bool = True,
    allow_lossy_cast: bool = False,
) -> Dict[str, Any]:
    from model_osmosis.obfuscation import obfuscate_osmo, ObfuscationConfig
    from model_osmosis.format.converters import UniversalConverter, HuggingFaceConverter

    output_path.mkdir(parents=True, exist_ok=True)
    report: Dict[str, Any] = {
        "missing_keys": [],
        "unexpected_keys": [],
        "lossy_casts": [],
        "dtype_transports": [],
    }

    if hasattr(model, "save_pretrained"):
        model.save_pretrained(output_path)
    else:
        if not hf_base_model:
            raise RuntimeError(
                "HF output requires --hf-base-model when the result is not already a "
                "HuggingFace model instance."
            )
        converter = UniversalConverter()
        osmo = converter.to_osmo(
            model,
            name=name or output_path.name,
            allow_lossy_cast=allow_lossy_cast,
        )
        osmo = obfuscate_osmo(osmo, config=ObfuscationConfig(preserve_quality=True))
        hf_converter = HuggingFaceConverter()
        hf_model = hf_converter.from_osmo(
            osmo,
            model_id=hf_base_model,
            strict_load=strict_load,
            allow_lossy_cast=allow_lossy_cast,
        )
        report = hf_converter.last_report.to_dict()
        hf_model.save_pretrained(output_path)

    if tokenizer is not None and hasattr(tokenizer, "save_pretrained"):
        tokenizer.save_pretrained(output_path)
    return report


def _save_onnx_output(model: Any, output_path: Path, onnx_input_shape: Optional[Tuple[int, ...]]) -> None:
    if onnx_input_shape is None:
        raise RuntimeError(
            "ONNX output requires --onnx-input-shape, e.g. --onnx-input-shape 1,128"
        )
    try:
        import torch.nn as nn
    except ImportError as exc:
        raise RuntimeError("ONNX export requires torch installed.") from exc
    if not isinstance(model, nn.Module):
        raise RuntimeError(
            "ONNX export requires a torch-compatible model. "
            "Use output-type osmoimg/pytorch or provide a torch-based mode input."
        )

    try:
        from model_osmosis.adapters.onnx import ONNXAdapter
    except ImportError as exc:
        raise RuntimeError(
            "ONNX export requires onnxruntime. Install with: pip install \"osmosis[onnx]\""
        ) from exc

    output_path.parent.mkdir(parents=True, exist_ok=True)
    adapter = ONNXAdapter(verbose=False)
    adapter.from_pytorch(
        model,
        input_shape=onnx_input_shape,
        output_path=output_path,
        name=output_path.stem,
    )


def _to_osmo_model(
    model: Any,
    *,
    name: Optional[str],
    allow_lossy_cast: bool = False,
) -> Any:
    from model_osmosis.format.osmo_format import OsmoModel as _OsmoModel
    from model_osmosis.format.converters import UniversalConverter

    if isinstance(model, _OsmoModel):
        return model
    converter = UniversalConverter()
    return converter.to_osmo(model, name=name or "", allow_lossy_cast=allow_lossy_cast)


def _flatten_nested_mapping(data: Dict[str, Any], prefix: str = "") -> Dict[str, Any]:
    flat: Dict[str, Any] = {}
    for key in sorted(data.keys(), key=lambda item: str(item)):
        value = data[key]
        key_s = str(key)
        full = f"{prefix}/{key_s}" if prefix else key_s
        if isinstance(value, dict):
            flat.update(_flatten_nested_mapping(value, full))
        else:
            flat[full] = value
    return flat


def _save_named_arrays_npz(arrays: Dict[str, Any], output_path: Path) -> None:
    import json
    import numpy as np

    output_path.parent.mkdir(parents=True, exist_ok=True)
    safe_arrays: Dict[str, Any] = {}
    name_map: Dict[str, str] = {}

    for idx, (name, value) in enumerate(arrays.items()):
        safe = f"arr_{idx:08d}"
        safe_arrays[safe] = np.asarray(value)
        name_map[safe] = name

    np.savez(output_path, **safe_arrays)
    names_path = output_path.with_name(f"{output_path.stem}.names.json")
    with open(names_path, "w", encoding="utf-8") as fh:
        json.dump(name_map, fh, indent=2, ensure_ascii=False)


def _save_osmo_output(
    model: Any,
    output_path: Path,
    *,
    name: Optional[str],
    allow_lossy_cast: bool = False,
) -> None:
    from model_osmosis.format.serialization import save_osmo

    osmo = _to_osmo_model(
        model,
        name=name or output_path.stem,
        allow_lossy_cast=allow_lossy_cast,
    )
    save_osmo(osmo, output_path, compress=True, obfuscate=False)


def _save_safetensors_output(
    model: Any,
    output_path: Path,
    *,
    name: Optional[str],
    allow_lossy_cast: bool = False,
) -> None:
    try:
        from safetensors.numpy import save_file as save_safetensors_file
    except ImportError as exc:
        raise RuntimeError(
            "safetensors export requires safetensors. Install with: pip install safetensors"
        ) from exc

    osmo = _to_osmo_model(
        model,
        name=name or output_path.stem,
        allow_lossy_cast=allow_lossy_cast,
    )
    arrays = {k: v for k, v in osmo.parameters.items()}
    output_path.parent.mkdir(parents=True, exist_ok=True)
    save_safetensors_file(arrays, str(output_path))


def _save_tensorflow_output(
    model: Any,
    output_path: Path,
    *,
    name: Optional[str],
    strict_load: bool = True,
    allow_lossy_cast: bool = False,
) -> Dict[str, Any]:
    from model_osmosis.format.converters import UniversalConverter

    # If user already has a Keras model instance, save it directly.
    if hasattr(model, "save") and output_path.suffix in {".keras", ".h5"}:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        model.save(str(output_path))
        return {
            "missing_keys": [],
            "unexpected_keys": [],
            "lossy_casts": [],
            "dtype_transports": [],
        }

    osmo = _to_osmo_model(
        model,
        name=name or output_path.stem,
        allow_lossy_cast=allow_lossy_cast,
    )
    converter = UniversalConverter()
    tf_weights = converter.from_osmo(
        osmo,
        target_framework="tensorflow",
        strict_load=strict_load,
        allow_lossy_cast=allow_lossy_cast,
    )
    if not isinstance(tf_weights, dict):
        raise RuntimeError(
            "TensorFlow export produced unsupported object. "
            "Use output-type osmoimg/osmo/pytorch/hf/onnx/safetensors for this model."
        )
    flat = _flatten_nested_mapping(tf_weights)
    _save_named_arrays_npz(flat, output_path)
    return converter.last_report


def _save_jax_output(
    model: Any,
    output_path: Path,
    *,
    name: Optional[str],
    strict_load: bool = True,
    allow_lossy_cast: bool = False,
) -> Dict[str, Any]:
    from model_osmosis.format.converters import UniversalConverter

    osmo = _to_osmo_model(
        model,
        name=name or output_path.stem,
        allow_lossy_cast=allow_lossy_cast,
    )
    converter = UniversalConverter()
    jax_params = converter.from_osmo(
        osmo,
        target_framework="jax",
        strict_load=strict_load,
        allow_lossy_cast=allow_lossy_cast,
    )
    if isinstance(jax_params, dict):
        flat = _flatten_nested_mapping(jax_params)
    else:
        flat = {"params": jax_params}
    _save_named_arrays_npz(flat, output_path)
    return converter.last_report


def _save_torchscript_output(
    model: Any,
    output_path: Path,
    *,
    trace_input_shape: Optional[Tuple[int, ...]],
) -> None:
    try:
        import torch
        import torch.nn as nn
    except ImportError as exc:
        raise RuntimeError("TorchScript export requires torch installed.") from exc

    if not isinstance(model, nn.Module):
        raise RuntimeError(
            "TorchScript export requires a torch-compatible model. "
            "Use output-type osmoimg/osmo/pytorch/hf/onnx/tensorflow/jax/safetensors for non-torch models."
        )

    module = model.eval()
    scripted = None
    try:
        scripted = torch.jit.script(module)
    except Exception as script_exc:
        if trace_input_shape is None:
            raise RuntimeError(
                "TorchScript export failed with scripting. "
                "Provide --trace-input-shape (for trace fallback), e.g. 1,128."
            ) from script_exc
        dummy = torch.randn(*trace_input_shape)
        scripted = torch.jit.trace(module, dummy, strict=False)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    scripted.save(str(output_path))


def _validate_pipeline_args(args) -> None:
    if args.mode not in PIPELINE_MODES:
        raise ValueError(f"Unsupported mode: {args.mode}")
    if args.output_type not in PIPELINE_OUTPUT_TYPES:
        raise ValueError(f"Unsupported output type: {args.output_type}")

    if args.mode in {"merge", "stack", "weight-merge"} and not args.models:
        raise ValueError(f"`--models` is required for mode={args.mode}")
    if args.mode == "absorb" and (not args.source or not args.target):
        raise ValueError("`--source` and `--target` are required for mode=absorb")
    if args.mode == "weight-merge":
        if len(args.models or []) < 2:
            raise ValueError("`--mode weight-merge` requires at least two values in `--models`")
        floor = float(args.coverage_floor)
        if floor < 0.0 or floor > 1.0:
            raise ValueError("`--coverage-floor` must be in range [0.0, 1.0]")
    if int(args.top_k) <= 0:
        raise ValueError("`--top-k` must be a positive integer")


def _run_unified_pipeline(args) -> Dict[str, Any]:
    from model_osmosis.utils.determinism import set_global_seed

    _validate_pipeline_args(args)
    signer_backend = getattr(args, "signer_backend", "local-keyring")
    self_check_level = _resolve_pipeline_self_check_level(args)
    dry_run = bool(getattr(args, "dry_run", False))
    set_global_seed(args.seed)
    sample_data = _load_sample_data(args.sample_data)
    output_path = Path(args.output) if args.output else _default_pipeline_output(
        args.mode, args.output_type, args.name
    )
    output_path = output_path.resolve()

    refs_for_summary: List[str] = []
    loaded_refs: List[LoadedModelRef] = []

    if args.mode == "merge":
        loaded_refs = [_load_model_ref(item) for item in args.models]
        refs_for_summary.extend([item.raw for item in loaded_refs])
    elif args.mode == "weight-merge":
        loaded_refs = [_load_model_ref(item) for item in args.models]
        refs_for_summary.extend([item.raw for item in loaded_refs])
    elif args.mode == "absorb":
        source_ref = _load_model_ref(args.source)
        target_ref = _load_model_ref(args.target)
        loaded_refs = [source_ref, target_ref]
        refs_for_summary.extend([source_ref.raw, target_ref.raw])
    elif args.mode == "stack":
        loaded_refs = [_load_model_ref(item) for item in args.models]
        refs_for_summary.extend([item.raw for item in loaded_refs])
    else:
        raise ValueError(f"Unsupported mode: {args.mode}")

    preflight = _pipeline_preflight(args, loaded_refs=loaded_refs)

    details: Dict[str, Any] = {
        "output_path": output_path,
        "route": preflight["route"],
        "route_reason": preflight["route_reason"],
        "resolved_inputs": preflight["resolved_inputs"],
        "compression_applied": False,
        "payload_codec": "n/a",
        "compression_profile": "n/a",
        "compression_engine_mode": _resolve_compression_engine_mode(
            _normalize_mode_arg(
                getattr(args, "compression_engine", "auto"),
                COMPRESSION_ENGINE_MODES,
                "auto",
            )
        ),
        "trust_status": "unsigned",
        "signer_id": None,
        "key_id": None,
        "signature_alg": None,
        "signer_backend": "local-keyring",
        "transparency_status": None,
        "artifact_id": "",
        "manifest_sha256": "",
        "tensor_index_sha256": "",
        "protected": False,
        "weight_merge_report": None,
        "strategy_engine_mode": _resolve_strategy_engine_mode(getattr(args, "strategy_engine", "auto")),
        "cross_opt_stage": str(getattr(args, "cross_opt_stage", "auto")),
        "strategy_name": "",
        "strategy_kwargs": {},
        "recipe_id": "",
        "recipe_pack_version": "",
        "selector_version": "",
        "optimizer_mode": _resolve_strategy_engine_mode(getattr(args, "strategy_engine", "auto")),
        "profile_id": "",
        "pack_version": "",
        "canary_status": "shadow",
        "rollback_reason": "",
        "selection_latency_ms": 0.0,
        "fallback_reason": "",
        "conversion_report": {
            "missing_keys": [],
            "unexpected_keys": [],
            "lossy_casts": [],
            "dtype_transports": [],
        },
        "reliability_mode": "strict",
        "self_check_status": "skipped",
        "self_check_level": self_check_level,
        "self_check_checks": [],
        "self_check_skips": [],
        "dry_run": dry_run,
        "status": "PASS",
    }

    if args.allow_partial_load or args.allow_lossy_cast:
        details["reliability_mode"] = "degraded"

    if dry_run:
        details["self_check_status"] = "skipped_dry_run"
        details["self_check_level"] = "off"
        details["sources"] = refs_for_summary
        return details

    from model_osmosis import OsmosisPipeline, OsmosisConfig
    from model_osmosis.obfuscation import ObfuscationConfig
    from model_osmosis.format.image import OsmoImage

    config = OsmosisConfig(verbose=True, seed=args.seed)
    pipeline = OsmosisPipeline(config)

    result = None
    tokenizer = None
    parity_reference_model: Any = None

    if args.mode == "merge":
        loaded_models = [item.model for item in loaded_refs]
        source_refs: List[Any] = []
        for item in loaded_refs:
            if item.scheme == "hf":
                source_refs.append(item.value)
            else:
                source_refs.append(item.model)
        merge_result = _run_merge_backend(
            pipeline,
            loaded_models=loaded_models,
            strategy=args.merge_strategy,
            weights=args.weights,
            sample_data=sample_data,
            source_refs=source_refs,
            force_cross=details["route"] == "cross-merge",
            epochs=args.epochs,
            seed=args.seed,
            strategy_engine_mode=getattr(args, "strategy_engine", "auto"),
            cross_opt_stage=getattr(args, "cross_opt_stage", "auto"),
            return_metadata=True,
        )
        if isinstance(merge_result, tuple) and len(merge_result) == 2 and isinstance(merge_result[1], dict):
            result = merge_result[0]
            details.update(
                {
                    "strategy_engine_mode": str(merge_result[1].get("strategy_engine_mode", details["strategy_engine_mode"])),
                    "cross_opt_stage": str(merge_result[1].get("cross_opt_stage", details["cross_opt_stage"])),
                    "strategy_name": str(merge_result[1].get("strategy_name", "")),
                    "strategy_kwargs": dict(merge_result[1].get("strategy_kwargs", {}) or {}),
                    "recipe_id": str(merge_result[1].get("recipe_id", "")),
                    "recipe_pack_version": str(merge_result[1].get("recipe_pack_version", "")),
                    "selector_version": str(merge_result[1].get("selector_version", "")),
                    "optimizer_mode": str(merge_result[1].get("optimizer_mode", merge_result[1].get("strategy_engine_mode", details["optimizer_mode"]))),
                    "profile_id": str(merge_result[1].get("profile_id", merge_result[1].get("recipe_id", ""))),
                    "pack_version": str(merge_result[1].get("pack_version", merge_result[1].get("recipe_pack_version", ""))),
                    "canary_status": str(merge_result[1].get("canary_status", details["canary_status"])),
                    "rollback_reason": str(merge_result[1].get("rollback_reason", details["rollback_reason"])),
                    "selection_latency_ms": float(merge_result[1].get("selection_latency_ms", 0.0) or 0.0),
                    "fallback_reason": str(merge_result[1].get("fallback_reason", "")),
                }
            )
        else:
            result = merge_result
    elif args.mode == "weight-merge":
        weight_merge = _run_weight_merge_backend(
            loaded_refs=loaded_refs,
            strategy=args.merge_strategy,
            weights=args.weights,
            coverage_floor=float(args.coverage_floor),
            output_type=args.output_type,
            seed=args.seed,
            allow_lossy_cast=args.allow_lossy_cast,
            strategy_engine_mode=getattr(args, "strategy_engine", "auto"),
        )
        result = weight_merge["merged_osmo"]
        details["weight_merge_report"] = weight_merge.get("report", {})
        details.update(
            {
                "strategy_engine_mode": str(details["weight_merge_report"].get("strategy_engine_mode", details["strategy_engine_mode"])),
                "strategy_name": str(details["weight_merge_report"].get("strategy_used", "")),
                "strategy_kwargs": dict(details["weight_merge_report"].get("strategy_kwargs", {}) or {}),
                "recipe_id": str(details["weight_merge_report"].get("recipe_id", "")),
                "recipe_pack_version": str(details["weight_merge_report"].get("recipe_pack_version", "")),
                "selector_version": str(details["weight_merge_report"].get("selector_version", "")),
                "optimizer_mode": str(details["weight_merge_report"].get("optimizer_mode", details["weight_merge_report"].get("strategy_engine_mode", details["optimizer_mode"]))),
                "profile_id": str(details["weight_merge_report"].get("profile_id", details["weight_merge_report"].get("recipe_id", ""))),
                "pack_version": str(details["weight_merge_report"].get("pack_version", details["weight_merge_report"].get("recipe_pack_version", ""))),
                "canary_status": str(details["weight_merge_report"].get("canary_status", details["canary_status"])),
                "rollback_reason": str(details["weight_merge_report"].get("rollback_reason", details["rollback_reason"])),
                "selection_latency_ms": float(details["weight_merge_report"].get("selection_latency_ms", 0.0) or 0.0),
                "fallback_reason": str(details["weight_merge_report"].get("fallback_reason", "")),
            }
        )
    elif args.mode == "absorb":
        source_ref = loaded_refs[0]
        target_ref = loaded_refs[1]
        target_model = _ensure_absorb_target_supported(target_ref.model, target_ref.raw)
        result = pipeline.absorb(
            source=source_ref.model,
            target=target_model,
            sample_data=sample_data,
            auto_grow=args.auto_grow,
        )
    elif args.mode == "stack":
        models = [_ensure_stack_expert(item.model, item.raw) for item in loaded_refs]
        result = pipeline.stack(
            models=models,
            sample_data=sample_data,
            train_gate=args.train_gate,
            top_k=args.top_k,
        )
    else:
        raise ValueError(f"Unsupported mode: {args.mode}")

    if isinstance(result, tuple):
        model_result, tokenizer = result
    else:
        model_result = result

    if args.output_type == "osmoimg":
        output_path.parent.mkdir(parents=True, exist_ok=True)
        image = OsmoImage.build(
            model=model_result,
            name=args.name or output_path.stem,
            compression_profile=args.compression_profile,
            compression_engine=_normalize_mode_arg(
                getattr(args, "compression_engine", "auto"),
                COMPRESSION_ENGINE_MODES,
                "auto",
            ),
            trust_sign=args.trust_sign,
            provenance_mode=args.provenance,
            repro_mode=args.repro_mode,
            signer_id=args.signer_id,
            signer_backend=signer_backend,
        )
        if args.mode in {"merge", "weight-merge"}:
            provenance = dict(getattr(image.runtime, "provenance", {}) or {})
            provenance["merge_recipe"] = {
                "strategy_engine_mode": str(details.get("strategy_engine_mode", "")),
                "strategy_name": str(details.get("strategy_name", "")),
                "strategy_kwargs": dict(details.get("strategy_kwargs", {}) or {}),
                "recipe_id": str(details.get("recipe_id", "")),
                "recipe_pack_version": str(details.get("recipe_pack_version", "")),
                "selector_version": str(details.get("selector_version", "")),
                "selection_latency_ms": float(details.get("selection_latency_ms", 0.0) or 0.0),
                "fallback_reason": str(details.get("fallback_reason", "")),
            }
            image.runtime.provenance = provenance
        saved = image.save(
            output_path,
            compression_profile=args.compression_profile,
            compression_engine=_normalize_mode_arg(
                getattr(args, "compression_engine", "auto"),
                COMPRESSION_ENGINE_MODES,
                "auto",
            ),
            trust_sign=args.trust_sign,
            provenance_mode=args.provenance,
            repro_mode=args.repro_mode,
            signer_id=args.signer_id,
            signer_backend=signer_backend,
        )
        runtime_compression_stats = getattr(image.runtime, "compression_stats", {})
        compression_optimizer = (
            runtime_compression_stats.get("optimizer", {})
            if isinstance(runtime_compression_stats, dict)
            else {}
        )
        details.update(
            {
                "output_path": Path(saved).resolve(),
                "compression_applied": True,
                "payload_codec": image.runtime.payload_codec,
                "compression_profile": image.runtime.compression_profile,
                "compression_engine_mode": _resolve_compression_engine_mode(
                    _normalize_mode_arg(
                        getattr(args, "compression_engine", "auto"),
                        COMPRESSION_ENGINE_MODES,
                        "auto",
                    )
                ),
                "trust_status": getattr(image.runtime, "trust_status", "unsigned"),
                "signer_id": getattr(image.runtime, "signer_id", None),
                "key_id": getattr(image.runtime, "key_id", None),
                "signature_alg": getattr(image.runtime, "signature_alg", None),
                "signer_backend": getattr(image.runtime, "signer_backend", "local-keyring"),
                "transparency_status": getattr(image.runtime, "transparency_status", None),
                "artifact_id": getattr(image.runtime, "artifact_id", ""),
                "manifest_sha256": getattr(image.runtime, "manifest_sha256", ""),
                "tensor_index_sha256": getattr(image.runtime, "tensor_index_sha256", ""),
                "protected": bool(image.is_encrypted),
                "optimizer_mode": str(compression_optimizer.get("optimizer_mode", details.get("optimizer_mode", "auto"))),
                "profile_id": str(compression_optimizer.get("profile_id", details.get("profile_id", ""))),
                "pack_version": str(compression_optimizer.get("pack_version", details.get("pack_version", ""))),
                "selector_version": str(compression_optimizer.get("selector_version", details.get("selector_version", ""))),
                "canary_status": str(compression_optimizer.get("canary_status", details.get("canary_status", "shadow"))),
                "rollback_reason": str(compression_optimizer.get("rollback_reason", details.get("rollback_reason", ""))),
                "selection_latency_ms": float(compression_optimizer.get("selection_latency_ms", details.get("selection_latency_ms", 0.0)) or 0.0),
                "fallback_reason": str(compression_optimizer.get("fallback_reason", details.get("fallback_reason", ""))),
            }
        )
        parity_reference_model = image
    else:
        if args.mode == "weight-merge":
            export_model = model_result
            parity_reference_model = model_result
            details["protected"] = False
        else:
            obf_cfg = ObfuscationConfig(preserve_quality=True)
            export_model = pipeline.obfuscate(
                model_result,
                obfuscation_config=obf_cfg,
            )
            parity_reference_model = export_model
            details["protected"] = True

        if args.output_type == "pytorch":
            output_path.parent.mkdir(parents=True, exist_ok=True)
            if args.mode == "weight-merge":
                from model_osmosis.format.converters import UniversalConverter

                converter = UniversalConverter()
                state_dict = converter.from_osmo(
                    _to_osmo_model(
                        export_model,
                        name=args.name or output_path.stem,
                        allow_lossy_cast=args.allow_lossy_cast,
                    ),
                    target_framework="pytorch",
                    strict_load=not args.allow_partial_load,
                    allow_lossy_cast=args.allow_lossy_cast,
                )
                details["conversion_report"] = dict(converter.last_report or {})
                details["conversion_report"].setdefault("missing_keys", [])
                details["conversion_report"].setdefault("unexpected_keys", [])
                details["conversion_report"].setdefault("lossy_casts", [])
                details["conversion_report"].setdefault("dtype_transports", [])
                _save_model(state_dict, str(output_path))
            else:
                _save_model(export_model, str(output_path))
        elif args.output_type == "osmo":
            _save_osmo_output(
                export_model,
                output_path,
                name=args.name,
                allow_lossy_cast=args.allow_lossy_cast,
            )
        elif args.output_type == "hf":
            details["conversion_report"] = _save_hf_output(
                export_model,
                output_path,
                hf_base_model=args.hf_base_model,
                name=args.name,
                tokenizer=tokenizer,
                strict_load=not args.allow_partial_load,
                allow_lossy_cast=args.allow_lossy_cast,
            )
        elif args.output_type == "onnx":
            shape = _parse_onnx_input_shape(args.onnx_input_shape)
            if args.mode == "weight-merge":
                torch_model, report = _reconstruct_torch_module_for_weight_merge(
                    merged_osmo=_to_osmo_model(
                        export_model,
                        name=args.name or output_path.stem,
                        allow_lossy_cast=args.allow_lossy_cast,
                    ),
                    loaded_refs=loaded_refs,
                    strict_load=not args.allow_partial_load,
                    allow_lossy_cast=args.allow_lossy_cast,
                )
                details["conversion_report"] = report
                _save_onnx_output(torch_model, output_path, shape)
            else:
                _save_onnx_output(export_model, output_path, shape)
        elif args.output_type in {"tensorflow", "keras"}:
            details["conversion_report"] = _save_tensorflow_output(
                export_model,
                output_path,
                name=args.name,
                strict_load=not args.allow_partial_load,
                allow_lossy_cast=args.allow_lossy_cast,
            )
        elif args.output_type in {"jax", "flax"}:
            details["conversion_report"] = _save_jax_output(
                export_model,
                output_path,
                name=args.name,
                strict_load=not args.allow_partial_load,
                allow_lossy_cast=args.allow_lossy_cast,
            )
        elif args.output_type == "safetensors":
            _save_safetensors_output(
                export_model,
                output_path,
                name=args.name,
                allow_lossy_cast=args.allow_lossy_cast,
            )
        elif args.output_type == "torchscript":
            trace_shape = _parse_onnx_input_shape(
                args.trace_input_shape or args.onnx_input_shape
            )
            if args.mode == "weight-merge":
                torch_model, report = _reconstruct_torch_module_for_weight_merge(
                    merged_osmo=_to_osmo_model(
                        export_model,
                        name=args.name or output_path.stem,
                        allow_lossy_cast=args.allow_lossy_cast,
                    ),
                    loaded_refs=loaded_refs,
                    strict_load=not args.allow_partial_load,
                    allow_lossy_cast=args.allow_lossy_cast,
                )
                details["conversion_report"] = report
                _save_torchscript_output(
                    torch_model,
                    output_path,
                    trace_input_shape=trace_shape,
                )
            else:
                _save_torchscript_output(
                    export_model,
                    output_path,
                    trace_input_shape=trace_shape,
                )
        else:
            raise ValueError(f"Unsupported output type: {args.output_type}")

    if self_check_level != "off":
        self_check_result = _pipeline_self_check(
            output_type=args.output_type,
            output_path=Path(details["output_path"]),
            self_check_level=self_check_level,
            reference_model=parity_reference_model,
        )
        details.update(self_check_result)
    else:
        details["self_check_status"] = "skipped_disabled"
        details["self_check_level"] = "off"

    details["sources"] = refs_for_summary
    return details


def cmd_pipeline(args):
    """Unified pipeline command (merge/absorb/stack + protected output)."""
    print(_logo())
    print(f"\n  {Colors.BOLD}PIPELINE{Colors.RESET}\n")

    signer_backend = getattr(args, "signer_backend", "local-keyring")
    start = time.perf_counter()
    _print_info("Mode:", args.mode)
    _print_info("Output type:", args.output_type)
    _print_info("Compression profile:", args.compression_profile)
    _print_info(
        "Compression engine:",
        _resolve_compression_engine_mode(
            _normalize_mode_arg(
                getattr(args, "compression_engine", "auto"),
                COMPRESSION_ENGINE_MODES,
                "auto",
            )
        ),
    )
    _print_info("Trust sign:", "yes" if args.trust_sign else "no")
    _print_info("Signer backend:", signer_backend)
    _print_info("Provenance mode:", args.provenance)
    _print_info("Repro mode:", args.repro_mode)
    _print_info("Seed:", str(args.seed))
    _print_info("Strict load:", "no" if args.allow_partial_load else "yes")
    _print_info("Allow lossy cast:", "yes" if args.allow_lossy_cast else "no")
    if args.mode == "weight-merge":
        _print_info("Coverage floor:", f"{float(args.coverage_floor):.2f}")
    _print_info("Dry run:", "yes" if args.dry_run else "no")
    _print_info("Self-check:", "yes" if args.self_check else "no")
    _print_info("Self-check level:", _resolve_pipeline_self_check_level(args))
    _print_info("Report file:", args.report_file or "none")
    print()
    if args.output_type == "osmoimg" and not args.trust_sign:
        _warn_insecure_output("pipeline")
        print()

    details: Dict[str, Any]
    status_label = f"{Colors.GREEN}PASS{Colors.RESET}"
    failure_reason = ""
    try:
        details = _run_unified_pipeline(args)
    except Exception as exc:
        details = {
            "output_path": Path(args.output) if args.output else _default_pipeline_output(
                args.mode, args.output_type, args.name
            ),
            "compression_applied": args.output_type == "osmoimg",
            "payload_codec": "n/a",
            "compression_profile": args.compression_profile if args.output_type == "osmoimg" else "n/a",
            "compression_engine_mode": _resolve_compression_engine_mode(
                _normalize_mode_arg(
                    getattr(args, "compression_engine", "auto"),
                    COMPRESSION_ENGINE_MODES,
                    "auto",
                )
            ),
            "trust_status": "unknown",
            "signer_id": None,
            "key_id": None,
            "signature_alg": None,
            "signer_backend": signer_backend,
            "artifact_id": "",
            "manifest_sha256": "",
            "tensor_index_sha256": "",
            "transparency_status": None,
            "reliability_mode": "degraded" if (args.allow_partial_load or args.allow_lossy_cast) else "strict",
            "strategy_engine_mode": _resolve_strategy_engine_mode(getattr(args, "strategy_engine", "auto")),
            "cross_opt_stage": str(getattr(args, "cross_opt_stage", "auto")),
            "strategy_name": "",
            "strategy_kwargs": {},
            "recipe_id": "",
            "recipe_pack_version": "",
            "selector_version": "",
            "optimizer_mode": _resolve_strategy_engine_mode(getattr(args, "strategy_engine", "auto")),
            "profile_id": "",
            "pack_version": "",
            "canary_status": "shadow",
            "rollback_reason": "",
            "selection_latency_ms": 0.0,
            "fallback_reason": "",
            "conversion_report": {},
            "protected": False,
            "route": "unknown",
            "route_reason": "error_before_route",
            "self_check_status": "failed",
            "self_check_level": _resolve_pipeline_self_check_level(args),
            "self_check_checks": [],
            "self_check_skips": [],
            "dry_run": bool(args.dry_run),
            "status": "FAIL",
            "resolved_inputs": [],
        }
        status_label = f"{Colors.RED}FAIL{Colors.RESET}"
        failure_reason = str(exc)
        elapsed = time.perf_counter() - start
        report_payload = _pipeline_report_payload(args, details=details, elapsed_seconds=elapsed)
        report_payload["status"] = "FAIL"
        report_payload["error"] = failure_reason
        if args.report_file:
            _write_pipeline_report(args.report_file, report_payload)
        print()
        print(f"  {Colors.BOLD}Pipeline Summary{Colors.RESET}")
        _print_info("Status:", status_label)
        _print_info("Mode:", args.mode)
        _print_info("Output type:", args.output_type)
        _print_info("Output path:", str(details["output_path"]))
        _print_info("Route:", str(details.get("route", "unknown")))
        _print_info("Route reason:", str(details.get("route_reason", "unknown")))
        _print_info("Self-check:", str(details.get("self_check_status", "failed")))
        _print_info("Self-check level:", str(details.get("self_check_level", "off")))
        _print_info("Error:", failure_reason)
        _print_info("Elapsed:", f"{elapsed:.1f}s")
        print()
        raise

    elapsed = time.perf_counter() - start
    report_payload = _pipeline_report_payload(args, details=details, elapsed_seconds=elapsed)
    report_payload["status"] = "PASS"
    if args.report_file:
        _write_pipeline_report(args.report_file, report_payload)

    print()
    print(f"  {Colors.BOLD}Pipeline Summary{Colors.RESET}")
    _print_info("Status:", status_label)
    _print_info("Mode:", args.mode)
    _print_info("Output type:", args.output_type)
    _print_info("Output path:", str(details["output_path"]))
    _print_info("Route:", str(details.get("route", "unknown")))
    _print_info("Route reason:", str(details.get("route_reason", "unknown")))
    _print_info("Protected:", "yes" if details.get("protected") else "no")
    _print_info(
        "Compression:",
        "osmoimg container enabled" if details.get("compression_applied")
        else "compression applies to osmoimg only",
    )
    _print_info("Payload codec:", str(details.get("payload_codec", "n/a")))
    _print_info("Compression profile:", str(details.get("compression_profile", "n/a")))
    _print_info("Optimizer mode:", str(details.get("optimizer_mode", "auto")))
    _print_info("Optimizer profile:", str(details.get("profile_id", "n/a")))
    _print_info("Optimizer pack:", str(details.get("pack_version", "n/a")))
    _print_info("Optimizer selector:", str(details.get("selector_version", "n/a")))
    _print_info("Canary status:", str(details.get("canary_status", "shadow")))
    _print_info("Rollback reason:", str(details.get("rollback_reason", "")) or "none")
    _print_info("Selection latency:", f"{float(details.get('selection_latency_ms', 0.0) or 0.0):.4f} ms")
    _print_info("Trust status:", str(details.get("trust_status", "unsigned")))
    _print_info("Signer ID:", str(details.get("signer_id") or "none"))
    _print_info("Key ID:", str(details.get("key_id") or "none"))
    _print_info("Signature alg:", str(details.get("signature_alg") or "none"))
    _print_info("Signer backend:", str(details.get("signer_backend") or "local-keyring"))
    _print_info("Artifact ID:", str(details.get("artifact_id") or "n/a"))
    _print_info("Manifest SHA256:", str(details.get("manifest_sha256") or "n/a"))
    _print_info("Tensor index SHA256:", str(details.get("tensor_index_sha256") or "n/a"))
    _print_info("Transparency:", str(details.get("transparency_status") or "n/a"))
    _print_info("Reliability mode:", str(details.get("reliability_mode", "strict")))
    conv_report = details.get("conversion_report", {}) or {}
    _print_info("Missing keys:", str(len(conv_report.get("missing_keys", []))))
    _print_info("Unexpected keys:", str(len(conv_report.get("unexpected_keys", []))))
    _print_info("Lossy casts:", str(len(conv_report.get("lossy_casts", []))))
    _print_info("Dtype transports:", str(len(conv_report.get("dtype_transports", []))))
    if args.mode == "weight-merge":
        wm = details.get("weight_merge_report", {}) or {}
        _print_info("Compatible tensors:", str(int(wm.get("compatible_tensor_count", 0) or 0)))
        _print_info("Skipped tensors:", str(int(wm.get("skipped_tensor_count", 0) or 0)))
        _print_info("Merged bytes:", str(int(wm.get("merged_bytes", 0) or 0)))
        _print_info("Min coverage:", f"{float(wm.get('min_coverage_ratio', 0.0) or 0.0):.4f}")
        _print_info("Coverage floor:", f"{float(wm.get('coverage_floor', 0.0) or 0.0):.4f}")
        skip_counts = wm.get("skip_reason_counts", {}) or {}
        _print_info(
            "Skip reasons:",
            (
                f"missing={int(skip_counts.get('missing', 0) or 0)}, "
                f"shape_mismatch={int(skip_counts.get('shape_mismatch', 0) or 0)}, "
                f"dtype_mismatch={int(skip_counts.get('dtype_mismatch', 0) or 0)}"
            ),
        )
    _print_info(
        "Fallbacks:",
        _reliability_fallbacks(args.allow_partial_load, args.allow_lossy_cast),
    )
    _print_info("Cross opt stage:", str(details.get("cross_opt_stage", "auto")))
    _print_info("Self-check:", str(details.get("self_check_status", "skipped")))
    _print_info("Self-check level:", str(details.get("self_check_level", "off")))
    skips = details.get("self_check_skips", []) or []
    _print_info("Self-check skips:", ", ".join(skips) if skips else "none")
    _print_info("Dry run:", "yes" if details.get("dry_run") else "no")
    if args.report_file:
        _print_info("Report file:", str(Path(args.report_file)))
    _print_info("Elapsed:", f"{elapsed:.1f}s")
    print()

def _load_model(path: str):
    """Load a model from path or HuggingFace ID."""
    path_obj = Path(path)

    # PyTorch file
    if path_obj.suffix in (".pt", ".pth", ".bin"):
        import torch
        return torch.load(path, map_location="cpu", weights_only=False)

    # SafeTensors directory or HuggingFace model
    if path_obj.is_dir() or "/" in path:
        try:
            from transformers import AutoModelForCausalLM
            _print_step("📥", f"Loading {path}...")
            return AutoModelForCausalLM.from_pretrained(
                path, trust_remote_code=True
            )
        except ImportError:
            pass  # transformers not installed, try other loaders
        except Exception:
            pass  # HF loading failed, fall through to other loaders

    # ONNX
    if path_obj.suffix == ".onnx":
        import onnxruntime
        return onnxruntime.InferenceSession(path)

    # Fallback: try torch
    import torch
    return torch.load(path, map_location="cpu", weights_only=False)


def _save_model(model, output: str):
    """Save model to the output path."""
    output_path = Path(output)

    if hasattr(model, "save_pretrained"):
        output_path.mkdir(parents=True, exist_ok=True)
        model.save_pretrained(output_path)
        _print_step("💾", f"Saved (HuggingFace format)")
    elif isinstance(model, dict):
        import torch
        output_path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(model, output_path)
        _print_step("💾", f"Saved (PyTorch state dict dict)")
    elif hasattr(model, "state_dict"):
        import torch
        output_path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(model.state_dict(), output_path)
        _print_step("💾", f"Saved (PyTorch state dict)")
    else:
        import pickle
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as f:
            pickle.dump(model, f)
        _print_step("💾", f"Saved (pickle)")


def _parse_weights(value: str) -> List[float]:
    """Parse weight string like '0.4,0.6' or '0.4 0.6'."""
    parts = value.replace(",", " ").split()
    return [float(x) for x in parts]


# ─── Unified 4-Command UX ────────────────────────────────────────────────

def cmd_make(args):
    """Unified model-operation entrypoint (merge/weight-merge/absorb/stack/cross-merge)."""
    print(_logo())
    print(f"\n  {Colors.BOLD}MAKE{Colors.RESET}\n")

    _print_info("Mode:", args.mode)
    _print_info("Output type:", args.output_type)
    _print_info("Merge strategy:", args.merge_strategy)
    _print_info("Coverage floor:", f"{float(args.coverage_floor):.2f}")
    _print_info("Seed:", str(args.seed))
    print()
    _print_step("↪", "Routing to unified pipeline backend...")

    from model_osmosis.services.unified_workflow import build_pipeline_args_from_make

    pipeline_args = build_pipeline_args_from_make(args)
    cmd_pipeline(pipeline_args)


def cmd_protect(args):
    """Package and protect a model into trusted .osmoimg output."""
    print(_logo())
    print(f"\n  {Colors.BOLD}PROTECT{Colors.RESET}\n")

    _print_info("Model:", args.model)
    _print_info("Output:", args.output)
    _print_info("Compression profile:", args.compression_profile)
    _print_info("Trust sign:", "yes" if args.trust_sign else "no")
    _print_info("Provenance:", args.provenance)
    _print_info("Repro mode:", args.repro_mode)
    _print_info("Signer backend:", args.signer_backend)
    print()
    _print_step("↪", "Routing to build backend...")

    from model_osmosis.services.unified_workflow import build_build_args_from_protect

    build_args = build_build_args_from_protect(args)
    cmd_build(build_args)


def cmd_check(args):
    """Unified validation gate (inspect + trust verify + self-check + optional benchmarks/gates)."""
    print(_logo())
    print(f"\n  {Colors.BOLD}CHECK{Colors.RESET}\n")

    from model_osmosis.services.unified_workflow import evaluate_check_workflow

    trust_profile = _resolve_trust_profile(getattr(args, "trust_profile", "standard"))
    require_trust = _effective_require_trust(bool(getattr(args, "require_trust", False)), trust_profile)

    _print_info("Image:", args.image)
    _print_info("Self-check level:", args.self_check_level)
    _print_info("Trust profile:", trust_profile)
    _print_info("Require trust:", "yes" if require_trust else "no")
    _print_info("Policy file:", args.policy_file or "none")
    _print_info("JSON:", "yes" if args.json else "no")
    print()

    payload = evaluate_check_workflow(
        image=str(args.image),
        self_check_level=str(args.self_check_level),
        trust_profile=trust_profile,
        require_trust=bool(require_trust),
        policy_file=args.policy_file,
        benchmark_runtime=bool(args.benchmark_runtime),
        benchmark_profile=str(args.benchmark_profile),
        runtime_profile=str(args.runtime_profile),
        device=str(args.device),
        runtime_quantization=str(args.runtime_quantization),
        run_release_check=bool(args.release_check),
        trust_checker=_check_trust_for_action,
        self_checker=_pipeline_self_check,
    )
    optional_checks = payload.get("optional_checks", {}) or {}
    failures = list(payload.get("failed_checks", []) or [])

    if args.json:
        print(json.dumps(payload, indent=2, sort_keys=True))
    else:
        _print_info("Status:", f"{Colors.GREEN}PASS{Colors.RESET}" if not failures else f"{Colors.RED}FAIL{Colors.RESET}")
        _print_info("Trust status:", payload["trust_status"])
        _print_info("Payload codec:", payload["payload_codec"])
        _print_info("Compression profile:", payload["compression_profile"])
        _print_info("Self-check:", payload["self_check_status"])
        _print_info("Self-check level:", payload["self_check_level"])
        _print_info(
            "Self-check skips:",
            ", ".join(payload["self_check_skips"]) if payload["self_check_skips"] else "none",
        )
        if optional_checks:
            _print_info("Optional checks:", ", ".join(sorted(optional_checks.keys())))
        if failures:
            _print_info("Failed checks:", ", ".join(sorted(set(failures))))
        print()

    if failures:
        raise RuntimeError("check failed: " + ", ".join(sorted(set(failures))))


def cmd_launch(args):
    """Unified runtime entrypoint (run/serve/export)."""
    from model_osmosis.services.unified_workflow import build_launch_action_namespace

    action, launch_args = build_launch_action_namespace(args)
    if action == "run":
        cmd_run(launch_args)
        return
    if action == "serve":
        cmd_serve(launch_args)
        return
    if action == "export":
        cmd_export(launch_args)
        return
    raise RuntimeError(f"Unsupported launch action '{action}'.")


def cmd_expand(args):
    """Expand and improve a model with Net2Net++ style local adaptation."""
    print(_logo())
    print(f"\n  {Colors.BOLD}EXPAND{Colors.RESET}\n")

    from model_osmosis import OsmosisConfig, OsmosisPipeline
    from model_osmosis.format.image import OsmoImage
    from model_osmosis.services.expand_data import infer_expand_task, plan_expand_data
    from model_osmosis.services.expand_revisions import (
        append_revision,
        default_revision_output_path,
        next_revision_id,
    )
    from model_osmosis.utils.determinism import set_global_seed

    loaded_ref = _load_expand_model_ref(args.model_ref)
    runtime_model, runtime_source, route = _prepare_expand_runtime_model(loaded_ref)
    model_name = _normalize_expand_name(args.name or _default_expand_name(loaded_ref))
    resolved_task = infer_expand_task(
        task=args.task,
        model_ref=loaded_ref.raw,
        model_obj=runtime_model,
    )
    feature_dim = _infer_expand_feature_dim(runtime_model)
    sample_data, data_report = plan_expand_data(
        task=resolved_task,
        dataset=args.dataset,
        seed=int(args.seed),
        feature_dim=int(feature_dim),
        model_ref=loaded_ref.raw,
        data_mode=str(args.data_mode),
        max_download_gb=float(args.max_download_gb),
        data_ttl_hours=float(args.data_ttl_hours),
        allow_licenses=list(getattr(args, "allow_licenses", []) or []),
        offline=bool(args.offline),
    )

    revision_id = next_revision_id(model_name)
    default_output = default_revision_output_path(model_name, revision_id)
    output_path = Path(args.output) if args.output else default_output
    if output_path.suffix != ".osmoimg":
        output_path = output_path.with_suffix(".osmoimg")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    _print_info("Model:", loaded_ref.raw)
    _print_info("Route:", route)
    _print_info("Runtime source:", runtime_source)
    _print_info("Strict mode:", "on")
    _print_info("Fallback mode:", str(args.fallback_mode))
    _print_info("Task:", resolved_task)
    _print_info("Data mode:", str(args.data_mode))
    _print_info("Dataset source:", str(data_report.get("dataset_source", "synthetic")))
    _print_info("Selection reason:", str(data_report.get("selection_reason", "")) or "n/a")
    _print_info("Planner route:", str(data_report.get("planner_route", "")) or "n/a")
    _print_info("Route reason:", str(data_report.get("route_reason", "")) or "n/a")
    _print_info("Task batch kind:", str(data_report.get("task_batch_kind", "")) or "n/a")
    _print_info("Task batch shape:", str(data_report.get("task_batch_shape", "")) or "n/a")
    _print_info("Dataset ID:", str(data_report.get("dataset_id", "")) or "n/a")
    _print_info("Source hub:", str(data_report.get("source_hub", "")) or "n/a")
    _print_info("Download MB:", f"{float(data_report.get('download_mb', 0.0) or 0.0):.4f}")
    _print_info("Target params:", str(args.target_params or "x1.2"))
    _print_info("Budget:", f"{int(args.budget_minutes)} min / {float(args.budget_ram_gb):.1f} GB")
    _print_info("Optimizer mode:", str(args.optimizer_mode))
    _print_info("Output:", str(output_path))
    _print_info("Revision ID:", revision_id)
    print()

    set_global_seed(int(args.seed))
    pipeline = OsmosisPipeline(OsmosisConfig(verbose=True, seed=int(args.seed)))
    expanded_model, expand_report = pipeline.expand(
        runtime_model,
        sample_data=sample_data,
        task=resolved_task,
        target_params=args.target_params,
        budget_minutes=int(args.budget_minutes),
        budget_ram_gb=float(args.budget_ram_gb),
        seed=int(args.seed),
        optimizer_mode=str(args.optimizer_mode),
        fallback_mode=str(args.fallback_mode),
        strict_mode=True,
        route=str(route),
        runtime_source=str(runtime_source),
    )

    if bool(expand_report.get("hard_failure", False)):
        reason = str(
            expand_report.get("error")
            or expand_report.get("can_expand_reason")
            or expand_report.get("rollback_reason")
            or "strict_expand_quality_gate_failed"
        )
        raise _error_with_next_command(
            f"Strict expand failed: {reason}",
            f"osmosis expand {loaded_ref.raw} --fallback-mode distill",
        )

    _print_step("📦", "Building trusted revision artifact...")
    image = OsmoImage.build(
        model=expanded_model,
        name=model_name,
        compression_profile="lossless-max",
        compression_engine="auto",
        trust_sign=True,
        provenance_mode="private",
        repro_mode="strict",
    )
    image.runtime.provenance = dict(image.runtime.provenance or {})
    image.runtime.provenance["expand"] = {
        "revision_id": revision_id,
        "model_ref": loaded_ref.raw,
        "route": str(route),
        "runtime_source": str(runtime_source),
        "strict_mode": True,
        "fallback_mode": str(args.fallback_mode),
        "degraded": bool(expand_report.get("degraded", False)),
        "task": resolved_task,
        "data_source": str(data_report.get("dataset_source", "synthetic")),
        "quality_gate": dict(expand_report.get("quality_gate", {}) or {}),
        "quality_delta": float(expand_report.get("quality_delta", 0.0) or 0.0),
        "optimizer_mode": str(expand_report.get("optimizer_mode", "auto")),
        "profile_id": str(expand_report.get("profile_id", "")),
        "pack_version": str(expand_report.get("pack_version", "")),
        "selector_version": str(expand_report.get("selector_version", "")),
        "canary_status": str(expand_report.get("canary_status", "shadow")),
        "rollback_reason": str(expand_report.get("rollback_reason", "")),
        "selection_latency_ms": float(expand_report.get("selection_latency_ms", 0.0) or 0.0),
        "self_check_status": str(expand_report.get("self_check_status", "n/a")),
    }
    saved_path = image.save(
        output_path,
        compression_profile="lossless-max",
        compression_engine="auto",
        trust_sign=True,
        provenance_mode="private",
        repro_mode="strict",
    )
    self_check_checks = ["artifact_saved", "trust_envelope_present", "route_recorded"]
    self_check_skips: List[str] = []
    self_check_status = "pass"
    if not Path(saved_path).exists():
        self_check_status = "fail"
    if str(getattr(image.runtime, "trust_status", "")) != "trusted":
        self_check_status = "fail"
    expand_report = dict(expand_report or {})
    expand_report["self_check_status"] = str(self_check_status)
    expand_report["self_check_level"] = "smoke"
    expand_report["self_check_checks"] = list(self_check_checks)
    expand_report["self_check_skips"] = list(self_check_skips)

    accepted = bool(expand_report.get("accepted", False))
    revision_entry = append_revision(
        name=model_name,
        artifact_path=str(saved_path),
        accepted=accepted,
        baseline_metric=float(expand_report.get("baseline_metric", 0.0) or 0.0),
        post_expand_metric=float(expand_report.get("post_expand_metric", 0.0) or 0.0),
        quality_delta=float(expand_report.get("quality_delta", 0.0) or 0.0),
        report={
            "expand": dict(expand_report),
            "data": dict(data_report),
            "model_ref": loaded_ref.raw,
            "route": str(route),
            "runtime_source": runtime_source,
        },
    )

    extra_export_path = ""
    if args.export_format != "osmoimg":
        export_fmt = str(args.export_format)
        _print_step("🔄", f"Exporting expanded model to {export_fmt}...")
        if output_path.suffix == ".osmoimg":
            base = output_path.with_suffix("")
        else:
            base = output_path
        suffix = _expand_export_suffix(export_fmt)
        if export_fmt == "hf":
            export_path = base.with_name(f"{base.name}{suffix}")
        else:
            export_path = Path(f"{str(base)}{suffix}")

        if export_fmt == "pytorch":
            _save_model(expanded_model, str(export_path))
        elif export_fmt == "hf":
            _save_hf_output(
                expanded_model,
                export_path,
                hf_base_model=(loaded_ref.value if loaded_ref.scheme == "hf" else None),
                name=model_name,
                tokenizer=None,
                strict_load=True,
                allow_lossy_cast=False,
            )
        elif export_fmt == "onnx":
            _save_onnx_output(expanded_model, export_path, (1, int(feature_dim)))
        elif export_fmt == "tensorflow":
            _save_tensorflow_output(
                expanded_model,
                export_path,
                name=model_name,
                strict_load=True,
                allow_lossy_cast=False,
            )
        elif export_fmt == "jax":
            _save_jax_output(
                expanded_model,
                export_path,
                name=model_name,
                strict_load=True,
                allow_lossy_cast=False,
            )
        else:
            raise RuntimeError(f"Unsupported expand export format: {export_fmt}")
        extra_export_path = str(export_path)

    status_color = Colors.GREEN if accepted else Colors.YELLOW
    status_text = "accepted" if accepted else "rejected"
    print()
    print(f"  {Colors.BOLD}Expand Summary{Colors.RESET}")
    _print_info("Status:", f"{status_color}{status_text}{Colors.RESET}")
    _print_info("Route:", str(expand_report.get("route", route)))
    _print_info("Runtime source:", str(expand_report.get("runtime_source", runtime_source)))
    _print_info("Strict mode:", "on")
    _print_info("Fallback mode:", str(expand_report.get("fallback_mode", args.fallback_mode)))
    _print_info("Degraded:", "yes" if bool(expand_report.get("degraded", False)) else "no")
    _print_info("Quality gate:", "pass" if bool((expand_report.get("quality_gate", {}) or {}).get("passed", False)) else "fail")
    _print_info("Revision:", str(revision_entry.get("revision_id", revision_id)))
    _print_info("Parent:", str(revision_entry.get("parent_revision_id") or "none"))
    _print_info("Artifact:", str(saved_path))
    _print_info("Baseline metric:", f"{float(expand_report.get('baseline_metric', 0.0) or 0.0):.6f}")
    _print_info("Post metric:", f"{float(expand_report.get('post_expand_metric', 0.0) or 0.0):.6f}")
    _print_info("Quality delta:", f"{float(expand_report.get('quality_delta', 0.0) or 0.0):.6f}")
    _print_info("Profile ID:", str(expand_report.get("profile_id", "")))
    _print_info("Pack version:", str(expand_report.get("pack_version", "")))
    _print_info("Selector:", str(expand_report.get("selector_version", "")))
    _print_info("Canary status:", str(expand_report.get("canary_status", "shadow")))
    _print_info("Rollback reason:", str(expand_report.get("rollback_reason", "")) or "none")
    _print_info("Fallback reason:", str(expand_report.get("fallback_reason", "")) or "none")
    _print_info("Self-check status:", str(expand_report.get("self_check_status", "n/a")))
    _print_info("Selection latency:", f"{float(expand_report.get('selection_latency_ms', 0.0) or 0.0):.4f} ms")
    _print_info("Trust status:", str(image.runtime.trust_status))
    if extra_export_path:
        _print_info("Extra export:", extra_export_path)
    print()


def cmd_expand_status(args):
    """Show local expansion revision chain and optimizer metadata."""
    print(_logo())
    print(f"\n  {Colors.BOLD}EXPAND STATUS{Colors.RESET}\n")

    from model_osmosis.optimizer.self_optimize import get_optimizer_status
    from model_osmosis.services.expand_revisions import active_revision, list_revisions, revision_state_path

    revisions = list_revisions(name=args.name)
    active = active_revision(name=args.name)
    optimizer = get_optimizer_status()
    expand_meta = dict((optimizer.get("subsystems", {}) or {}).get("expand", {}))

    payload = {
        "schema_version": 1,
        "state_path": str(revision_state_path()),
        "name_filter": str(args.name or ""),
        "active_revision": active,
        "revision_count": len(revisions),
        "revisions": revisions,
        "optimizer": expand_meta,
    }

    if args.json:
        print(json.dumps(payload, indent=2, sort_keys=True))
        return

    _print_info("State path:", payload["state_path"])
    _print_info("Name filter:", payload["name_filter"] or "none")
    _print_info("Revision count:", str(payload["revision_count"]))
    _print_info("Active revision:", str((active or {}).get("revision_id", "none")))
    _print_info("Canary status:", str(expand_meta.get("canary_status", "shadow")))
    _print_info("Promoted profile:", str(expand_meta.get("promoted_profile_id", "none") or "none"))
    _print_info("Rollback reason:", str(expand_meta.get("rollback_reason", "")) or "none")
    print()

    for item in revisions[-10:]:
        marker = "✓" if bool(item.get("accepted", False)) else "✗"
        print(
            f"  {marker} {item.get('revision_id')} "
            f"(parent={item.get('parent_revision_id') or 'none'}, "
            f"delta={float(item.get('quality_delta', 0.0) or 0.0):.6f})"
        )
    if not revisions:
        print("  (no revisions yet)")
    print()


def cmd_expand_rollback(args):
    """Set active expand revision pointer to a previous revision id."""
    print(_logo())
    print(f"\n  {Colors.BOLD}EXPAND ROLLBACK{Colors.RESET}\n")

    from model_osmosis.services.expand_revisions import rollback_revision

    entry = rollback_revision(args.to)
    _print_step("✓", "Active expand revision updated")
    _print_info("Revision:", str(entry.get("revision_id", "")))
    _print_info("Artifact:", str(entry.get("artifact_path", "")))
    _print_info("Accepted:", "yes" if bool(entry.get("accepted", False)) else "no")
    print()


# ─── Main Parser ──────────────────────────────────────────────────────────

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="osmosis",
        description="Model Osmosis — Docker for Machine Learning",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=f"""
{Colors.BOLD}Image Commands (Docker-style):{Colors.RESET}
  osmosis make --mode merge --models MODEL_A MODEL_B --output-type osmo
  osmosis protect MODEL -o model.osmoimg
  osmosis check model.osmoimg --self-check-level smoke
  osmosis launch model.osmoimg --action run --prompt "Hello world"
  osmosis artifact track model.osmoimg
  osmosis artifact verify --all
  osmosis demo
  osmosis doctor
  osmosis build user/Model -o model.osmoimg
  osmosis run model.osmoimg --prompt "Hello world"
  osmosis inspect model.osmoimg
  osmosis export model.osmoimg --format pytorch -o model.pt
  osmosis images ./models
  osmosis pull user/Model -o model.osmoimg
  osmosis push model.osmoimg user/repo --private
  osmosis register

{Colors.BOLD}Model Commands:{Colors.RESET}
  osmosis pipeline --mode merge --output-type osmoimg --models MODEL_A MODEL_B
  osmosis merge modelA.pt modelB.pt -s dare -o merged.pt
  osmosis obfuscate model.pt -o protected.pt
  osmosis compare original.pt obfuscated.pt
  osmosis benchmark modelA.pt modelB.pt modelC.pt
  osmosis benchmark-runtime model.osmoimg --profile i5_2gb --json
  osmosis info model.pt
  osmosis convert model.pt --to tensorflow -o model.tf
  osmosis diff a.osmoimg b.osmoimg
  osmosis validate model.osmoimg
  osmosis verify model.osmoimg --require-trust
  osmosis quantize model.osmoimg --bits 8
  osmosis transcode model.osmoimg --compression-profile lossless-max
  osmosis trust keygen --name default
  osmosis serve model.osmoimg --port 8080
  osmosis history model.osmoimg
  osmosis release-check --json
""",
    )

    sub = parser.add_subparsers(dest="command", help="Command to run")

    # ── make ──
    make = sub.add_parser(
        "make",
        help="Unified model operation command",
        description="Build/optimize model side with merge/weight-merge/absorb/stack/cross-merge",
    )
    make.add_argument(
        "--mode",
        required=True,
        choices=sorted(MAKE_MODES),
        help="Operation mode (required): merge, weight-merge, absorb, stack, or cross-merge",
    )
    make.add_argument(
        "--output-type",
        choices=sorted(PIPELINE_OUTPUT_TYPES),
        default="osmo",
        help="Output artifact type (default: osmo)",
    )
    make.add_argument(
        "--models",
        nargs="+",
        default=None,
        help="Model refs for merge/weight-merge/stack/cross-merge",
    )
    make.add_argument("--source", default=None, help="Source model ref for absorb mode")
    make.add_argument("--target", default=None, help="Target model ref for absorb mode")
    make.add_argument("-o", "--output", default=None, help="Output path")
    make.add_argument("-n", "--name", default=None, help="Artifact/model name override")
    make.add_argument(
        "--merge-strategy",
        default="auto",
        choices=["weighted_average", "slerp", "ties", "dare", "breadcrumbs", "consensus", "auto"],
        help="Merge strategy for merge/weight-merge/cross-merge (default: auto)",
    )
    make.add_argument("--weights", type=_parse_weights, default=None, help="Merge weights, e.g. 0.6,0.4")
    make.add_argument(
        "--coverage-floor",
        type=float,
        default=0.70,
        help="Minimum per-model merged byte coverage for weight-merge (default: 0.70)",
    )
    make.add_argument("--seed", type=int, default=42, help="Deterministic seed (default: 42)")
    make.add_argument(
        "--strategy-engine",
        choices=sorted(STRATEGY_ENGINE_MODES),
        default="auto",
        help="Smart strategy engine mode (default: auto)",
    )
    make.add_argument(
        "--cross-opt-stage",
        choices=["auto", "shadow", "active"],
        default="auto",
        help="Cross-merge optimizer stage control (default: auto)",
    )
    make.add_argument("--sample-data", default=None, help="Optional .npy sample data path")
    make.add_argument("--auto-grow", dest="auto_grow", action="store_true", default=True,
                      help="Enable auto-grow for absorb mode (default: enabled)")
    make.add_argument("--no-auto-grow", dest="auto_grow", action="store_false",
                      help="Disable auto-grow for absorb mode")
    make.add_argument("--top-k", type=int, default=2, help="Number of active experts for stack (default: 2)")
    make.add_argument("--train-gate", dest="train_gate", action="store_true", default=True,
                      help="Train gate for stack mode (default: enabled)")
    make.add_argument("--no-train-gate", dest="train_gate", action="store_false",
                      help="Skip gate training for stack mode")
    make.add_argument("--hf-base-model", default=None, help="HF base model ID for HF export reconstruction")
    make.add_argument("--onnx-input-shape", default=None, help="Input shape for ONNX export, e.g. 1,128")
    make.add_argument("--trace-input-shape", default=None, help="Input shape for TorchScript trace fallback")
    make.add_argument("--epochs", type=int, default=3, help="Epochs for cross-merge path (default: 3)")
    make.add_argument("--allow-partial-load", action="store_true",
                      help="Allow partial framework restore")
    make.add_argument("--allow-lossy-cast", action="store_true",
                      help="Allow degraded numeric restoration for lossy-cast tensors")
    make.add_argument("--dry-run", action="store_true",
                      help="Run preflight and routing only without writing outputs")
    make.add_argument("--report-file", default=None, help="Optional path to write operation report JSON")
    make.add_argument(
        "--self-check-level",
        choices=["off", "smoke", "parity"],
        default="smoke",
        help="Post-write self-check depth (default: smoke)",
    )
    make.add_argument(
        "--self-check",
        dest="self_check",
        action="store_true",
        default=True,
        help="Alias for --self-check-level smoke (default: enabled)",
    )
    make.add_argument(
        "--no-self-check",
        dest="self_check",
        action="store_false",
        help="Alias for --self-check-level off",
    )
    make.set_defaults(func=cmd_make)

    # ── protect ──
    protect = sub.add_parser(
        "protect",
        help="Package and protect model as trusted .osmoimg",
        description="Create a trusted .osmoimg with compression/signing/provenance controls",
    )
    protect.add_argument("model", help="Model source (path, directory, or HF model ID)")
    protect.add_argument("-o", "--output", default="model.osmoimg",
                         help="Output path (default: model.osmoimg)")
    protect.add_argument("-n", "--name", default=None, help="Artifact name override")
    protect.add_argument(
        "--compression-profile",
        choices=["legacy", "lossless-max"],
        default="lossless-max",
        help="Compression profile (default: lossless-max)",
    )
    protect.add_argument(
        "--trust-sign",
        dest="trust_sign",
        action="store_true",
        default=True,
        help="Sign artifact trust envelope (default: enabled)",
    )
    protect.add_argument(
        "--no-trust-sign",
        dest="trust_sign",
        action="store_false",
        help="Disable artifact trust signing (insecure mode)",
    )
    protect.add_argument(
        "--provenance",
        choices=["private", "full"],
        default="private",
        help="Provenance mode (default: private)",
    )
    protect.add_argument(
        "--repro-mode",
        choices=["best-effort", "strict"],
        default="strict",
        help="Reproducibility mode (default: strict)",
    )
    protect.add_argument("--signer-id", default="default",
                         help="Signer id used for trust signing (default: default)")
    protect.add_argument(
        "--signer-backend",
        choices=["local-keyring", "aws-kms"],
        default="local-keyring",
        help="Signer backend (default: local-keyring)",
    )
    protect.set_defaults(func=cmd_protect)

    # ── check ──
    check = sub.add_parser(
        "check",
        help="Unified artifact validation gate",
        description="Inspect + verify trust + self-check with optional runtime benchmark and release checks",
    )
    check.add_argument("image", help="Path to .osmoimg artifact")
    check.add_argument(
        "--self-check-level",
        choices=["off", "smoke", "parity"],
        default="smoke",
        help="Self-check depth (default: smoke)",
    )
    check.add_argument(
        "--trust-profile",
        choices=sorted(TRUST_PROFILES),
        default="standard",
        help="Trust enforcement profile (default: standard)",
    )
    check.add_argument("--require-trust", action="store_true",
                       help="Fail unless artifact trust status is trusted")
    check.add_argument("--policy-file", default=None,
                       help="Optional trust policy JSON file")
    check.add_argument("--benchmark-runtime", action="store_true",
                       help="Also run runtime benchmark gate")
    check.add_argument(
        "--benchmark-profile",
        choices=["i5_2gb", "ci_standard", "ci_gpu_standard"],
        default="ci_standard",
        help="Runtime benchmark profile when --benchmark-runtime is enabled (default: ci_standard)",
    )
    check.add_argument("--release-check", action="store_true",
                       help="Also run aggregated release-check gate")
    check.add_argument("--runtime-profile", choices=["default", "cpu-lowram", "gpu-fast"], default="default",
                       help="Runtime profile for optional benchmark (default: default)")
    check.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto",
                       help="Device policy for optional benchmark (default: auto)")
    check.add_argument("--runtime-quantization", choices=["none", "int8"], default="none",
                       help="Runtime quantization for optional benchmark (default: none)")
    check.add_argument("--json", action="store_true", help="Print JSON result")
    check.set_defaults(func=cmd_check)

    # ── launch ──
    launch = sub.add_parser(
        "launch",
        help="Unified runtime command (run/serve/export)",
        description="Launch model runtime path as run, serve, or export from one command",
    )
    launch.add_argument("image", help="Path to .osmoimg artifact")
    launch.add_argument(
        "--action",
        choices=["run", "serve", "export"],
        default="run",
        help="Runtime action (default: run)",
    )
    launch.add_argument("-p", "--prompt", default=None, help="Prompt used when --action run")
    launch.add_argument("-i", "--input", default=None, help="Input .npy path used when --action run")
    launch.add_argument("--output-file", default=None,
                        help="Output .npy path for --action run")
    launch.add_argument("--max-tokens", type=int, default=256,
                        help="Max tokens for run action (default: 256)")
    launch.add_argument("--temperature", type=float, default=0.7,
                        help="Sampling temperature for run action (default: 0.7)")
    launch.add_argument("--include-prompt", action="store_true",
                        help="Include original prompt in run output text")
    launch.add_argument("--host", default="0.0.0.0",
                        help="Host for serve action (default: 0.0.0.0)")
    launch.add_argument("--port", type=int, default=8080,
                        help="Port for serve action (default: 8080)")
    launch.add_argument("--warmup", action="store_true",
                        help="Warm up runtime before serving requests")
    launch.add_argument(
        "--format",
        default="pytorch",
        choices=["pytorch", "tensorflow", "jax", "onnx", "huggingface"],
        help="Export format for --action export (default: pytorch)",
    )
    launch.add_argument("-o", "--output", default=None,
                        help="Output path for --action export")
    launch.add_argument("--runtime-profile", choices=["default", "cpu-lowram", "gpu-fast"], default="default",
                        help="Runtime profile (default: default)")
    launch.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto",
                        help="Device policy (default: auto)")
    launch.add_argument("--runtime-quantization", choices=["none", "int8"], default="none",
                        help="Runtime quantization mode (default: none)")
    launch.add_argument("--runtime-cache", dest="runtime_cache", action="store_true", default=True,
                        help="Enable runtime cache (default: enabled)")
    launch.add_argument("--no-runtime-cache", dest="runtime_cache", action="store_false",
                        help="Disable runtime cache")
    launch.add_argument("--require-trust", action="store_true",
                        help="Fail if trust status is not trusted")
    launch.add_argument("--trust-profile", choices=sorted(TRUST_PROFILES), default="standard",
                        help="Trust enforcement profile (default: standard)")
    launch.add_argument("--policy-file", default=None, help="Optional trust policy JSON file")
    launch.set_defaults(func=cmd_launch)

    # ── expand ──
    expand = sub.add_parser(
        "expand",
        help="Adaptive Net2Net++ model expansion workflow",
        description="Grow and improve a model with quality-gated local learning and versioned .osmoimg revisions",
    )
    expand.add_argument(
        "model_ref",
        help=(
            "Model reference (hf:/pt:/onnx:/tf:/jax:/vision:/audio:/diffusion:/multimodal:/lora:/mlx:/"
            "osmo:/osmoimg: or auto path)"
        ),
    )
    expand.add_argument(
        "--task",
        choices=sorted(EXPAND_TASKS),
        default="auto",
        help="Expansion task hint (default: auto)",
    )
    expand.add_argument(
        "--dataset",
        default=None,
        help="Optional dataset path/URI (.npy/.npz recommended)",
    )
    expand.add_argument(
        "--data-mode",
        choices=["auto", "off"],
        default="auto",
        help="Automatic data planner mode (default: auto)",
    )
    expand.add_argument(
        "--max-download-gb",
        type=float,
        default=1.0,
        help="Max hub data download size in GB (default: 1.0)",
    )
    expand.add_argument(
        "--data-ttl-hours",
        type=float,
        default=24.0 * 7.0,
        help="Cached data freshness TTL in hours (default: 168)",
    )
    expand.add_argument(
        "--allow-license",
        action="append",
        dest="allow_licenses",
        default=[],
        help="Allowed data license (repeatable). Also merges with OSMOSIS_EXPAND_ALLOWED_LICENSES.",
    )
    expand.add_argument(
        "--offline",
        action="store_true",
        help="Use cached expansion data only (no hub downloads)",
    )
    expand.add_argument(
        "--target-params",
        default=None,
        help="Target parameter envelope (int) or ratio spec like x1.5",
    )
    expand.add_argument(
        "--budget-minutes",
        type=int,
        default=5,
        help="Expansion budget in minutes (default: 5)",
    )
    expand.add_argument(
        "--budget-ram-gb",
        type=float,
        default=1.5,
        help="Extra RAM budget in GB (default: 1.5)",
    )
    expand.add_argument("--seed", type=int, default=42, help="Deterministic seed (default: 42)")
    expand.add_argument("-o", "--output", default=None, help="Output .osmoimg path (default: versioned revision)")
    expand.add_argument("-n", "--name", default=None, help="Logical model/revision name override")
    expand.add_argument(
        "--export-format",
        choices=sorted(EXPAND_EXPORT_FORMATS),
        default="osmoimg",
        help="Optional additional export format (default: osmoimg only)",
    )
    expand.add_argument(
        "--optimizer-mode",
        choices=["auto", "off"],
        default="auto",
        help="Expansion optimizer mode (default: auto)",
    )
    expand.add_argument(
        "--fallback-mode",
        choices=sorted(EXPAND_FALLBACK_MODES),
        default="off",
        help="Fallback behavior when strict expansion fails (default: off)",
    )
    expand.set_defaults(func=cmd_expand)

    # ── expand-status ──
    expand_status = sub.add_parser(
        "expand-status",
        help="Show local expansion revision/optimizer status",
        description="Display expand revision chain, active revision, and local optimizer metadata",
    )
    expand_status.add_argument("--name", default=None, help="Optional revision name filter")
    expand_status.add_argument("--json", action="store_true", help="Print JSON output")
    expand_status.set_defaults(func=cmd_expand_status)

    # ── expand-rollback ──
    expand_rollback = sub.add_parser(
        "expand-rollback",
        help="Rollback active expansion revision",
        description="Set active expand revision pointer to a previous revision id",
    )
    expand_rollback.add_argument("--to", required=True, help="Target revision id (required)")
    expand_rollback.set_defaults(func=cmd_expand_rollback)

    # ── build ──
    build = sub.add_parser(
        "build", help="Build an .osmoimg image from any model",
        description="Convert any model into a portable, compressed, encrypted .osmoimg image",
    )
    build.add_argument("model", help="Model to build (path, directory, or HF model ID)")
    build.add_argument("-o", "--output", default="model.osmoimg",
                       help="Output path (default: model.osmoimg)")
    build.add_argument("-n", "--name", default=None,
                       help="Image name (default: auto-detected)")
    build.add_argument(
        "--compression-profile",
        choices=["legacy", "lossless-max"],
        default="legacy",
        help="Compression profile (default: legacy)",
    )
    build.add_argument(
        "--compression-engine",
        choices=sorted(COMPRESSION_ENGINE_MODES),
        default="auto",
        help="Compression optimizer mode (default: auto, env: OSMOSIS_OPTIMIZER)",
    )
    build.add_argument(
        "--trust-sign",
        dest="trust_sign",
        action="store_true",
        default=True,
        help="Sign artifact trust envelope (default: enabled)",
    )
    build.add_argument(
        "--no-trust-sign",
        dest="trust_sign",
        action="store_false",
        help="Disable artifact trust signing (insecure mode)",
    )
    build.add_argument(
        "--provenance",
        choices=["private", "full"],
        default="private",
        help="Provenance mode (default: private)",
    )
    build.add_argument(
        "--repro-mode",
        choices=["best-effort", "strict"],
        default="strict",
        help="Reproducibility mode (default: strict)",
    )
    build.add_argument(
        "--signer-id",
        default="default",
        help="Local signer id used for trust signing (default: default)",
    )
    build.add_argument(
        "--signer-backend",
        choices=["local-keyring", "aws-kms"],
        default="local-keyring",
        help="Signer backend (default: local-keyring)",
    )
    build.set_defaults(func=cmd_build)

    # ── demo ──
    demo = sub.add_parser(
        "demo", help="Run one-command first-time demo",
        description="Build and validate a runnable .osmoimg in one command for first-time users",
    )
    demo.add_argument(
        "--model-id",
        default=DEMO_DEFAULT_MODEL_ID,
        help=f"HuggingFace model ID (default: {DEMO_DEFAULT_MODEL_ID})",
    )
    demo.add_argument(
        "--output",
        default=DEMO_DEFAULT_OUTPUT,
        help=f"Output .osmoimg path (default: {DEMO_DEFAULT_OUTPUT})",
    )
    demo.add_argument(
        "--compression-profile",
        choices=["legacy", "lossless-max"],
        default="lossless-max",
        help="Compression profile for demo image (default: lossless-max)",
    )
    demo.add_argument(
        "--max-tokens",
        type=int,
        default=64,
        help="Max tokens for demo smoke generation (default: 64)",
    )
    demo.add_argument(
        "--temperature",
        type=float,
        default=0.2,
        help="Sampling temperature for demo smoke generation (default: 0.2)",
    )
    demo.add_argument(
        "--timeout-minutes",
        type=float,
        default=30.0,
        help="Maximum allowed end-to-end demo runtime in minutes (default: 30)",
    )
    demo.set_defaults(func=cmd_demo)

    # ── doctor ──
    doctor = sub.add_parser(
        "doctor", help="Run environment preflight diagnostics",
        description="Check dependencies, network, disk, and RAM with actionable remediation",
    )
    doctor.add_argument(
        "--compression-profile",
        choices=["legacy", "lossless-max"],
        default="lossless-max",
        help="Profile used for dependency checks (default: lossless-max)",
    )
    doctor.set_defaults(func=cmd_doctor)

    # ── run ──
    run = sub.add_parser(
        "run", help="Run inference from an .osmoimg image",
        description="Load an .osmoimg image and run inference",
    )
    run.add_argument("image", help="Path to .osmoimg image")
    run.add_argument("-p", "--prompt", default=None,
                     help="Text prompt for language models")
    run.add_argument("-i", "--input", default=None,
                     help="Path to numpy input file (.npy)")
    run.add_argument("--output-file", default=None,
                     help="Save output to file")
    run.add_argument("--max-tokens", type=int, default=256,
                     help="Maximum tokens to generate (default: 256)")
    run.add_argument("--temperature", type=float, default=0.7,
                     help="Sampling temperature (default: 0.7)")
    run.add_argument(
        "--include-prompt",
        action="store_true",
        help="Include original prompt in output text (legacy behavior)",
    )
    run.add_argument(
        "--runtime-profile",
        choices=["default", "cpu-lowram", "gpu-fast"],
        default="default",
        help="Runtime execution profile (default: default)",
    )
    run.add_argument(
        "--device",
        choices=["auto", "cpu", "cuda"],
        default="auto",
        help="Device policy (default: auto)",
    )
    run.add_argument(
        "--runtime-quantization",
        choices=["none", "int8"],
        default="none",
        help="Runtime-only quantization mode (default: none)",
    )
    run.add_argument(
        "--runtime-cache",
        dest="runtime_cache",
        action="store_true",
        default=True,
        help="Enable in-memory runtime cache (default: enabled)",
    )
    run.add_argument(
        "--no-runtime-cache",
        dest="runtime_cache",
        action="store_false",
        help="Disable in-memory runtime cache",
    )
    run.add_argument(
        "--require-trust",
        action="store_true",
        help="Fail if artifact trust status is not trusted",
    )
    run.add_argument(
        "--trust-profile",
        choices=sorted(TRUST_PROFILES),
        default="standard",
        help="Trust enforcement profile (default: standard)",
    )
    run.add_argument(
        "--policy-file",
        default=None,
        help="Optional trust policy JSON file enforced before execution",
    )
    run.set_defaults(func=cmd_run)

    # ── inspect ──
    insp = sub.add_parser(
        "inspect", help="Inspect an .osmoimg image metadata",
        description="View image metadata without loading model weights",
    )
    insp.add_argument("image", help="Path to .osmoimg image")
    insp.set_defaults(func=cmd_inspect)

    # ── export ──
    exp = sub.add_parser(
        "export", help="Export an .osmoimg image to any framework",
        description="Convert .osmoimg back to PyTorch, TensorFlow, JAX, ONNX, or HuggingFace",
    )
    exp.add_argument("image", help="Path to .osmoimg image")
    exp.add_argument("-f", "--format", default="pytorch",
                     choices=["pytorch", "tensorflow", "jax", "onnx", "huggingface"],
                     help="Target framework (default: pytorch)")
    exp.add_argument("-o", "--output", default=None,
                     help="Output path")
    exp.add_argument(
        "--require-trust",
        action="store_true",
        help="Fail if artifact trust status is not trusted",
    )
    exp.add_argument(
        "--trust-profile",
        choices=sorted(TRUST_PROFILES),
        default="standard",
        help="Trust enforcement profile (default: standard)",
    )
    exp.add_argument(
        "--policy-file",
        default=None,
        help="Optional trust policy JSON file enforced before export",
    )
    exp.set_defaults(func=cmd_export)

    # ── register ──
    reg = sub.add_parser(
        "register", help="Register .osmoimg file format with the OS",
        description="Register .osmoimg so your OS recognizes it (double-click to run)",
    )
    reg.add_argument("--unregister", action="store_true",
                     help="Remove registration instead")
    reg.add_argument(
        "--icon",
        default=None,
        help="Optional .ico path for file icon (Windows)",
    )
    reg.set_defaults(func=cmd_register)

    # ── merge ──
    merge = sub.add_parser(
        "merge", help="Merge multiple models into one",
        description="Merge 2+ models using advanced strategies",
    )
    merge.add_argument("models", nargs="+", help="Models to merge (paths or HF IDs)")
    merge.add_argument("-o", "--output", default="merged_model",
                       help="Output path (default: merged_model)")
    merge.add_argument("-s", "--strategy", default="auto",
                       choices=["weighted_average", "slerp", "ties",
                                "dare", "breadcrumbs", "consensus", "auto"],
                       help="Merge strategy (default: auto)")
    merge.add_argument(
        "--strategy-engine",
        choices=sorted(STRATEGY_ENGINE_MODES),
        default="auto",
        help="Smart strategy engine mode (default: auto, env: OSMOSIS_STRATEGY_ENGINE)",
    )
    merge.add_argument(
        "--cross-opt-stage",
        choices=["auto", "shadow", "active"],
        default="auto",
        help="Cross-merge optimizer stage control (default: auto)",
    )
    merge.add_argument("-w", "--weights", type=_parse_weights, default=None,
                       help="Model weights, e.g. '0.4,0.6'")
    merge.add_argument("--cross", action="store_true",
                       help="Cross-architecture merge (knowledge distillation)")
    merge.add_argument("--obfuscate", action="store_true",
                       help="Obfuscate after merging")
    merge.add_argument("--epochs", type=int, default=3,
                       help="Training epochs for cross-merge (default: 3)")
    merge.add_argument("--samples", type=str, default=None,
                       help="Sample texts for cross-merge, separated by |")
    merge.add_argument("--seed", type=int, default=42,
                       help="Deterministic seed for merge flows (default: 42)")
    merge.set_defaults(func=cmd_merge)

    # ── pipeline ──
    pipeline_cmd = sub.add_parser(
        "pipeline",
        help="Unified merge/weight-merge/absorb/stack pipeline with protected outputs",
        description=(
            "Run merge/weight-merge/absorb/stack from one command and export as "
            "osmoimg/osmo/pytorch/hf/onnx/tensorflow/keras/jax/flax/safetensors/torchscript "
            "with invisible protection."
        ),
    )
    pipeline_cmd.add_argument(
        "--mode",
        required=True,
        choices=sorted(PIPELINE_MODES),
        help="Pipeline mode (required): merge, absorb, stack, or weight-merge",
    )
    pipeline_cmd.add_argument(
        "--output-type",
        required=True,
        choices=sorted(PIPELINE_OUTPUT_TYPES),
        help="Output artifact type (required)",
    )
    pipeline_cmd.add_argument(
        "--models",
        nargs="+",
        default=None,
        help="Model refs for merge/stack. Supports hf:/pt:/onnx:/tf:/jax: prefixes",
    )
    pipeline_cmd.add_argument(
        "--source",
        default=None,
        help="Source model ref for absorb mode",
    )
    pipeline_cmd.add_argument(
        "--target",
        default=None,
        help="Target model ref for absorb mode",
    )
    pipeline_cmd.add_argument(
        "-o",
        "--output",
        default=None,
        help="Output path (default is deterministic by mode + output type)",
    )
    pipeline_cmd.add_argument(
        "-n",
        "--name",
        default=None,
        help="Artifact/model name override",
    )
    pipeline_cmd.add_argument(
        "--compression-profile",
        choices=["legacy", "lossless-max"],
        default="lossless-max",
        help="Compression profile for osmoimg output (default: lossless-max)",
    )
    pipeline_cmd.add_argument(
        "--compression-engine",
        choices=sorted(COMPRESSION_ENGINE_MODES),
        default="auto",
        help="Compression optimizer mode for osmoimg output (default: auto)",
    )
    pipeline_cmd.add_argument(
        "--sample-data",
        default=None,
        help="Optional .npy sample data path",
    )
    pipeline_cmd.add_argument(
        "--auto-grow",
        dest="auto_grow",
        action="store_true",
        default=True,
        help="Enable auto-grow for absorb mode (default: enabled)",
    )
    pipeline_cmd.add_argument(
        "--no-auto-grow",
        dest="auto_grow",
        action="store_false",
        help="Disable auto-grow for absorb mode",
    )
    pipeline_cmd.add_argument(
        "--merge-strategy",
        default="auto",
        choices=["weighted_average", "slerp", "ties", "dare", "breadcrumbs", "consensus", "auto"],
        help="Merge strategy when mode=merge/weight-merge (default: auto)",
    )
    pipeline_cmd.add_argument(
        "--strategy-engine",
        choices=sorted(STRATEGY_ENGINE_MODES),
        default="auto",
        help="Smart strategy engine mode for merge/weight-merge (default: auto)",
    )
    pipeline_cmd.add_argument(
        "--cross-opt-stage",
        choices=["auto", "shadow", "active"],
        default="auto",
        help="Cross-merge optimizer stage control for mode=merge (default: auto)",
    )
    pipeline_cmd.add_argument(
        "--weights",
        type=_parse_weights,
        default=None,
        help="Merge weights, e.g. '0.5,0.5'",
    )
    pipeline_cmd.add_argument(
        "--coverage-floor",
        type=float,
        default=0.70,
        help="Minimum per-model merged byte coverage for mode=weight-merge (range: 0..1, default: 0.70)",
    )
    pipeline_cmd.add_argument(
        "--top-k",
        type=int,
        default=2,
        help="Number of active experts for stack mode (default: 2)",
    )
    pipeline_cmd.add_argument(
        "--train-gate",
        dest="train_gate",
        action="store_true",
        default=True,
        help="Train gate for stack mode (default: enabled)",
    )
    pipeline_cmd.add_argument(
        "--no-train-gate",
        dest="train_gate",
        action="store_false",
        help="Skip gate training for stack mode",
    )
    pipeline_cmd.add_argument(
        "--hf-base-model",
        default=None,
        help="HF base model ID for reconstructing HF export when needed",
    )
    pipeline_cmd.add_argument(
        "--onnx-input-shape",
        default=None,
        help="Input shape for ONNX export, e.g. 1,128",
    )
    pipeline_cmd.add_argument(
        "--trace-input-shape",
        default=None,
        help="Input shape for TorchScript trace fallback, e.g. 1,128",
    )
    pipeline_cmd.add_argument(
        "--epochs",
        type=int,
        default=3,
        help="Epochs used when merge auto-routes to cross-merge (default: 3)",
    )
    pipeline_cmd.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Deterministic seed for pipeline execution (default: 42)",
    )
    pipeline_cmd.add_argument(
        "--allow-partial-load",
        action="store_true",
        help="Allow partial framework restore (disables strict key matching)",
    )
    pipeline_cmd.add_argument(
        "--allow-lossy-cast",
        action="store_true",
        help="Allow degraded numeric restoration for lossy-cast tensors",
    )
    pipeline_cmd.add_argument(
        "--trust-sign",
        dest="trust_sign",
        action="store_true",
        default=True,
        help="Sign osmoimg output trust envelope (default: enabled)",
    )
    pipeline_cmd.add_argument(
        "--no-trust-sign",
        dest="trust_sign",
        action="store_false",
        help="Disable trust signing for osmoimg output (insecure mode)",
    )
    pipeline_cmd.add_argument(
        "--provenance",
        choices=["private", "full"],
        default="private",
        help="Provenance mode for osmoimg output (default: private)",
    )
    pipeline_cmd.add_argument(
        "--repro-mode",
        choices=["best-effort", "strict"],
        default="strict",
        help="Reproducibility mode for osmoimg output (default: strict)",
    )
    pipeline_cmd.add_argument(
        "--signer-id",
        default="default",
        help="Local signer id for osmoimg trust signature (default: default)",
    )
    pipeline_cmd.add_argument(
        "--signer-backend",
        choices=["local-keyring", "aws-kms"],
        default="local-keyring",
        help="Signer backend for osmoimg signing (default: local-keyring)",
    )
    pipeline_cmd.add_argument(
        "--dry-run",
        action="store_true",
        help="Run preflight and route resolution only; do not execute operations or write outputs",
    )
    pipeline_cmd.add_argument(
        "--report-file",
        default=None,
        help="Optional path to write deterministic pipeline operation report JSON",
    )
    pipeline_cmd.add_argument(
        "--self-check-level",
        choices=["off", "smoke", "parity"],
        default="smoke",
        help="Post-write self-check depth (default: smoke)",
    )
    pipeline_cmd.add_argument(
        "--self-check",
        dest="self_check",
        action="store_true",
        default=True,
        help="Alias for --self-check-level smoke (default: enabled)",
    )
    pipeline_cmd.add_argument(
        "--no-self-check",
        dest="self_check",
        action="store_false",
        help="Alias for --self-check-level off",
    )
    pipeline_cmd.set_defaults(func=cmd_pipeline)

    # ── obfuscate ──
    obfuscate = sub.add_parser(
        "obfuscate", help="Obfuscate model weights (invisible encryption)",
        description="Apply function-preserving transforms to prevent weight tracing",
    )
    obfuscate.add_argument("model", help="Model to obfuscate (path or HF ID)")
    obfuscate.add_argument("-o", "--output", default=None,
                           help="Output path")
    obfuscate.add_argument("--preserve-quality", action="store_true", default=True,
                           help="Use exact transforms only (default: True)")
    obfuscate.add_argument("--no-preserve-quality", dest="preserve_quality",
                           action="store_false",
                           help="Allow rotation (may reduce quality slightly)")
    obfuscate.add_argument("--seed", type=int, default=None,
                           help="Random seed for reproducibility")
    obfuscate.add_argument("--rounds", type=int, default=1,
                           help="Number of obfuscation rounds (default: 1)")
    obfuscate.set_defaults(func=cmd_obfuscate)

    # ── compare ──
    compare = sub.add_parser(
        "compare", help="Compare two models' outputs",
        description="Run both models on sample data and report similarity",
    )
    compare.add_argument("model_a", help="First model")
    compare.add_argument("model_b", help="Second model")
    compare.add_argument("-n", "--samples", type=int, default=100,
                         help="Number of random samples (default: 100)")
    compare.set_defaults(func=cmd_compare)

    # ── benchmark ──
    bench = sub.add_parser(
        "benchmark", help="Benchmark all merge strategies",
        description="Try all strategies and rank by output consistency",
    )
    bench.add_argument("models", nargs="+", help="Models to benchmark")
    bench.set_defaults(func=cmd_benchmark)

    # ── benchmark-runtime ──
    bench_runtime = sub.add_parser(
        "benchmark-runtime",
        help="Benchmark runtime KPI metrics for an .osmoimg",
        description="Measure cold start, memory, latency, and throughput against a baseline profile",
    )
    bench_runtime.add_argument("image", help="Path to .osmoimg artifact")
    bench_runtime.add_argument(
        "--profile",
        choices=["i5_2gb", "ci_standard", "ci_gpu_standard"],
        default="i5_2gb",
        help="Runtime profile baseline (default: i5_2gb)",
    )
    bench_runtime.add_argument(
        "--baseline-file",
        default=str(RUNTIME_BENCHMARK_BASELINES),
        help=f"Baseline threshold JSON file (default: {RUNTIME_BENCHMARK_BASELINES})",
    )
    bench_runtime.add_argument(
        "--iterations",
        type=int,
        default=5,
        help="Latency measurement iterations for p95 (default: 5)",
    )
    bench_runtime.add_argument(
        "--prompt",
        default=RUNTIME_BENCHMARK_DEFAULT_PROMPT,
        help="Prompt used for runtime measurement",
    )
    bench_runtime.add_argument(
        "--max-tokens",
        type=int,
        default=32,
        help="Tokens used for throughput/latency measurements (default: 32)",
    )
    bench_runtime.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="Sampling temperature used for benchmark runs (default: 0.0)",
    )
    bench_runtime.add_argument(
        "--warmup-runs",
        type=int,
        default=2,
        help="Warmup runs before warm-metric measurements (default: 2)",
    )
    bench_runtime.add_argument(
        "--runtime-profile",
        choices=["default", "cpu-lowram", "gpu-fast"],
        default="default",
        help="Runtime execution profile (default: default)",
    )
    bench_runtime.add_argument(
        "--device",
        choices=["auto", "cpu", "cuda"],
        default="auto",
        help="Device policy (default: auto)",
    )
    bench_runtime.add_argument(
        "--runtime-quantization",
        choices=["none", "int8"],
        default="none",
        help="Runtime-only quantization mode (default: none)",
    )
    bench_runtime.add_argument(
        "--runtime-cache",
        dest="runtime_cache",
        action="store_true",
        default=True,
        help="Enable in-memory runtime cache (default: enabled)",
    )
    bench_runtime.add_argument(
        "--no-runtime-cache",
        dest="runtime_cache",
        action="store_false",
        help="Disable in-memory runtime cache",
    )
    bench_runtime.add_argument(
        "--json",
        action="store_true",
        help="Print benchmark result JSON",
    )
    bench_runtime.set_defaults(func=cmd_benchmark_runtime)

    # ── info ──
    info = sub.add_parser(
        "info", help="Show model information",
        description="Display model architecture, parameters, and layers",
    )
    info.add_argument("model", help="Model to inspect (path or HF ID)")
    info.set_defaults(func=cmd_info)

    # ── version ──
    ver = sub.add_parser("version", help="Print version")
    ver.set_defaults(func=cmd_version)

    # ── images ──
    images = sub.add_parser(
        "images", help="List all .osmoimg files in a directory",
        description="Scan a directory for .osmoimg files and display a summary table",
    )
    images.add_argument("directory", nargs="?", default=".",
                        help="Directory to scan (default: current directory)")
    images.set_defaults(func=cmd_images)

    # ── pull ──
    pull = sub.add_parser(
        "pull", help="Download a HuggingFace model as .osmoimg",
        description="Download a model from HuggingFace, encrypt, compress, and save as .osmoimg",
    )
    pull.add_argument("model", help="HuggingFace model ID (e.g. user/ModelName)")
    pull.add_argument("-o", "--output", default=None,
                      help="Output path (default: <model_name>.osmoimg)")
    pull.add_argument(
        "--compression-profile",
        choices=["legacy", "lossless-max"],
        default="legacy",
        help="Compression profile (default: legacy)",
    )
    pull.add_argument(
        "--compression-engine",
        choices=sorted(COMPRESSION_ENGINE_MODES),
        default="auto",
        help="Compression optimizer mode (default: auto, env: OSMOSIS_OPTIMIZER)",
    )
    pull.add_argument(
        "--trust-sign",
        dest="trust_sign",
        action="store_true",
        default=True,
        help="Sign artifact trust envelope (default: enabled)",
    )
    pull.add_argument(
        "--no-trust-sign",
        dest="trust_sign",
        action="store_false",
        help="Disable artifact trust signing (insecure mode)",
    )
    pull.add_argument(
        "--provenance",
        choices=["private", "full"],
        default="private",
        help="Provenance mode (default: private)",
    )
    pull.add_argument(
        "--repro-mode",
        choices=["best-effort", "strict"],
        default="strict",
        help="Reproducibility mode (default: strict)",
    )
    pull.add_argument(
        "--signer-id",
        default="default",
        help="Local signer id used for trust signing (default: default)",
    )
    pull.add_argument(
        "--signer-backend",
        choices=["local-keyring", "aws-kms"],
        default="local-keyring",
        help="Signer backend (default: local-keyring)",
    )
    pull.set_defaults(func=cmd_pull)

    # ── transcode ──
    transcode = sub.add_parser(
        "transcode", help="Recompress an .osmoimg with a selected profile",
        description="Load an existing .osmoimg and rewrite it using legacy or lossless-max compression",
    )
    transcode.add_argument("image", help="Path to source .osmoimg file")
    transcode.add_argument(
        "-o", "--output", default=None,
        help="Output path (default: <input>.v2.osmoimg)",
    )
    transcode.add_argument(
        "--compression-profile",
        choices=["legacy", "lossless-max"],
        default="lossless-max",
        help="Compression profile for output (default: lossless-max)",
    )
    transcode.add_argument(
        "--compression-engine",
        choices=sorted(COMPRESSION_ENGINE_MODES),
        default="auto",
        help="Compression optimizer mode (default: auto, env: OSMOSIS_OPTIMIZER)",
    )
    transcode.add_argument(
        "--trust-sign",
        dest="trust_sign",
        action="store_true",
        default=True,
        help="Sign output artifact trust envelope (default: enabled)",
    )
    transcode.add_argument(
        "--no-trust-sign",
        dest="trust_sign",
        action="store_false",
        help="Disable artifact trust signing for transcode output (insecure mode)",
    )
    transcode.add_argument(
        "--upgrade-trust",
        action="store_true",
        help="Alias for trust upgrade path (equivalent to forcing trust-sign on output)",
    )
    transcode.add_argument(
        "--signer-id",
        default="default",
        help="Signer id used when --upgrade-trust is enabled (default: default)",
    )
    transcode.add_argument(
        "--signer-backend",
        choices=["local-keyring", "aws-kms"],
        default="local-keyring",
        help="Signer backend used for trust upgrade (default: local-keyring)",
    )
    transcode.add_argument(
        "--migration-report",
        default=None,
        help="Optional path to write migration report JSON",
    )
    transcode.set_defaults(func=cmd_transcode)

    # ── push ──
    push = sub.add_parser(
        "push", help="Upload .osmoimg to HuggingFace Hub",
        description="Upload an .osmoimg file to a HuggingFace repository",
    )
    push.add_argument("image", help="Path to .osmoimg file")
    push.add_argument("repo", help="HuggingFace repo ID (e.g. user/repo)")
    push.add_argument("--private", action="store_true",
                      help="Create a private repository")
    push.set_defaults(func=cmd_push)

    # ── convert ──
    conv = sub.add_parser(
        "convert", help="Convert model between frameworks",
        description="Convert any model file or .osmoimg to another framework format",
    )
    conv.add_argument("model", help="Model to convert (path or .osmoimg)")
    conv.add_argument("--to", required=True,
                      choices=["pytorch", "tensorflow", "jax", "onnx", "huggingface"],
                      help="Target framework")
    conv.add_argument("-o", "--output", default=None,
                      help="Output path")
    conv.add_argument("--seed", type=int, default=42,
                      help="Deterministic seed for conversion (default: 42)")
    conv.add_argument("--allow-partial-load", action="store_true",
                      help="Allow partial framework restore (disables strict key matching)")
    conv.add_argument("--allow-lossy-cast", action="store_true",
                      help="Allow degraded numeric restoration for lossy-cast tensors")
    conv.set_defaults(func=cmd_convert)

    # ── diff ──
    diff = sub.add_parser(
        "diff", help="Compare two .osmoimg image metadata",
        description="Show side-by-side differences in metadata, architecture, and params",
    )
    diff.add_argument("image_a", help="First .osmoimg image")
    diff.add_argument("image_b", help="Second .osmoimg image")
    diff.set_defaults(func=cmd_diff)

    # ── validate ──
    val = sub.add_parser(
        "validate", help="Verify .osmoimg integrity",
        description="Check magic bytes, header, zip payload, manifest, and parameters",
    )
    val.add_argument("image", help="Path to .osmoimg file")
    val.set_defaults(func=cmd_validate)

    # ── verify ──
    verify = sub.add_parser(
        "verify", help="Verify trusted-signature status of .osmoimg",
        description="Verify trust signature, signer trust, and payload integrity",
    )
    verify.add_argument("image", help="Path to .osmoimg file")
    verify.add_argument("--require-trust", action="store_true",
                        help="Return non-zero unless trust status is trusted")
    verify.add_argument("--policy-file", default=None,
                        help="Optional trust policy JSON file enforced during verification")
    verify.add_argument("--repro-check", action="store_true",
                        help="Run strict reproducibility check by deterministic re-serialization")
    verify.add_argument("--json", action="store_true",
                        help="Print JSON verification result")
    verify.set_defaults(func=cmd_verify)

    # ── quantize ──
    quant = sub.add_parser(
        "quantize", help="Quantize model weights",
        description="Reduce model size by quantizing weights to lower precision",
    )
    quant.add_argument("image", help="Path to .osmoimg file")
    quant.add_argument("--bits", type=int, required=True, choices=[4, 8, 16],
                       help="Quantization bits (4, 8, or 16)")
    quant.add_argument("-o", "--output", default=None,
                       help="Output path (default: overwrite input)")
    quant.set_defaults(func=cmd_quantize)

    # ── serve ──
    serve = sub.add_parser(
        "serve", help="Start HTTP inference server",
        description="Load an .osmoimg and serve inference via HTTP endpoints",
    )
    serve.add_argument("image", help="Path to .osmoimg file")
    serve.add_argument("--host", default="0.0.0.0",
                       help="Host to bind (default: 0.0.0.0)")
    serve.add_argument("--port", type=int, default=8080,
                       help="Port to listen on (default: 8080)")
    serve.add_argument(
        "--runtime-profile",
        choices=["default", "cpu-lowram", "gpu-fast"],
        default="default",
        help="Runtime execution profile (default: default)",
    )
    serve.add_argument(
        "--device",
        choices=["auto", "cpu", "cuda"],
        default="auto",
        help="Device policy (default: auto)",
    )
    serve.add_argument(
        "--runtime-quantization",
        choices=["none", "int8"],
        default="none",
        help="Runtime-only quantization mode (default: none)",
    )
    serve.add_argument(
        "--runtime-cache",
        dest="runtime_cache",
        action="store_true",
        default=True,
        help="Enable in-memory runtime cache (default: enabled)",
    )
    serve.add_argument(
        "--no-runtime-cache",
        dest="runtime_cache",
        action="store_false",
        help="Disable in-memory runtime cache",
    )
    serve.add_argument(
        "--warmup",
        action="store_true",
        help="Run one startup warmup inference call before serving requests",
    )
    serve.add_argument("--require-trust", action="store_true",
                       help="Fail if artifact trust status is not trusted")
    serve.add_argument(
        "--trust-profile",
        choices=sorted(TRUST_PROFILES),
        default="standard",
        help="Trust enforcement profile (default: standard)",
    )
    serve.add_argument("--policy-file", default=None,
                       help="Optional trust policy JSON file enforced before serving")
    serve.set_defaults(func=cmd_serve)

    # ── history ──
    hist = sub.add_parser(
        "history", help="Show build/export history",
        description="Display build timestamp, source framework, encryption status, and metadata",
    )
    hist.add_argument("image", help="Path to .osmoimg file")
    hist.set_defaults(func=cmd_history)

    # ── release-check ──
    release_check = sub.add_parser(
        "release-check",
        help="Run aggregated release-blocking checks (local + CI by default)",
        description=(
            "Run local reliability gates (pipeline/trust/runtime/cross-framework) and "
            "strictly validate required GitHub workflow status for the current commit SHA."
        ),
    )
    release_check.add_argument(
        "--json",
        action="store_true",
        help="Print deterministic JSON report",
    )
    release_check.add_argument(
        "--local-only",
        action="store_true",
        help="Skip CI workflow checks and run local gates only",
    )
    release_check.add_argument(
        "--with-ci",
        action="store_true",
        default=True,
        help="Include GitHub Actions workflow status summary via gh CLI",
    )
    release_check.add_argument(
        "--strict-ci",
        action="store_true",
        default=True,
        help="Fail if CI status cannot be fetched when --with-ci is used",
    )
    release_check.set_defaults(func=cmd_release_check)

    # ── optimize-status ──
    optimize_status = sub.add_parser(
        "optimize-status",
        help="Show local self-optimizer status",
        description="Inspect deterministic local optimizer state for strategy/compression tuning",
    )
    optimize_status.add_argument(
        "--json",
        action="store_true",
        help="Print optimizer status as JSON",
    )
    optimize_status.set_defaults(func=cmd_optimize_status)

    # ── optimize-reset ──
    optimize_reset = sub.add_parser(
        "optimize-reset",
        help="Reset local self-optimizer state",
        description="Clear local optimizer state under .model_osmosis/self_optimize",
    )
    optimize_reset.add_argument(
        "--json",
        action="store_true",
        help="Print post-reset status as JSON",
    )
    optimize_reset.set_defaults(func=cmd_optimize_reset)

    # ── artifact ──
    artifact_cmd = sub.add_parser(
        "artifact",
        help="Track and verify .osmoimg artifacts in git workflows",
        description="Artifact lock tracking, verification, and git hook installation",
    )
    artifact_sub = artifact_cmd.add_subparsers(dest="artifact_command", help="Artifact command")

    artifact_track = artifact_sub.add_parser(
        "track",
        help="Track one or more .osmoimg artifacts into lock file",
    )
    artifact_track.add_argument(
        "images",
        nargs="+",
        help="Artifact path(s) to track",
    )
    artifact_track.add_argument(
        "--lock-file",
        default=str(ARTIFACT_LOCKFILE_DEFAULT),
        help=f"Artifact lock file path (default: {ARTIFACT_LOCKFILE_DEFAULT})",
    )
    artifact_track.set_defaults(func=cmd_artifact_track)

    artifact_verify = artifact_sub.add_parser(
        "verify",
        help="Verify tracked artifact metadata and integrity drift",
    )
    artifact_verify.add_argument(
        "image",
        nargs="?",
        default=None,
        help="Single tracked artifact path to verify (omit with --all)",
    )
    artifact_verify.add_argument(
        "--all",
        action="store_true",
        help="Verify all artifacts listed in lock file",
    )
    artifact_verify.add_argument(
        "--lock-file",
        default=str(ARTIFACT_LOCKFILE_DEFAULT),
        help=f"Artifact lock file path (default: {ARTIFACT_LOCKFILE_DEFAULT})",
    )
    artifact_verify.add_argument(
        "--json",
        action="store_true",
        help="Print verification result as JSON",
    )
    artifact_verify.set_defaults(func=cmd_artifact_verify)

    artifact_hook = artifact_sub.add_parser(
        "hook",
        help="Manage git hooks for artifact verification",
    )
    artifact_hook_sub = artifact_hook.add_subparsers(dest="artifact_hook_command", help="Hook command")
    artifact_hook_install = artifact_hook_sub.add_parser(
        "install",
        help="Install pre-commit and pre-push hooks that run artifact verification",
    )
    artifact_hook_install.add_argument(
        "--force",
        action="store_true",
        help="Overwrite existing non-managed hooks",
    )
    artifact_hook_install.add_argument(
        "--lock-file",
        default=str(ARTIFACT_LOCKFILE_DEFAULT),
        help=f"Artifact lock file path for hook command (default: {ARTIFACT_LOCKFILE_DEFAULT})",
    )
    artifact_hook_install.set_defaults(func=cmd_artifact_hook_install)

    # ── trust ──
    trust_cmd = sub.add_parser(
        "trust", help="Manage local trust keys and trusted signers",
        description="Key management for trusted .osmoimg signatures",
    )
    trust_sub = trust_cmd.add_subparsers(dest="trust_command", help="Trust command to run")

    trust_keygen = trust_sub.add_parser("keygen", help="Generate a local signer keypair")
    trust_keygen.add_argument("--name", required=True, help="Signer name/id")
    trust_keygen.add_argument("--overwrite", action="store_true",
                              help="Overwrite existing signer keypair")
    trust_keygen.set_defaults(func=cmd_trust_keygen)

    trust_list = trust_sub.add_parser("list-keys", help="List local keys and trusted signers")
    trust_list.set_defaults(func=cmd_trust_list_keys)

    trust_import = trust_sub.add_parser("import-signer", help="Import trusted signer public key")
    trust_import.add_argument("--name", required=True, help="Signer name/id")
    trust_import.add_argument("--pubkey", required=True, help="Public key file path")
    trust_import.set_defaults(func=cmd_trust_import_signer)

    trust_revoke = trust_sub.add_parser("revoke-signer", help="Revoke trusted signer")
    trust_revoke.add_argument("--name", required=True, help="Signer name/id")
    trust_revoke.set_defaults(func=cmd_trust_revoke_signer)

    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()

    if args.command is None:
        print(_logo())
        print()
        parser.print_help()
        return
    if not hasattr(args, "func"):
        parser.print_help()
        return

    try:
        _ensure_windows_auto_register(getattr(args, "command", None))
        args.func(args)
    except KeyboardInterrupt:
        print(f"\n  {Colors.YELLOW}Interrupted{Colors.RESET}")
        sys.exit(1)
    except Exception as e:
        message = str(e).strip()
        if "Next command:" not in message:
            message = f"{message}\n{_next_command_line(_infer_next_command(args, message))}"
        _print_error(message)
        sys.exit(1)


if __name__ == "__main__":
    main()
