#!/usr/bin/env python3
"""Package FSI_Edge source for Colab upload."""
import zipfile, os

SRC = '/FSI_Edge'
FILES = [
    'src/model.py', 'src/data.py',
    'training/train.py',
    'export/export_gguf.py',
    'data/prepare_data.py', 'data/generate_coldstart.py',
]
DEST = '/FSI_Edge/fsi_edge_colab.zip'

with zipfile.ZipFile(DEST, 'w', zipfile.ZIP_DEFLATED) as z:
    for f in FILES:
        z.write(os.path.join(SRC, f), arcname=f'FSI_Edge/{f}')
    # Also include the tokenizer files
    tok_dir = os.path.join(SRC, 'fsi_edge_tokenizer')
    if os.path.exists(tok_dir):
        for root, _, files in os.walk(tok_dir):
            for fn in files:
                full = os.path.join(root, fn)
                arcname = os.path.relpath(full, SRC)
                z.write(full, arcname=f'FSI_Edge/{arcname}')

print(f'Created {DEST} ({os.path.getsize(DEST)/1024:.0f} KB)')
print('Upload to Colab: files.upload() or drag to /content/')