"""CAPTURE-RGB-01: SDK retrieve/extract is RGB; no model-input conversion here."""
import os
import time

_from_dlpack = None
_preparation = {'ready': False}


def preparation_status():
    return dict(_preparation)


def prepare_capture():
    """Opt-in lifespan preparation, before any Stream can run; failures disable capture only."""
    global _from_dlpack, _preparation
    if _preparation.get('ready'):
        return preparation_status()
    started = time.monotonic_ns()
    try:
        import torch
        from torch.utils.dlpack import from_dlpack
        imported = time.monotonic_ns()
        # Warm this process's debug transfer path, not a model or SDK-owned surface.
        value = torch.zeros(1, dtype=torch.uint8, device='cuda:'+os.environ.get('GPU_ID', '0'))
        value.cpu()
        del value
        finished = time.monotonic_ns()
        _from_dlpack = from_dlpack
        _preparation = dict(ready=True, import_ms=(imported-started)/1e6,
                            cuda_transfer_prepare_ms=(finished-imported)/1e6)
    except Exception as error:
        _preparation = dict(ready=False, error_type=type(error).__name__,
                            elapsed_ms=(time.monotonic_ns()-started)/1e6)
    return preparation_status()


def copy_sdk_rgb(tensor):
    """Copy RGB HWC bytes before SDK buffer return; do not retain GPU references."""
    convert = _from_dlpack
    if convert is None:
        from torch.utils.dlpack import from_dlpack
        convert = from_dlpack
    return convert(tensor.__dlpack__(None)).cpu().numpy()