"""
BoccaccioAI - Merge tokenized binary data

Concatenates new tokenized data (e.g. Wikipedia) to the existing
pre-training binary files (train.bin, val.bin) and updates metadata.

Usage:
    python -m src.data.merge_tokenized \
        --existing-dir data/tokenized/pretrain \
        --new-dir data/tokenized/wiki \
        --output-dir data/tokenized/pretrain_v2

De Lauretis Tech
"""

from __future__ import annotations

import argparse
import json
import logging
import shutil
from pathlib import Path

import numpy as np

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("boccaccio.data.merge_tokenized")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Merge tokenized binary datasets for BoccaccioAI pre-training."
    )
    parser.add_argument(
        "--existing-dir",
        type=str,
        required=True,
        help="Directory with existing tokenized data (train.bin, val.bin, meta.json).",
    )
    parser.add_argument(
        "--new-dir",
        type=str,
        required=True,
        help="Directory with new tokenized data to append.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        required=True,
        help="Output directory for merged data.",
    )
    parser.add_argument(
        "--sequence-length",
        type=int,
        default=2048,
        help="Sequence length for alignment (default: 2048).",
    )
    return parser.parse_args()


def load_bin(path: Path) -> np.ndarray:
    """Load a binary file as uint16 array."""
    return np.memmap(path, dtype=np.uint16, mode="r")


def main() -> None:
    args = parse_args()
    existing_dir = Path(args.existing_dir)
    new_dir = Path(args.new_dir)
    output_dir = Path(args.output_dir)
    seq_len = args.sequence_length

    logger.info("=== BoccaccioAI - Merge Tokenized Data ===")
    logger.info("Existing: %s", existing_dir)
    logger.info("New:      %s", new_dir)
    logger.info("Output:   %s", output_dir)

    # Validate inputs
    existing_train = existing_dir / "train.bin"
    existing_val = existing_dir / "val.bin"
    new_train = new_dir / "train.bin"
    new_val = new_dir / "val.bin"

    for f in [existing_train, existing_val, new_train, new_val]:
        if not f.exists():
            raise FileNotFoundError(f"Required file not found: {f}")

    # Load existing data
    logger.info("Loading existing train.bin ...")
    old_train = load_bin(existing_train)
    logger.info("  %s tokens", f"{len(old_train):,}")

    logger.info("Loading existing val.bin ...")
    old_val = load_bin(existing_val)
    logger.info("  %s tokens", f"{len(old_val):,}")

    # Load new data
    logger.info("Loading new train.bin ...")
    add_train = load_bin(new_train)
    logger.info("  %s tokens", f"{len(add_train):,}")

    logger.info("Loading new val.bin ...")
    add_val = load_bin(new_val)
    logger.info("  %s tokens", f"{len(add_val):,}")

    # Concatenate
    merged_train_tokens = len(old_train) + len(add_train)
    merged_val_tokens = len(old_val) + len(add_val)

    # Align to sequence length
    merged_train_tokens = (merged_train_tokens // seq_len) * seq_len
    merged_val_tokens = (merged_val_tokens // seq_len) * seq_len

    logger.info("Merged train: %s tokens (%d sequences)",
                f"{merged_train_tokens:,}", merged_train_tokens // seq_len)
    logger.info("Merged val:   %s tokens (%d sequences)",
                f"{merged_val_tokens:,}", merged_val_tokens // seq_len)

    # Write output
    output_dir.mkdir(parents=True, exist_ok=True)

    # Train
    logger.info("Writing merged train.bin ...")
    out_train = np.memmap(
        output_dir / "train.bin", dtype=np.uint16, mode="w+",
        shape=(merged_train_tokens,)
    )
    old_n = min(len(old_train), merged_train_tokens)
    out_train[:old_n] = old_train[:old_n]
    remaining = merged_train_tokens - old_n
    if remaining > 0:
        out_train[old_n:old_n + remaining] = add_train[:remaining]
    out_train.flush()
    del out_train

    # Val
    logger.info("Writing merged val.bin ...")
    out_val = np.memmap(
        output_dir / "val.bin", dtype=np.uint16, mode="w+",
        shape=(merged_val_tokens,)
    )
    old_vn = min(len(old_val), merged_val_tokens)
    out_val[:old_vn] = old_val[:old_vn]
    remaining_v = merged_val_tokens - old_vn
    if remaining_v > 0:
        out_val[old_vn:old_vn + remaining_v] = add_val[:remaining_v]
    out_val.flush()
    del out_val

    # Load existing metadata if available
    total_docs = 0
    existing_meta = existing_dir / "meta.json"
    new_meta = new_dir / "meta.json"
    if existing_meta.exists():
        with open(existing_meta) as f:
            total_docs += json.load(f).get("total_documents", 0)
    if new_meta.exists():
        with open(new_meta) as f:
            total_docs += json.load(f).get("total_documents", 0)

    # Write metadata
    meta = {
        "vocab_size": 32000,
        "train_tokens": merged_train_tokens,
        "val_tokens": merged_val_tokens,
        "sequence_length": seq_len,
        "total_documents": total_docs,
        "sources": [str(existing_dir), str(new_dir)],
    }
    with open(output_dir / "meta.json", "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)

    logger.info("=== Merge complete ===")
    logger.info("Total train tokens: %s", f"{merged_train_tokens:,}")
    logger.info("Total val tokens:   %s", f"{merged_val_tokens:,}")
    logger.info("Output:             %s", output_dir)


if __name__ == "__main__":
    main()