"""Compose pinned Q4 and complete MLP Q3 payloads without dense expansion."""
from pathlib import Path
import argparse, hashlib, json, os, shutil, struct, sys
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
import native_model_export as native

Q4_MANIFEST='a068c0532c24b208d95de2b02af4533701f21bdb2436936178a35c8e6d219f45'
Q4_WEIGHTS='478a045ad7f447d36205d411c93cf5e5847e0879d950c9ee7dcfc4cf8d289694'

def main():
    p=argparse.ArgumentParser();p.add_argument('--bundle',type=Path,required=True);p.add_argument('--overlay',type=Path,required=True);p.add_argument('--overlay-sha256',required=True);p.add_argument('--assets',type=Path,required=True);p.add_argument('--output',type=Path,required=True);a=p.parse_args()
    assert native.sha256(a.bundle/'manifest.json')==Q4_MANIFEST
    assert native.sha256(a.bundle/'weights.bin')==Q4_WEIGHTS
    assert native.sha256(a.overlay/'manifest.json')==a.overlay_sha256
    parent=json.loads((a.bundle/'manifest.json').read_text());overlay=json.loads((a.overlay/'manifest.json').read_text())
    assert overlay['status']=='complete' and overlay['format']=='audio8-q3-mlp-overlay-v1' and overlay['identity']['training_only']
    targets={f'language_model.model.layers.{i}.mlp.{n}.weight' for i in range(36) for n in ('gate_proj','up_proj','down_proj')}
    assert set(overlay['tensors'])==targets and overlay['complete_target_count']==108
    assert parent['tie_word_embeddings'] and parent['aliases']=={'language_model.lm_head.weight':'language_model.model.embed_tokens.weight'}
    expected=native.schema();records={r['name']:dict(r) for r in parent['tensors']}
    assert set(records)==set(expected) and len(parent['tensors'])==len(expected)
    sources={};hashes={}
    for name,r in records.items():
        if name in targets:
            q=overlay['tensors'][name];path=a.overlay/q['file']
            assert q['file']==name+'.q3f8' and q['precision']=='q3_full8' and q['shape']==r['shape']
            assert q['grid']==dict(bits=3,zero_point=4,max_code=7,padding_code=4)
            assert path.stat().st_size==q['bytes'] and native.sha256(path)==q['sha256']
            r.update(precision='q3_full8',offset=0,group_size=64,codes_bytes=q['codes_bytes'],scales_bytes=q['bytes']-q['codes_bytes'],scales_offset=q['codes_bytes'],scales_dtype='F16',bytes=q['bytes'])
            sources[name]=(path,0);hashes[name]=q['sha256']
        else:
            sources[name]=(a.bundle/'weights.bin',r['offset']);hashes[name]=r['provenance']['payload_sha256']
    metadata={n:native.record_metadata(r,expected[n]) for n,r in records.items()}
    total=sum(x[2] for x in metadata.values());assert total==1882263104
    for name in native.ASSETS:
        assert native.sha256(a.assets/name)==parent['external_assets_not_included'][name]['sha256']
    assert shutil.disk_usage(a.output.parent).free>total+1024**3
    a.output.mkdir(exist_ok=False);out=a.output/'model.a8m';digest=hashlib.sha256();actual={}
    with out.open('xb') as dst:
        def write(data):dst.write(data);digest.update(data)
        write(native.HEADER.pack(native.MAGIC,len(expected),total))
        for name,(shape,_) in expected.items():
            kind,grid,size=metadata[name];encoded=name.encode('ascii')
            write(native.RECORD.pack(len(encoded),kind,len(shape),*grid,size));write(struct.pack('<'+'I'*len(shape),*shape));write(encoded)
            path,offset=sources[name];h=hashlib.sha256()
            with path.open('rb') as src:
                src.seek(offset);remaining=size
                while remaining:
                    block=src.read(min(native.CHUNK,remaining));assert block
                    write(block);h.update(block);remaining-=len(block)
            assert h.hexdigest()==hashes[name],name
            actual[name]=h.hexdigest()
        dst.write(digest.digest());dst.flush();os.fsync(dst.fileno())
    readback=native.verify_stream(out)
    assert readback['stream_body_sha256']==digest.hexdigest() and readback['records']==938
    assert native.sha256(a.bundle/'weights.bin')==Q4_WEIGHTS and native.sha256(a.overlay/'manifest.json')==a.overlay_sha256
    for name in native.ASSETS:
        shutil.copyfile(a.assets/name,a.output/name)
        assert native.sha256(a.output/name)==parent['external_assets_not_included'][name]['sha256']
    report=dict(format='A8MOD001',status='complete',production_ready=False,accuracy_validated=False,source_manifest_sha256=Q4_MANIFEST,source_weights_sha256=Q4_WEIGHTS,overlay_manifest_sha256=a.overlay_sha256,exporter_sha256=native.sha256(Path(__file__)),helper_sha256=native.sha256(Path(native.__file__)),model_sha256=native.sha256(out),model_bytes=out.stat().st_size,payload_bytes=total,weight_bytes_saved=304349184,tensor_sha256=actual,readback=readback,aliases=parent['aliases'],claim_scope='All 938 tensor payloads independently hash verified; no inference or accuracy acceptance.')
    (a.output/'manifest.json').write_text(json.dumps(report,indent=2)+'\n');print(json.dumps({k:v for k,v in report.items() if k not in ('tensor_sha256','readback')}))

if __name__=='__main__':main()