import sys, time, json
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
import numpy as np
sys.path.insert(0, "/workspace/fractus-cte-atom-main")
from fractus.atom_tokenizer import AtomFractusTokenizer

OUT = Path("/workspace/atom_corpus")
OUT.mkdir(exist_ok=True)

def convert_one(path):
    tok = AtomFractusTokenizer()
    out_path = OUT / (Path(path).stem + ".i16")
    if out_path.exists() and out_path.stat().st_size > 0:
        return Path(path).name, 0, 0, True
    ids = []
    n = 0
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            text = "\n".join(m.get("content","") for m in obj.get("messages",[]) if isinstance(m, dict))
            if not text.strip():
                continue
            enc, _ = tok.encode_with_features(text)
            ids.extend(enc)
            n += 1
    if ids:
        arr = np.asarray(ids, dtype=np.int16)
        arr.tofile(out_path)
    return Path(path).name, n, len(ids), False

files = sorted(Path("/workspace/fractus-datasets").rglob("*.jsonl"))
print(f"files {len(files)}", flush=True)
total = 0
skipped = 0
t0 = time.time()
with ProcessPoolExecutor(max_workers=8) as ex:
    futs = {ex.submit(convert_one, str(p)): p for p in files}
    for i, fut in enumerate(as_completed(futs)):
        name, n, s, was_skip = fut.result()
        if was_skip:
            skipped += 1
        else:
            total += s
        if i % 5 == 0:
            print(f"{i}/{len(files)} {name} docs={n} ids={s} new={total:,} skipped={skipped} {total/max(time.time()-t0,1):.0f}/s", flush=True)
print(f"DONE new={total:,} skipped={skipped}", flush=True)