"""Export only inference weights and public project code; never upload credentials or remote config."""
import argparse
import hashlib
import json
from pathlib import Path
import shutil
import torch
from safetensors.torch import save_file
from tinyquery.model import Config,TinyQuery


def main():
    p=argparse.ArgumentParser(); p.add_argument('--checkpoint',required=True); p.add_argument('--data',required=True)
    p.add_argument('--out',required=True); args=p.parse_args()
    source=Path(args.checkpoint); data=Path(args.data); out=Path(args.out); out.mkdir(parents=True,exist_ok=True)
    if source.suffix=='.safetensors':
        config=json.loads((source.parent/'config.json').read_text())
        shutil.copy2(source,out/'model.safetensors')
        info=source.parent/'checkpoint-info.json'
        if not info.exists():info=source.parent/'best-info.json'
        if info.exists(): shutil.copy2(info,out/'checkpoint-info.json')
    else:
        checkpoint=torch.load(source,map_location='cpu',weights_only=False)
        config=checkpoint['config']
        state={k:v.to(torch.bfloat16).contiguous() for k,v in checkpoint['model'].items()}
        save_file(state,str(out/'model.safetensors'))
        info={k:checkpoint[k] for k in ['step','processed_tokens','response_tokens','training_seconds','random_initialization']}
        (out/'checkpoint-info.json').write_text(json.dumps(info,indent=2))
    (out/'config.json').write_text(json.dumps(config,indent=2))
    shutil.copy2(data/'tokenizer.json',out/'tokenizer.json')
    package=Path(__file__).parent
    shutil.copytree(package,out/'tinyquery',dirs_exist_ok=True,ignore=shutil.ignore_patterns('__pycache__','*.pyc'))
    shutil.copy2(package/'requirements-inference.txt',out/'requirements.txt')
    for name in ['tokenization.json','grounding-stats.json','split-audit.json']:
        if (data/name).exists(): shutil.copy2(data/name,out/name)
    model=TinyQuery(Config(**config)); count=sum(p.numel() for p in model.parameters())
    from safetensors.torch import load_file
    model.load_state_dict(load_file(str(out/'model.safetensors')),strict=True)
    manifest={'parameters':count,'format':'Custom native PyTorch; see tinyquery/model.py, not a Transformers AutoModel checkpoint',
              'random_initialization':True,'files':{}}
    for file in sorted(out.rglob('*')):
        if file.is_file() and file.name!='manifest.json':
            digest=hashlib.sha256()
            with file.open('rb') as stream:
                for block in iter(lambda:stream.read(4*1024*1024),b''): digest.update(block)
            manifest['files'][str(file.relative_to(out))]={'bytes':file.stat().st_size,'sha256':digest.hexdigest()}
    (out/'manifest.json').write_text(json.dumps(manifest,indent=2)); print(json.dumps({'parameters':count,'files':len(manifest['files'])}))


if __name__=='__main__': main()