"""Train a byte-level BPE and build audited, streaming 1024-token LM data."""
from __future__ import annotations
import argparse
from collections import Counter
import json
import os
os.environ.setdefault('RAYON_NUM_THREADS', '4')
import random
from pathlib import Path
import time
import numpy as np
import pyarrow.parquet as pq
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
from common import *

CHAT_TEMPLATE = "{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"

def rows(path, seed=SEED):
    pf = pq.ParquetFile(path)
    groups = list(range(pf.num_row_groups))
    random.Random(seed).shuffle(groups)
    for group in groups:
        row_offset = sum(pf.metadata.row_group(i).num_rows for i in range(group))
        for batch in pf.iter_batches(batch_size=128, row_groups=[group],
                                     columns=['source','kind','prompt','think','answer']):
            for i, row in enumerate(batch.to_pylist()):
                yield row_offset + i, row
            row_offset += batch.num_rows

def train_tokenizer(source, manifest, out, sample_bytes):
    from transformers import PreTrainedTokenizerFast
    counts = Counter()
    byte_counts = Counter()
    def texts():
        # Each source file gets the same byte allowance, avoiding a tokenizer
        # dominated by the two largest distillation files.
        per_file = max(65536, sample_bytes // len(manifest['files']))
        for item in manifest['files']:
            for _, row in rows(source / item['name']):
                if split_for(row['prompt']) != 'train': continue
                text = render_row(row)
                if text is None: continue
                yield text
                counts[item['name']] += 1
                byte_counts[item['name']] += len(text.encode('utf-8'))
                if byte_counts[item['name']] >= per_file: break
            print('TOKENIZER_SAMPLE', item['name'], counts[item['name']], flush=True)
    tokenizer = Tokenizer(models.BPE())
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True)
    tokenizer.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(vocab_size=VOCAB, min_frequency=2,
                                 special_tokens=SPECIALS,
                                 initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
                                 show_progress=False)
    tokenizer.train_from_iterator(texts(), trainer=trainer)
    if tokenizer.get_vocab_size() < VOCAB:
        tokenizer.add_tokens([f'<|reserved_{i}|>' for i in range(VOCAB-tokenizer.get_vocab_size())])
    assert tokenizer.get_vocab_size() == VOCAB
    fast = PreTrainedTokenizerFast(tokenizer_object=tokenizer,
                                  pad_token=SPECIALS[0], bos_token=SPECIALS[1], eos_token=SPECIALS[2],
                                  additional_special_tokens=SPECIALS[3:], model_max_length=CONTEXT,
                                  clean_up_tokenization_spaces=False)
    fast.chat_template = CHAT_TEMPLATE
    out.mkdir(parents=True, exist_ok=True)
    fast.save_pretrained(out)
    report = tokenizer_checks(fast)
    report.update({'training_split_only': True, 'sample_rows_by_file': dict(counts),
                   'sample_utf8_bytes_by_file': dict(byte_counts),
                   'normalizer': None, 'algorithm': 'byte-level BPE',
                   'tokenizer_sha256': sha(out/'tokenizer.json')})
    dump(out/'tokenizer_audit.json', report)
    print('TOKENIZER_VERIFIED', report['vocab_size'], flush=True)
    return fast

def tokenizer_checks(fast):
    samples = ['Hello, world!\n  Two spaces stay.',
               'x² + y² = π; √2 ≈ 1.414; café — naïve 🚀',
               'def add(a, b):\n    return a + b\n',
               r'\frac{13}{6} = 2.1667; $x \\ y$',
               '<|im_start|>assistant\n<think>\nCheck.\n</think>\n42<|im_end|>']
    for text in samples:
        ids = fast.encode(text, add_special_tokens=False)
        assert fast.decode(ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) == text
        assert all(0 <= t < VOCAB for t in ids)
    special_ids = {s: fast.encode(s, add_special_tokens=False) for s in SPECIALS}
    assert special_ids == {s:[i] for i,s in enumerate(SPECIALS)}, special_ids
    assert fast.model_max_length == CONTEXT
    return {'vocab_size': len(fast), 'special_token_ids': special_ids,
            'round_trip_cases_passed': len(samples), 'context': CONTEXT,
            'no_unknown_token': fast.unk_token_id is None}

class TokenWriter:
    # 512 MB per shard: below the requested 7 GB ceiling.
    def __init__(self, directory, split, shard_tokens=256_000_000):
        self.directory = Path(directory); self.directory.mkdir(parents=True, exist_ok=True)
        self.split = split; self.shard_tokens = shard_tokens
        self.real_tokens = 0; self.total_tokens = 0; self.current_tokens = 0
        self.files = []; self.handle = None; self.path = None
    def add(self, ids, padding=False):
        ids = np.asarray(ids, dtype='<u2')
        if ids.size: assert int(ids.max()) < VOCAB
        if not padding: self.real_tokens += len(ids)
        pos = 0
        while pos < len(ids):
            if self.handle is None:
                self.path = self.directory/f'{self.split}-{len(self.files):04d}.bin'
                self.handle = self.path.open('wb'); self.current_tokens = 0
            n = min(len(ids)-pos, self.shard_tokens-self.current_tokens)
            self.handle.write(ids[pos:pos+n].tobytes())
            self.current_tokens += n; self.total_tokens += n; pos += n
            if self.current_tokens == self.shard_tokens: self.close_shard()
    def close_shard(self):
        if self.handle is None: return
        self.handle.close(); self.handle = None
        self.files.append({'name': self.path.name, 'tokens': self.current_tokens,
                           'bytes': self.path.stat().st_size, 'sha256': sha(self.path)})
    def finish(self):
        if self.total_tokens:
            # Each sample uses 1024 inputs and the next 1024 target tokens.
            # Pad only the final tail; padding targets are ignored by the loss.
            self.add([0] * ((1-self.total_tokens) % CONTEXT), padding=True)
        self.close_shard()
        return {'real_tokens': self.real_tokens, 'stored_tokens': self.total_tokens,
                'padding_tokens': self.total_tokens-self.real_tokens,
                'blocks': max(0,(self.total_tokens-1)//CONTEXT), 'files': self.files}

def prepare(source, manifest_path, output, train_tokens, sample_bytes, verify=True, tokenizer_path=None):
    source = Path(source); output = Path(output); output.mkdir(parents=True, exist_ok=True)
    manifest = json.loads(Path(manifest_path).read_text(encoding='utf-8'))
    assert manifest['total_bytes'] < MAX_SOURCE_BYTES
    assert all(f['bytes'] < MAX_FILE_BYTES for f in manifest['files'])
    for item in manifest['files']:
        path = source/item['name']
        assert path.stat().st_size == item['bytes'], path
        if verify:
            assert sha(path) == item['sha256'], path
            print('SOURCE_HASH_OK', item['name'], flush=True)
    dump(output/'source_manifest.json', manifest)
    if tokenizer_path:
        import shutil
        from transformers import AutoTokenizer
        shutil.copytree(tokenizer_path,output/'tokenizer',dirs_exist_ok=True)
        fast=AutoTokenizer.from_pretrained(output/'tokenizer',local_files_only=True)
        tokenizer_checks(fast)
    else:
        fast = train_tokenizer(source, manifest, output/'tokenizer', sample_bytes)
    tokenizer = fast.backend_tokenizer
    writers = {s:TokenWriter(output, s) for s in ['train','validation','test']}
    selected = Counter(); kinds = Counter(); source_tokens = Counter()
    sample_audit = []; duplicates = 0; scanned = 0
    # Store compact hashes, not full documents. This set is bounded by selected records.
    seen = set()
    items = manifest['files'].copy(); random.Random(SEED).shuffle(items)
    quota_basis = sum(x['bytes'] for x in items)
    started = time.time()
    for item in items:
        quota = max(CONTEXT*4, round(train_tokens * item['bytes']/quota_basis))
        file_train = 0; pending = []; pending_meta = []
        def flush():
            nonlocal file_train, duplicates
            if not pending: return
            encodings = tokenizer.encode_batch(pending, add_special_tokens=False)
            for text, encoding, meta in zip(pending, encodings, pending_meta):
                idx, row, split = meta
                digest = __import__('hashlib').sha256(text.encode('utf-8')).digest()
                if digest in seen:
                    duplicates += 1; continue
                seen.add(digest)
                ids = encoding.ids
                writers[split].add(ids)
                selected[split] += 1; kinds[row['kind']] += 1
                source_tokens[item['name'] + '/' + split] += len(ids)
                if split == 'train': file_train += len(ids)
                if len(sample_audit) < 100 or split != 'train':
                    # No source text in audit: positions and hashes identify records.
                    sample_audit.append({'file':item['name'], 'row':idx, 'split':split,
                                         'prompt_sha256':prompt_key(row['prompt']).hex(),
                                         'record_sha256':digest.hex(), 'tokens':len(ids)})
            pending.clear(); pending_meta.clear()
        for idx, row in rows(source/item['name']):
            scanned += 1
            text = render_row(row)
            if text is None: continue
            split = split_for(row['prompt'])
            pending.append(text); pending_meta.append((idx,row,split))
            if len(pending) >= 64:
                flush()
                if file_train >= quota: break
        flush()
        print('TOKENIZED_FILE', item['name'], 'train_tokens', file_train,
              'total_train', writers['train'].real_tokens, 'seconds', round(time.time()-started), flush=True)
    splits = {s:w.finish() for s,w in writers.items()}
    assert all(splits[s]['blocks'] > 0 for s in splits), 'Each split needs data'
    assert sum(x['bytes'] for s in splits.values() for x in s['files']) < MAX_SOURCE_BYTES
    assert all(x['bytes'] < MAX_FILE_BYTES for s in splits.values() for x in s['files'])
    result = {'created_utc':now(), 'seed':SEED, 'dtype':'little-endian uint16',
              'context':CONTEXT, 'vocab_size':VOCAB,
              'source_manifest_sha256':sha(manifest_path),
              'tokenizer_sha256':sha(output/'tokenizer'/'tokenizer.json'),
              'requested_train_tokens':train_tokens,
              'sampling':'Deterministically shuffled row groups; per-file train-token quota proportional to compressed file bytes; retain whole records including final answers. Quotas may overshoot by one encoding batch.',
              'split_rule':'SHA256(NFKC/casefold/whitespace-normalized prompt) modulo 1000: validation 0..4; test 5..9; train 10..999. Split before tokenizer fitting and packing.',
              'packing':'Separate continuous token streams per split; EOS separates records; long documents continue into later windows. Attention resets every 1024 tokens; cross-document attention within a window is allowed.',
              'loss':'All next-token targets except pad=0; no assistant-only loss at scratch-pretraining stage.',
              'rows_scanned':scanned, 'selected_rows':dict(selected), 'exact_duplicate_rows_removed':duplicates,
              'source_tokens':dict(source_tokens), 'kinds':dict(kinds), 'splits':splits}
    dump(output/'split_record_audit.json', sample_audit)
    dump(output/'corpus_manifest.json', result)
    model_config().save_pretrained(output/'model_config')
    print('PREPARE_COMPLETE', json.dumps({k:v['real_tokens'] for k,v in splits.items()}), flush=True)

if __name__ == '__main__':
    ap=argparse.ArgumentParser()
    ap.add_argument('--source',required=True); ap.add_argument('--manifest',required=True)
    ap.add_argument('--out',required=True)
    ap.add_argument('--train-tokens',type=int,default=1_520_000_000)
    ap.add_argument('--tokenizer-sample-bytes',type=int,default=128_000_000)
    ap.add_argument('--tokenizer')
    args=ap.parse_args()
    prepare(args.source,args.manifest,args.out,args.train_tokens,args.tokenizer_sample_bytes,tokenizer_path=args.tokenizer)