"""Load the extracted selected components with the recorded native library versions.

This example does not install the native router, warmups, input shell, controller
witnesses or temporal budget policy, and is not a full-pipeline runtime claim.
"""
import argparse
from importlib.metadata import version
import os
from pathlib import Path
import platform

EXPECTED_VERSIONS = {
    'Pillow': '12.3.0', 'accelerate': '1.14.0', 'av': '16.1.0',
    'bitsandbytes': '0.50.1', 'huggingface-hub': '1.29.0', 'peft': '0.20.0',
    'safetensors': '0.8.0', 'tokenizers': '0.22.2', 'torch': '2.13.0',
    'transformers': '5.14.0',
}


def load(snapshot):
    if platform.python_version() != '3.11.11':
        raise RuntimeError('Use the recorded Python 3.11.11 environment')
    actual = {name: version(name) for name in EXPECTED_VERSIONS}
    if actual != EXPECTED_VERSIONS:
        raise RuntimeError('Installed library versions differ from the retained image')
    for name in ('HF_HUB_OFFLINE', 'TRANSFORMERS_OFFLINE', 'HF_DATASETS_OFFLINE'):
        os.environ[name] = '1'
    from verify_files import verify
    verify(snapshot, restore_aliases=True)
    import torch
    from transformers import AutoModelForImageTextToText, AutoProcessor
    from peft import PeftModel
    if not torch.cuda.is_available():
        raise RuntimeError('The selected int8 component-loading example requires CUDA')
    artifacts = Path(snapshot).resolve(strict=True) / 'image_root' / 'app' / 'artifacts'
    processor = AutoProcessor.from_pretrained(
        str(artifacts / 'q38_base'), trust_remote_code=True, local_files_only=True)
    # The saved prequant config already declares bitsandbytes int8. Do not
    # supply a new BitsAndBytesConfig or re-quantize the base at load time.
    base = AutoModelForImageTextToText.from_pretrained(
        str(artifacts / 'q38_base_int8_prequant'),
        torch_dtype=torch.bfloat16, device_map='cuda', attn_implementation='sdpa',
        trust_remote_code=True, local_files_only=True)
    model = PeftModel.from_pretrained(
        base, str(artifacts / 'proc_q38_dense48_ck12910'),
        adapter_name='default', is_trainable=False, local_files_only=True)
    devices = sorted({p.device.index for p in model.parameters() if p.device.type == 'cuda'})
    with torch.random.fork_rng(devices=devices):
        result = model.load_adapter(
            str(artifacts / 'proc_q38_point_ck12910'), 'point_object',
            is_trainable=False, local_files_only=True, autocast_adapter_dtype=True)
    if result.missing_keys or result.unexpected_keys:
        raise RuntimeError('Point adapter load has missing or unexpected keys')
    if set(model.peft_config) != {'default', 'point_object'}:
        raise RuntimeError('The selected component pair was not loaded')
    model.set_adapter('default', inference_mode=True)
    model.eval()
    return processor, model


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('snapshot', type=Path)
    args = parser.parse_args()
    processor, model = load(args.snapshot)
    print({'processor_type': type(processor).__name__, 'model_type': type(model).__name__,
           'active_adapter': model.active_adapter, 'full_pipeline_qualified': False})


if __name__ == '__main__':
    main()