import json, pathlib, subprocess, sys, time
root = pathlib.Path('/content/sub1quant')
log = root / 'colab_runs' / 'setup_model_inner.log'

def emit(msg):
    print(msg, flush=True)
    with log.open('a', encoding='utf-8') as f:
        f.write(msg + '\n')

def run(cmd):
    emit('$ ' + ' '.join(cmd))
    p = subprocess.run(cmd, text=True, capture_output=True)
    emit('returncode=' + str(p.returncode))
    if p.stdout:
        emit('stdout:\n' + p.stdout[-4000:])
    if p.stderr:
        emit('stderr:\n' + p.stderr[-4000:])
    if p.returncode != 0:
        raise SystemExit(p.returncode)

emit('SETUP START ' + time.strftime('%Y-%m-%d %H:%M:%S'))
run([sys.executable, '-m', 'pip', 'install', '-q', '-U', 'huggingface_hub', 'datasets', 'safetensors', 'accelerate', 'transformers'])

data_dir = root / 'data'
data_dir.mkdir(exist_ok=True)
wiki = data_dir / 'wiki.test.txt'
if not wiki.exists() or wiki.stat().st_size < 1000:
    emit('Downloading Salesforce/wikitext wikitext-2-raw-v1 test split')
    from datasets import load_dataset
    ds = load_dataset('Salesforce/wikitext', 'wikitext-2-raw-v1', split='test')
    wiki.write_text('\n\n'.join(ds['text']), encoding='utf-8')
    emit(f'wrote {wiki} bytes={wiki.stat().st_size}')
else:
    emit(f'wikitext exists bytes={wiki.stat().st_size}')

model_dir = root / 'models' / 'gemma-4-E2B'
model_dir.mkdir(parents=True, exist_ok=True)
if not (model_dir / 'model.safetensors').exists():
    emit('Downloading google/gemma-4-E2B-it snapshot')
    from huggingface_hub import snapshot_download
    snapshot_download(
        repo_id='google/gemma-4-E2B-it',
        local_dir=str(model_dir),
        allow_patterns=['*.json', '*.model', '*.safetensors', '*.md'],
    )
else:
    emit(f'model exists bytes={(model_dir / "model.safetensors").stat().st_size}')
summary = {
    'wiki_exists': wiki.exists(),
    'wiki_size': wiki.stat().st_size if wiki.exists() else None,
    'model_exists': (model_dir / 'model.safetensors').exists(),
    'model_size': (model_dir / 'model.safetensors').stat().st_size if (model_dir / 'model.safetensors').exists() else None,
}
(root / 'colab_runs' / 'setup_summary.json').write_text(json.dumps(summary, indent=2), encoding='utf-8')
emit('SETUP DONE ' + json.dumps(summary))