"""Reference BF16 reconstruction of source-bound AXQ NVFP4 diffusion storage."""
import hashlib
import importlib.util
import sys
import json
from pathlib import Path, PurePosixPath

def digest(path):
    value=hashlib.sha256()
    with path.open('rb') as stream:
        while data:=stream.read(8*1024*1024): value.update(data)
    return value.hexdigest()

def verify(directory):
    root=Path(directory).resolve()
    manifest=json.loads((root/'axq-diffusion-manifest.json').read_text())
    schema_path=Path(__file__).with_name('axq_diffusion_schema.py')
    module_spec=importlib.util.spec_from_file_location('axq_pack_diffusion_schema',schema_path)
    schema=importlib.util.module_from_spec(module_spec);sys.modules[module_spec.name]=schema;module_spec.loader.exec_module(schema)
    schema.DiffusionManifest.model_validate(manifest)
    if manifest['schema_version']!='axq-diffusion.reference/1' or manifest['format']!='nvfp4':
        raise ValueError('Unsupported AXQ diffusion storage')
    for name,entry in manifest['files'].items():
        relative=PurePosixPath(name)
        if relative.is_absolute() or '..' in relative.parts: raise ValueError('Unsafe manifest member')
        path=root/name
        if not path.is_file() or path.is_symlink() or not path.resolve().is_relative_to(root):
            raise ValueError('Use a regular-file snapshot_download local_dir; symlinks are not accepted')
        if path.stat().st_size!=entry['size_bytes'] or digest(path)!=entry['sha256']:
            raise ValueError('Pack digest mismatch: '+name)
    return root,manifest

def load_transformer(directory, component='transformer', *, verify_files=True):
    import torch
    from safetensors.torch import load_file
    from axquant.nvfp4 import Nvfp4Tensor,dequantize_nvfp4
    import diffusers
    from accelerate import init_empty_weights
    root=Path(directory).resolve()
    root,manifest=verify(root) if verify_files else (root,json.loads((root/'axq-diffusion-manifest.json').read_text()))
    if component not in manifest['components']: raise ValueError('Component is not packaged')
    spec=manifest['components'][component]
    allowed={'QwenImage21Transformer2DModel','MiniMaxH3Transformer3DModel'}
    if spec['class_name'] not in allowed: raise ValueError('Unsupported transformer class')
    config=json.loads((root/component/'config.json').read_text())
    cls=getattr(diffusers,spec['class_name'])
    with init_empty_weights(include_buffers=False): model=cls.from_config(config)
    expected=set(model.state_dict())
    if set(spec['tensors'])!=expected: raise ValueError('Model state does not match pinned component inventory')
    state={}
    for filename in sorted(set(x['file'] for x in spec['tensors'].values())):
        stored=load_file(str(root/filename),device='cpu')
        consumed=set()
        for name,entry in spec['tensors'].items():
            if entry['file']!=filename: continue
            if entry['method']=='preserve':
                value=stored[name];consumed.add(name)
            elif entry['method']=='nvfp4':
                keys=[name+'.axq_packed',name+'.axq_scales']
                consumed.update(keys)
                shape=tuple(entry['shape'])
                value=torch.empty(shape,dtype=torch.bfloat16)
                for start in range(0,shape[0],128):
                    stop=min(start+128,shape[0])
                    decoded=dequantize_nvfp4(Nvfp4Tensor(stored[keys[0]][start:stop].numpy(),stored[keys[1]][start:stop].numpy(),entry['global_scale'],(stop-start,shape[1]),'numpy-reference'))
                    value[start:stop]=torch.from_numpy(decoded).to(torch.bfloat16)
            else: raise ValueError('Unknown tensor method')
            if list(value.shape)!=entry['shape']: raise ValueError('Shape mismatch')
            if str(value.dtype).removeprefix('torch.')!=entry['dtype']: raise ValueError('Unexpected tensor dtype')
            state[name]=value
        if consumed!=set(stored): raise ValueError('Unconsumed packed tensor keys')
    model.load_state_dict(state,strict=True,assign=True)
    del state
    model.eval()
    if any(x.is_meta for x in model.parameters()) or any(x.is_meta for x in model.buffers()): raise ValueError('Model has unloaded tensors')
    return model

def load_qwen_image_pipeline(directory):
    import torch
    from diffusers import QwenImage21Pipeline
    root,manifest=verify(directory)
    if manifest['source_model']!='Qwen/Qwen-Image-2.1-Turbo': raise ValueError('Wrong pipeline family')
    model=load_transformer(root,verify_files=False)
    return QwenImage21Pipeline.from_pretrained(str(root),transformer=model,torch_dtype=torch.bfloat16,local_files_only=True)