import os
import argparse
import time
from datasets import load_dataset
# HuggingFace Tokenizer to bootstrap vocab and merges
from transformers import AutoTokenizer

# The native Memory-Safe Rust Kernel
import sail_kernel

def main():
    parser = argparse.ArgumentParser(description="SAIL Fast Data Preprocessor (Native Rust Kernel)")
    parser.add_argument("--dataset", type=str, default="roneneldan/TinyStories", help="HF Dataset name")
    parser.add_argument("--split", type=str, default="train", help="Dataset split")
    parser.add_argument("--text-col", type=str, default="text", help="Text column name")
    parser.add_argument("--output-dir", type=str, default="sail/tokenized_cache", help="Output directory for shards")
    parser.add_argument("--batch-size", type=int, default=10000, help="Number of records per shard")
    parser.add_argument("--max-length", type=int, default=2048, help="Max tokens per sequence")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    
    # 1. Load HuggingFace Tokenizer to get vocab and merges
    print("[1/3] Loading Tokenizer configuration...")
    hf_tok = AutoTokenizer.from_pretrained("sail/sail_hf_model", trust_remote_code=True)
    vocab = hf_tok.get_vocab()
    
    # BPE models have merges. We pull them out. If not a BPE model, returns empty list.
    merges = []
    if hasattr(hf_tok, "bpe_ranks"):
        merges = [tuple(pair) for pair in hf_tok.bpe_ranks.keys()]
        
    print(f"      Vocab size: {len(vocab)}")
    print(f"      Merges: {len(merges)}")

    pad_id = vocab.get("<PAD>", hf_tok.pad_token_id or 0)
    unk_id = vocab.get("<UNK>", hf_tok.unk_token_id or 1)
    bos_id = vocab.get("<SOS>", hf_tok.bos_token_id or 2)
    eos_id = vocab.get("<EOS>", hf_tok.eos_token_id or 3)

    print("[2/3] Initializing Rust Memory-Safe Tokenizer Kernel...")
    # Initialize the high-speed Rust-based Tokenizer
    rust_tok = sail_kernel.RustTokenizer(
        vocab=vocab,
        merges=merges,
        unk_id=unk_id,
        pad_id=pad_id,
        eos_id=eos_id,
        bos_id=bos_id
    )

    print(f"[3/3] Streaming dataset '{args.dataset}' and compressing natively...")
    ds = load_dataset(args.dataset, split=args.split, streaming=True)
    
    batch_texts = []
    shard_count = 0
    total_records = 0
    start_time = time.time()

    def parse_item_to_text(item, text_field):
        # 1. Handle Chat/Message format (e.g., ultrachat_200k)
        if "messages" in item and isinstance(item["messages"], list):
            formatted_text = ""
            for msg in item["messages"]:
                role = msg.get('role', 'user').capitalize()
                content = msg.get('content', '')
                formatted_text += f"{role}: {content}\n"
            return formatted_text

        # 2. Handle Orca format (system, question, response)
        if "system_prompt" in item and "question" in item and "response" in item:
            return f"### System:\n{item['system_prompt']}\n\n### Question:\n{item['question']}\n\n### Response:\n{item['response']}"

        # 3. Handle Alpaca format (instruction, input, output)
        if "instruction" in item and "output" in item:
            text = f"### Instruction:\n{item['instruction']}"
            if item.get("input"):
                text += f"\n\n### Input:\n{item['input']}"
            text += f"\n\n### Response:\n{item['output']}"
            return text

        # 4. Fallback: Search for generic text fields
        if text_field in item:
            return str(item[text_field])
        
        for field in ["text", "instruction", "content", "response", "output"]:
            if field in item:
                return str(item[field])
                
        return ""

    for item in ds:
        text = parse_item_to_text(item, args.text_col)
        if not text: continue
        batch_texts.append(text)
        
        if len(batch_texts) >= args.batch_size:
            # 1. Rust Native: Clean Texts (parallel, GIL released)
            clean_texts = sail_kernel.clean_batch(batch_texts)
            
            # 2. Rust Native: Tokenize batch (parallel, GIL released)
            token_ids_batch = rust_tok.encode_batch(clean_texts, True, args.max_length)
            
            # 3. Rust Native: Write Binary Shard (zero Python objects)
            shard_path = os.path.join(args.output_dir, f"shard_{shard_count:05d}.bin")
            sail_kernel.write_arrow_shard(
                token_ids_batch, 
                shard_path, 
                pad_id=pad_id, 
                max_length=args.max_length
            )
            
            # 4. Rust Native: Zstd Compress
            zstd_path = f"{shard_path}.zst"
            sail_kernel.compress_shard(shard_path, zstd_path, 3)
            os.remove(shard_path)  # Keep only compressed
            
            shard_count += 1
            total_records += len(batch_texts)
            elapsed = time.time() - start_time
            print(f"  -> Wrote Shard {shard_count:05d} ({args.batch_size} records). Rate: {total_records/elapsed:.1f} rec/s")
            
            batch_texts.clear()
            
            # For demonstration, stop after 5 shards
            if shard_count >= 5:
                break

    # Process remaining
    if batch_texts:
        clean_texts = sail_kernel.clean_batch(batch_texts)
        token_ids_batch = rust_tok.encode_batch(clean_texts, True, args.max_length)
        shard_path = os.path.join(args.output_dir, f"shard_{shard_count:05d}.bin")
        sail_kernel.write_arrow_shard(token_ids_batch, shard_path, pad_id=pad_id, max_length=args.max_length)
        zstd_path = f"{shard_path}.zst"
        sail_kernel.compress_shard(shard_path, zstd_path, 3)
        os.remove(shard_path)
        total_records += len(batch_texts)
        shard_count += 1

    elapsed = time.time() - start_time
    print(f"\n[DONE] Successfully generated {shard_count} compressed shards.")
    print(f"Total Records: {total_records}")
    print(f"Time Taken: {elapsed:.2f} seconds ({total_records/elapsed:.1f} rec/s)")

if __name__ == "__main__":
    main()