#!/usr/bin/env python
# coding=utf-8
# Copyright 2024 Statistics and Machine Learning Research Group. All rights reserved.
#
# This script is designed for CPU-only, offline dataset preprocessing.
# It uses the exact same argument parsing and tokenization logic as the main
# training pipeline to ensure that the generated cache is 100% compatible.

import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

# Add lmflow to path
sys.path.insert(0, os.getcwd())
from lmflow.args import DatasetArguments, ModelArguments
from lmflow.datasets.dataset import Dataset
from transformers import HfArgumentParser, AutoTokenizer

logger = logging.getLogger(__name__)

# A dummy TrainingArguments class to satisfy HfArgumentParser
@dataclass
class PreprocessingTrainingArgs:
    main_process_first: bool = field(
        default=True,
        metadata={"help": "Dummy argument for main_process_first context manager."}
    )
    
def main():
    print("[preprocess_main] Starting robust CPU dataset preprocessing script.")

    # --- Use HfArgumentParser to ensure 100% argument compatibility ---
    parser = HfArgumentParser((ModelArguments, DatasetArguments, PreprocessingTrainingArgs))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    print("\n[preprocess_main] Parsed ModelArguments:\n", model_args)
    print("\n[preprocess_main] Parsed DatasetArguments:\n", data_args)

    # --- Load Dataset ---
    print(f"\n[preprocess_main] Loading dataset from {data_args.dataset_path}")
    dataset = Dataset(data_args=data_args)
    print(f"[preprocess_main] Dataset loaded with {len(dataset)} examples.")

    # --- Load Tokenizer ---
    print(f"\n[preprocess_main] Loading tokenizer for '{model_args.tokenizer_name or model_args.model_name_or_path}'...")
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name or model_args.model_name_or_path,
        trust_remote_code=model_args.trust_remote_code,
        cache_dir=model_args.cache_dir,
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    print("[preprocess_main] Tokenizer loaded successfully.")

    # --- This is the core logic from LMFlow's tokenization pipeline ---
    from lmflow.tokenization.hf_decoder_model import conversation_tokenize_function
    
    print(f"\n[preprocess_main] Starting tokenization with num_proc={data_args.preprocessing_num_workers}...")
    print(f"  This will generate a cache in {data_args.dataset_cache_dir or '~/.cache/huggingface/datasets'}")

    # Manually handle the logic of main_process_first since we are not in a trainer
    # This ensures that only one process downloads/creates the dataset if it's not cached
    tokenized_dataset = dataset.map(
        conversation_tokenize_function,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        remove_columns=dataset.backend_dataset.column_names,
        load_from_cache_file=not data_args.overwrite_cache,
        desc="Running tokenizer on dataset",
        fn_kwargs={
            "data_args": data_args,
            "tokenizer": tokenizer,
            "conversation_template": data_args.conversation_template,
            "column_names": dataset.backend_dataset.column_names
        },
    )

    # Persist tokenized HF dataset to disk for cache-only training
    save_dir = data_args.dataset_cache_dir or os.path.join(os.getcwd(), ".cache", "lmf_preprocessed", "dataset")
    os.makedirs(save_dir, exist_ok=True)
    print(f"\n[preprocess_main] Saving tokenized dataset to disk at: {save_dir}")
    tokenized_dataset.get_backend_dataset().save_to_disk(save_dir)

    print("\n" + "="*50)
    print("      ✓✓✓ PREPROCESSING COMPLETE! ✓✓✓")
    print("="*50 + "\n")
    print(f"  Total tokenized examples: {len(tokenized_dataset)}")
    print(f"  Saved pre-tokenized dataset to: {save_dir}")
    print(f"  You can now run training with only --dataset_cache_dir={save_dir} (dataset_path not required).")

if __name__ == "__main__":
    main()
