"""Local, reproducible model identities and safetensors inspection."""
import hashlib
import json
import math
import struct
from pathlib import Path

SOURCE = Path('/Volumes/ZX6 1TB/Qwen-Image-2.1/models/bf16')
REVISION = 'b3179ad355be050328e483a9dfdd9e60cd62adfa'
RUNTIME_REVISION = '95b01ccad2d9f65a9e87f6a87bd1c5df69626261'
NOTICE = ('Qwen is licensed under the Qwen RESEARCH LICENSE AGREEMENT, Copyright (c) 2026 '
          'Hangzhou Tongyi Laboratory Technology Co., Ltd. All Rights Reserved.')
MODIFICATION = ('Modified by ixim / iximbox for Image21-MLX: converted from the pinned '
                'BF16 source to MLX layout; eligible linear weights use groupwise affine '
                'quantization. See conversion.json for precision and exceptions. Built with Qwen.')

def write_json(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + '\n')

def sha256(path):
    with Path(path).open('rb') as f:
        return hashlib.file_digest(f, 'sha256').hexdigest()

def header(path):
    with Path(path).open('rb') as f:
        length = struct.unpack('<Q', f.read(8))[0]
        if length > 100_000_000:
            raise ValueError('Invalid safetensors header')
        return json.loads(f.read(length))

def tensors(root):
    for path in sorted(Path(root).glob('*.safetensors')):
        if path.name.startswith('._'):
            continue
        for name, info in header(path).items():
            if name != '__metadata__':
                yield path, name, info

def eligible(component, name, shape):
    # Preserve all visual encoder, embeddings, norms, modulation and boundary layers.
    if len(shape) != 2 or shape[-1] % 64 or not name.endswith('.weight'):
        return False
    if component == 'transformer':
        return name.startswith('transformer_blocks.') and ('.attn.to_' in name or '.img_mlp.' in name)
    if component == 'text_encoder':
        return 'language_model' in name and '.layers.' in name and ('.self_attn.' in name or '.mlp.' in name)
    return False

def inventory(root):
    result = {}
    for component in ('transformer', 'text_encoder', 'vae'):
        rows = list(tensors(Path(root) / component))
        params = sum(math.prod(v['shape']) for _, _, v in rows)
        qparams = sum(math.prod(v['shape']) for _, n, v in rows if eligible(component, n, v['shape']))
        nbytes = sum(v['data_offsets'][1] - v['data_offsets'][0] for _, _, v in rows)
        result[component] = dict(parameters=params, quantizable_parameters=qparams, tensor_bytes=nbytes,
                                 dtypes=sorted({v['dtype'] for _, _, v in rows}))
    return result