"""Measure an explicit two-component alias storage layout, not inference speed.

Payloads remain on the server under weights/ and must not be committed or cleaned
without the user's explicit deletion request. Manifest/byte/hash evidence may be archived.
"""
import argparse
import hashlib
import json
from pathlib import Path
import numpy as np
from pilot.tensors import load_tensors,decode
from .run import write_json


def payload(t,path):
    digest=hashlib.sha256()
    with path.open('wb') as f:
        for start in range(0,t.raw.size,262144):
            b=t.raw[start:start+262144].tobytes();f.write(b);digest.update(b)
            t.evict()
    return {'file':path.name,'dtype':t.dtype,'shape':list(t.shape),'bytes':path.stat().st_size,'sha256':digest.hexdigest()}


def main():
    p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);a=p.parse_args();root=a.root
    selected=json.loads((root/'selection.json').read_text())['selected']
    if not selected:
        write_json(root/'storage-evidence.json',{'measured':False,'reason':'No calibration candidate selected'});return
    c=selected[0]
    if c['kind']!='qk_norm':
        write_json(root/'storage-evidence.json',{'measured':False,'reason':'Only two-gain alias layout implemented in this bounded pilot; larger-module savings remain theoretical'});return
    tensors=load_tensors(root/'weights')
    prefix=f'model.language_model.layers.{c["i"]}.self_attn.'
    names=[prefix+'q_norm.weight',prefix+'k_norm.weight']
    # Original pair plus two alternative layouts; each is counted independently.
    path=root/'component-layout'/'weights';path.mkdir(parents=True,exist_ok=True)
    original={'format':'two_component_raw_v1','tensors':{}}
    for short,name in zip(('q','k'),names):original['tensors'][name]=payload(tensors[name],path/f'{short}.bin')
    original_bytes=json.dumps(original,separators=(',',':'),sort_keys=True).encode()
    (path/'original-layout.json').write_bytes(original_bytes)
    evidence=[]
    for target,donor in ((names[0],names[1]),(names[1],names[0])):
        entry=original['tensors'][donor]
        layout={'format':'two_component_raw_v1','tensors':{donor:entry},'aliases':{target:donor}}
        encoded=json.dumps(layout,separators=(',',':'),sort_keys=True).encode()
        file=path/('share-'+('k' if donor==names[1] else 'q')+'.json');file.write_bytes(encoded)
        t=tensors[donor]
        raw=np.fromfile(path/entry['file'],dtype=t.raw.dtype)
        reconstructed=decode(raw,t.dtype).reshape(t.shape)
        np.testing.assert_array_equal(reconstructed,t.read().reshape(t.shape));t.evict()
        before=sum(e['bytes'] for e in original['tensors'].values())+len(original_bytes)
        after=entry['bytes']+len(encoded)
        evidence.append({'target':target,'donor':donor,'original_payload_bytes':sum(e['bytes'] for e in original['tensors'].values()),
                         'shared_payload_bytes':entry['bytes'],'original_descriptor_bytes':len(original_bytes),
                         'shared_descriptor_bytes':len(encoded),'original_total_bytes':before,'shared_total_bytes':after,
                         'net_saved_bytes':before-after,'shared_descriptor_sha256':hashlib.sha256(encoded).hexdigest(),
                         'decoded_donor_exactly_matches_intervention_gain':True})
    write_json(root/'storage-evidence.json',{'measured':True,'candidate':c['id'],'layouts':evidence,
               'claim_limit':'Measured standalone component layout including descriptors, not an integrated compressed checkpoint, latency gain or full-model RSS reduction. Two sharing directions are alternatives and must not be added.'})


if __name__=='__main__':main()