import argparse
import hashlib
import json
from pathlib import Path
from transformers import AutoTokenizer

p = argparse.ArgumentParser()
p.add_argument('--weights', required=True)
p.add_argument('--root', required=True)
a = p.parse_args()
root = Path(a.root)
tokenizer = AutoTokenizer.from_pretrained(a.weights, local_files_only=True)
records = []
provenance = {'dataset': 'Salesforce/wikitext', 'config': 'wikitext-2-raw-v1', 'licenses': ['cc-by-sa-3.0', 'gfdl'], 'sequence_length': 256, 'selection': 'First 100 source rows per official split, contiguous non-overlapping token windows; no chat template or added special tokens', 'sources': {}}
for source_split, split, count in [('train', 'calibration', 16), ('validation', 'validation', 8), ('test', 'heldout', 8)]:
    source = root / (source_split + '-rows.json')
    raw = source.read_bytes()
    rows = json.loads(raw)['rows']
    text = '\n'.join(row['row']['text'] for row in rows)
    tokens = tokenizer.encode(text, add_special_tokens=False)
    if len(tokens) < count * 256:
        raise ValueError(f'{source_split}: insufficient tokens: {len(tokens)}')
    provenance['sources'][source_split] = {'sha256': hashlib.sha256(raw).hexdigest(), 'rows': len(rows), 'available_tokens': len(tokens), 'used_tokens': count * 256}
    for index in range(count):
        records.append({'id': f'{split}-{index:03d}', 'split': split, 'token_ids': tokens[index * 256:(index + 1) * 256]})
(root / 'corpus.jsonl').write_text(''.join(json.dumps(record) + '\n' for record in records))
smoke = dict(records[0], id='smoke', token_ids=records[0]['token_ids'][:32])
(root / 'smoke.jsonl').write_text(json.dumps(smoke) + '\n')
provenance['corpus_sha256'] = hashlib.sha256((root / 'corpus.jsonl').read_bytes()).hexdigest()
(root / 'corpus-provenance.json').write_text(json.dumps(provenance, indent=2) + '\n')
print(json.dumps(provenance, indent=2))