#!/usr/bin/env python3
"""Check a Swift 1.5 4/5-bit download and optional pinned Python runtime.

Does not download files, load weights, generate text, or certify GUI compatibility.
"""
import argparse
import hashlib
import json
from pathlib import Path
import struct
import sys

RELEASES = {4: {
    'model-00001-of-00003.safetensors': (5328325554, 'd91b70ca84dff315addeeb8a599dce2beccfec33c7483686ede2e8ff3cc459dc'),
    'model-00002-of-00003.safetensors': (5354185158, 'eddcff1a6ef0971990f7cd01dba77ab9a8c20d59821bb4cd7e25234fe3809346'),
    'model-00003-of-00003.safetensors': (5144253923, '644cc5322ff706c9608b42ea1ff83e0beb12b828d49b36c468a366f44593b8eb'),
}, 5: {
    'model-00001-of-00004.safetensors': (5360417859, '9da8750595ec42d8c5d2ff0ae0b695e387cbcb3c0beac40f262e6d9a7196b59d'),
    'model-00002-of-00004.safetensors': (5338181777, 'b592723bd4e9f01344335cd7d4b4f41a7c8cde98e2e074ad1d943cdb951a6835'),
    'model-00003-of-00004.safetensors': (5355925620, '0027e7fe05310d5b2d960414221ed0872de9b511aa7e6b1c3a8f0e4647558412'),
    'model-00004-of-00004.safetensors': (3227578201, 'afec577f16805b327471561d6841662eb4acf7d866ad5e761c5ef8cd3149d8f4'),
}}
EXPECTED_TENSORS = 2379
MODEL_HASHES = {
    4: {'f559e4559946bfad4ce20b93906028449b76e007f42f36bb83d2c2c5f6cb95a9',
        'ecda666d6d5f1059388a2ce4267d9ff8625398bb925d6c0840d0059ce17b676f'},
    5: {'ecda666d6d5f1059388a2ce4267d9ff8625398bb925d6c0840d0059ce17b676f'},
}
REQUIRED = (
    'config.json', 'model.safetensors.index.json', 'tokenizer.json',
    'tokenizer_config.json', 'chat_template.jinja', 'generation_config.json',
    'vocab.json', 'merges.txt', 'preprocessor_config.json', 'video_preprocessor_config.json',
)

def unique(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError('Duplicate JSON key')
        result[key] = value
    return result

def read_json(path):
    return json.loads(path.read_text(encoding='utf-8'), object_pairs_hook=unique)

def runtime_check(config):
    """Check dispatch and the audited architecture source, without model creation."""
    import inspect
    from importlib.metadata import version
    from mlx_lm.utils import _get_classes
    cls, args = _get_classes(config)
    args.from_dict(config)
    row = {'mlx': version('mlx'), 'mlx_lm': version('mlx-lm'),
           'selected_model_class': cls.__module__+'.'+cls.__name__}
    if cls.__module__ != 'mlx_lm.models.qwen3_5_full':
        raise ValueError('The required full-model MLX-LM patch is not active. Follow USAGE.md in this Python environment.')
    row['model_source_sha256'] = hashlib.sha256(Path(inspect.getsourcefile(cls)).read_bytes()).hexdigest()
    if row['model_source_sha256'] not in MODEL_HASHES[config['quantization']['bits']]:
        raise ValueError('This is not the pinned patched model implementation. For 5-bit, both patches are required. Follow USAGE.md.')
    return row

def check(folder, check_hashes=False, check_runtime=False):
    folder = Path(folder)
    errors = []
    report = {'files':[], 'weights_loaded':False, 'generation_tested':False,
              'scope':'File sizes, headers, index and requested checks only; no full-model or GUI test'}
    try:
        config = read_json(folder/'config.json')
        bits = config['quantization']['bits']
        if type(bits) is not int or bits not in RELEASES:
            raise ValueError('Expected 4-bit or 5-bit Swift release')
        expected = RELEASES[bits]
    except (OSError,ValueError,KeyError,TypeError) as exc:
        return dict(report,status='FAIL',errors=[f'Cannot identify the release: {exc}'])
    report.update(bits=bits,expected_shards=len(expected),expected_weight_bytes=sum(v[0] for v in expected.values()))
    for name in REQUIRED:
        path = folder/name
        if not path.is_file():
            errors.append(f'Missing required asset: {name}')
        elif path.stat().st_size == 0:
            errors.append(f'Empty required asset: {name}')
        elif name.endswith('.json'):
            try:
                read_json(path)
            except (ValueError, OSError) as exc:
                errors.append(f'Cannot parse {name}: {type(exc).__name__}')
    seen = {}
    sizes = {'U32':4, 'BF16':2, 'F16':2, 'F32':4}
    for name,(size,sha256) in expected.items():
        path = folder/name
        actual = path.stat().st_size if path.is_file() else None
        row = {'file':name,'expected_bytes':size,'actual_bytes':actual}
        report['files'].append(row)
        if actual != size:
            errors.append(f'{name}: expected {size} bytes, found {actual if actual is not None else "missing"}')
            continue
        try:
            with path.open('rb') as stream:
                length = struct.unpack('<Q',stream.read(8))[0]
                if not 0 < length <= 16*1024*1024:
                    raise ValueError('Invalid safetensors header size')
                headers = json.loads(stream.read(length),object_pairs_hook=unique)
            if not isinstance(headers,dict):
                raise ValueError('Invalid safetensors header object')
            headers.pop('__metadata__',None)
            ranges = []
            for tensor,entry in headers.items():
                if tensor in seen:
                    raise ValueError('Duplicate tensor across shards')
                start,end = entry['data_offsets']
                shape = entry['shape']
                if not isinstance(shape,list) or type(start) is not int or type(end) is not int:
                    raise ValueError('Invalid shape or offset types')
                count = 1
                for dim in shape:
                    if type(dim) is not int or dim < 0:
                        raise ValueError('Invalid tensor shape')
                    count *= dim
                if entry['dtype'] not in sizes or end-start != count*sizes[entry['dtype']]:
                    raise ValueError('Invalid tensor size or dtype')
                if not 0 <= start <= end <= actual-8-length:
                    raise ValueError('Invalid tensor data offsets')
                seen[tensor] = name
                ranges.append((start,end))
            offset = 0
            for start,end in sorted(ranges):
                if start != offset:
                    raise ValueError('Overlapping tensor data or a payload gap')
                offset = end
            if offset != actual-8-length:
                raise ValueError('Tensor payload does not fill shard')
            row['header_check'] = 'PASS'
            if check_hashes:
                digest = hashlib.sha256()
                with path.open('rb') as stream:
                    for block in iter(lambda:stream.read(8*1024*1024),b''):
                        digest.update(block)
                row['sha256'] = digest.hexdigest()
                row['sha256_check'] = 'PASS' if row['sha256'] == sha256 else 'FAIL'
                if row['sha256'] != sha256:
                    errors.append(f'SHA256 mismatch: {name}')
        except (ValueError,KeyError,TypeError,OSError,struct.error) as exc:
            errors.append(f'{name}: {exc}')
    extra = sorted(p.name for p in folder.glob('*.safetensors') if p.name not in expected)
    if extra:
        errors.append(f'Unexpected root weight files: {extra}')
    try:
        index = read_json(folder/'model.safetensors.index.json')
        if not isinstance(index['weight_map'],dict):
            raise ValueError('Invalid weight map')
        if set(index['weight_map'].values()) != set(expected):
            errors.append(f'Index must reference exactly the {len(expected)} expected shards')
        if seen != index['weight_map']:
            errors.append('Available tensor headers do not match the complete index')
        if len(index['weight_map']) != EXPECTED_TENSORS:
            errors.append(f'Expected {EXPECTED_TENSORS} saved tensor entries')
    except (OSError,ValueError,KeyError,TypeError):
        errors.append('Cannot validate the checkpoint index')
    try:
        if config.get('quantization') != {'bits':bits,'group_size':64,'mode':'affine'}:
            errors.append(f'Expected affine / {bits}-bit / group size 64')
        if config.get('language_model_only') is not False:
            errors.append('Expected the complete Swift architecture configuration')
        if check_runtime:
            try:
                report['runtime'] = runtime_check(config)
            except Exception as exc:
                errors.append(f'Runtime check failed: {type(exc).__name__}: {exc}')
    except (OSError,ValueError,KeyError,TypeError):
        errors.append('Cannot validate model configuration')
    report['hash_check'] = ('PASS' if all(row.get('sha256_check')=='PASS' for row in report['files']) else 'FAIL_OR_INCOMPLETE') if check_hashes else 'not run; add --hash to read and SHA256-check all weight bytes'
    report['runtime_check_requested'] = check_runtime
    report['errors'] = errors
    report['status'] = 'FAIL' if errors else 'PASS_DOWNLOAD_AND_REQUESTED_CHECKS_ONLY'
    return report

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('folder',type=Path)
    parser.add_argument('--hash',action='store_true',help='Stream-read every shard and compare its original SHA256')
    parser.add_argument('--runtime',action='store_true',help='Check the installed Python MLX-LM dispatcher; does not instantiate the model')
    args = parser.parse_args()
    result = check(args.folder,args.hash,args.runtime)
    print(json.dumps(result,indent=2))
    return 1 if result['errors'] else 0

if __name__ == '__main__':
    sys.exit(main())