"""Local SAM image detector from a SAM 3.1 multiplex checkpoint.

No downloads, service, video tracker, or silently uninitialized parameters.
"""
import gc
import hashlib
import importlib.metadata
import os
from pathlib import Path
import sys


def check_environment(code_dir):
    if Path(sys.prefix).name != 'ai-spine' or not (Path(sys.prefix) / 'conda-meta').is_dir():
        raise RuntimeError('Run inside Conda ai-spine: conda run -n ai-spine ...')
    code = Path(code_dir).resolve()
    if not (code / 'sam3' / 'model_builder.py').is_file():
        raise FileNotFoundError(f'SAM source checkout missing: {code}')
    sys.path.insert(0, str(code))
    os.environ['HF_HUB_OFFLINE'] = '1'
    os.environ['TRANSFORMERS_OFFLINE'] = '1'
    import torch
    import sam3
    from sam3.model_builder import build_sam3_image_model  # noqa: F401
    from sam3.model.sam3_image_processor import Sam3Processor  # noqa: F401
    if Path(sam3.__file__).resolve().parent != code / 'sam3':
        raise RuntimeError('Imported SAM package does not match --sam-code-dir')
    if not torch.cuda.is_available():
        raise RuntimeError('CUDA GPU unavailable; this stage requires GPU inference')
    source_hash = hashlib.sha256()
    for file in sorted((code / 'sam3').rglob('*')):
        if file.is_file() and (file.suffix == '.py' or file.name.endswith('.txt.gz')):
            source_hash.update(file.relative_to(code).as_posix().encode())
            source_hash.update(file.read_bytes())
    versions = {}
    for name in ['torch', 'torchvision', 'numpy', 'Pillow', 'opencv-python', 'timm', 'triton-windows']:
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = None
    return {'python': sys.executable, 'sam_source': str(code), 'sam_source_sha256': source_hash.hexdigest(),
            'versions': versions, 'cuda': torch.version.cuda, 'gpu': torch.cuda.get_device_name(0),
            'gpu_total_bytes': torch.cuda.get_device_properties(0).total_memory,
            'bf16_supported': torch.cuda.is_bf16_supported(), 'network_requested': False}


def load_detector(checkpoint, precision='bf16'):
    import torch
    from sam3.model_builder import build_sam3_image_model
    model = build_sam3_image_model(device='cpu', load_from_HF=False,
                                  enable_inst_interactivity=False, compile=False)
    weights = torch.load(checkpoint, map_location='cpu', weights_only=True, mmap=True)
    weights = weights.get('model', weights)
    detector = {key.removeprefix('detector.'): value for key, value in weights.items() if key.startswith('detector.')}
    if not detector:
        raise RuntimeError('Checkpoint has no detector.* weights; refusing an uninitialized model')
    # The image builder uses four scales then scalp=1 discards the last one.
    # SAM 3.1 stores only the three used detection scales. Remove the unused
    # fourth scale and scalp, preserving exactly the same three feature maps.
    multiplex = any(k.startswith('backbone.vision_backbone.interactive_convs.') for k in detector)
    if multiplex:
        neck = model.backbone.vision_backbone
        if len(neck.convs) != 4 or model.backbone.scalp != 1:
            raise RuntimeError('Unsupported SAM image backbone layout; review adapter for this source version')
        neck.convs = torch.nn.ModuleList(list(neck.convs)[:3])
        neck.scale_factors = list(neck.scale_factors)[:3]
        model.backbone.scalp = 0
    ignored_prefixes = ('backbone.vision_backbone.interactive_convs.',
                        'backbone.vision_backbone.propagation_convs.') if multiplex else ()
    selected = {k: v for k, v in detector.items() if not k.startswith(ignored_prefixes)}
    # strict=True checks every active parameter and every tensor shape.
    model.load_state_dict(selected, strict=True)
    info = {'checkpoint_kind': 'sam3.1_multiplex_detector' if multiplex else 'sam3_detector',
            'strict_load': True, 'loaded_tensors': len(selected),
            'ignored_video_neck_tensors': len(detector) - len(selected),
            'instance_interactivity': False, 'compile': False, 'precision': precision}
    del selected, detector, weights
    gc.collect()
    if precision == 'bf16' and not torch.cuda.is_bf16_supported():
        raise RuntimeError('BF16 unsupported on this GPU; use --sam-precision fp32')
    # Keep FP32 parameters and complex RoPE buffers intact. The upstream
    # decoder explicitly disables autocast for its FP32 FFN; casting the whole
    # model to BF16 would both break it and discard complex imaginary values.
    model = model.to(device='cuda').eval()
    return model, info