import argparse
import hashlib
import json
import os

import numpy as np
import pyarrow.parquet as pq
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
from transformers import PreTrainedTokenizerFast

DEFAULTS = {
    "dataset": "SlayerLab/polish-dynaword",
    "parquet": "data/european_hplt_v3_pl/european_hplt_v3_pl.parquet",
    "out": "data/",
    "vocab_size": 16384,
    "eos_token": "<|eos|>",
    "text_column": "text",
    "seed": 1337,
}


def load_config(path):
    cfg = dict(DEFAULTS)
    if path and os.path.isfile(path):
        with open(path) as f:
            raw = json.load(f)
        data = raw.get("data", {})
        tok = raw.get("tokenizer", {})
        if "dataset" in data:
            cfg["dataset"] = data["dataset"]
        if "parquet" in data:
            cfg["parquet"] = data["parquet"]
        if "out" in data:
            cfg["out"] = data["out"]
        if "vocab_size" in tok:
            cfg["vocab_size"] = tok["vocab_size"]
        if "eos_token" in tok:
            cfg["eos_token"] = tok["eos_token"]
    return cfg


def ensure_parquet(path, dataset, repo_root):
    if os.path.isfile(path):
        return path
    snapshot_download(
        repo_id=dataset,
        repo_type="dataset",
        allow_patterns=[path],
        local_dir=repo_root,
    )
    if not os.path.isfile(path):
        raise FileNotFoundError(path)
    return path


def iter_documents(path, column, limit):
    pf = pq.ParquetFile(path)
    seen = 0
    for batch in pf.iter_batches(batch_size=1000, columns=[column]):
        for doc in batch.column(0).to_pylist():
            if doc is None:
                continue
            yield doc
            seen += 1
            if limit is not None and seen >= limit:
                return


def train_tokenizer(tokenizer_path, docs, vocab_size, eos_token):
    tokenizer = Tokenizer(models.BPE())
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=[eos_token],
        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
        seed=DEFAULTS["seed"],
    )
    tokenizer.train_from_iterator(docs, trainer=trainer)
    tokenizer.save(tokenizer_path)


def encode_stream(path, docs, tokenizer, eos_id):
    num_tokens = 0
    buffer = np.empty(65536, dtype=np.uint16)
    filled = 0
    with open(path, "wb") as f:
        for batch in batch_documents(docs, 1000):
            encodings = tokenizer(batch, add_special_tokens=False)["input_ids"]
            for ids in encodings:
                ids.append(eos_id)
                if filled + len(ids) > buffer.size:
                    f.write(buffer[:filled].tobytes())
                    filled = 0
                    if len(ids) > buffer.size:
                        arr = np.asarray(ids, dtype=np.uint16)
                        f.write(arr.tobytes())
                        num_tokens += arr.size
                        continue
                buffer[filled : filled + len(ids)] = ids
                filled += len(ids)
                num_tokens += len(ids)
        if filled:
            f.write(buffer[:filled].tobytes())
    return num_tokens


def batch_documents(docs, size):
    batch = []
    for doc in docs:
        batch.append(doc)
        if len(batch) == size:
            yield batch
            batch = []
    if batch:
        yield batch


def sha256_of(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="config.json")
    parser.add_argument("--out", default=None)
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--dataset", default=None)
    parser.add_argument("--parquet", default=None)
    parser.add_argument("--text-column", default=None)
    args = parser.parse_args()

    cfg = load_config(args.config)
    if args.out is not None:
        cfg["out"] = args.out
    if args.dataset is not None:
        cfg["dataset"] = args.dataset
    if args.parquet is not None:
        cfg["parquet"] = args.parquet
    if args.text_column is not None:
        cfg["text_column"] = args.text_column

    repo_root = os.path.dirname(os.path.abspath(__file__))
    out_dir = cfg["out"]
    os.makedirs(out_dir, exist_ok=True)

    parquet_path = ensure_parquet(cfg["parquet"], cfg["dataset"], repo_root)
    tokenizer_path = os.path.join(out_dir, "tokenizer.json")
    tokens_path = os.path.join(out_dir, "tokens.bin")
    meta_path = os.path.join(out_dir, "meta.json")

    if not os.path.isfile(tokenizer_path):
        train_tokenizer(
            tokenizer_path,
            iter_documents(parquet_path, cfg["text_column"], args.limit),
            cfg["vocab_size"],
            cfg["eos_token"],
        )

    tokenizer = PreTrainedTokenizerFast(
        tokenizer_file=tokenizer_path, eos_token=cfg["eos_token"]
    )
    tokenizer.save_pretrained(out_dir)

    eos_id = tokenizer.convert_tokens_to_ids(cfg["eos_token"])
    num_tokens = encode_stream(
        tokens_path,
        iter_documents(parquet_path, cfg["text_column"], args.limit),
        tokenizer,
        eos_id,
    )

    meta = {
        "num_tokens": num_tokens,
        "eos_id": eos_id,
        "vocab_size": cfg["vocab_size"],
        "sha256": sha256_of(tokens_path),
    }
    with open(meta_path, "w") as f:
        json.dump(meta, f, indent=2)

    print(f"num_tokens {num_tokens}")
    print(f"sha256     {meta['sha256']}")


if __name__ == "__main__":
    main()