"""Check the tested four-GPU environment and shard headers without loading weights."""
import importlib.metadata
import json
import os
import struct
from pathlib import Path

import torch

from flashinfer.fused_moe import B12xMoEWrapper
from transformers import AutoConfig, AutoTokenizer


def main():
    versions = {name: importlib.metadata.version(name) for name in (
        'vllm', 'torch', 'transformers', 'accelerate', 'tokenizers', 'huggingface-hub',
        'safetensors', 'flashinfer-python', 'nvidia-cutlass-dsl', 'triton')}
    expected = {'vllm': '0.28.0', 'torch': '2.13.0+cu130', 'transformers': '5.17.0', 'accelerate': '1.15.0',
                'tokenizers': '0.23.2', 'huggingface-hub': '1.33.0', 'safetensors': '0.8.0',
                'flashinfer-python': '0.7.0', 'nvidia-cutlass-dsl': '4.6.2', 'triton': '3.7.1'}
    assert versions == expected, {'expected': expected, 'actual': versions}
    assert torch.cuda.device_count() == 4, 'This deployment recipe requires four visible GPUs.'
    gpus = []
    for index in range(4):
        props = torch.cuda.get_device_properties(index)
        assert (props.major, props.minor) == (12, 0), 'The fused B12x recipe targets SM120 Blackwell.'
        assert props.total_memory >= 90 * 1024**3, 'The tested configuration uses four 96 GB GPUs.'
        gpus.append({'index': index, 'name': props.name, 'memory_bytes': props.total_memory})
    cache = Path('/cache')
    cache.mkdir(parents=True, exist_ok=True)
    probe = cache / '.write-check'
    probe.write_text('ok\n')
    probe.unlink()
    checkpoint = Path(os.environ['MODEL_PATH'])
    config = AutoConfig.from_pretrained(checkpoint, trust_remote_code=True)
    assert config.model_type == 'naive_n05_flash'
    tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
    assert tokenizer.encode('Hello')
    weight_map = json.loads((checkpoint / 'model.safetensors.index.json').read_text())['weight_map']
    shards = sorted(set(weight_map.values()))
    assert len(shards) == 49
    discovered = {}
    for name in shards:
        shard = checkpoint / name
        with shard.open('rb') as stream:
            header_size = struct.unpack('<Q', stream.read(8))[0]
            assert header_size < 100_000_000
            header = json.loads(stream.read(header_size))
        payload_end = 0
        for key, metadata in header.items():
            if key == '__metadata__':
                continue
            assert key not in discovered
            discovered[key] = name
            payload_end = max(payload_end, metadata['data_offsets'][1])
        assert shard.stat().st_size == 8 + header_size + payload_end, name
    assert discovered == weight_map, 'Shard headers and weight index differ.'
    print(json.dumps({'passed': True, 'versions': versions, 'gpus': gpus,
                      'shards': len(shards), 'tensors': len(discovered)}, indent=2))


if __name__ == '__main__':
    main()