#!/usr/bin/env python3
"""
Prepare datasets for EthioBBPE tokenizer training.
Extracts text from parquet files and creates combined training corpus.
"""

import pandas as pd
from pathlib import Path
import argparse


def extract_text_from_parquet(parquet_path: str, text_columns: list = None) -> list:
    """
    Extract text from a parquet file.
    
    Args:
        parquet_path: Path to parquet file
        text_columns: List of column names to extract. If None, auto-detect text columns.
    
    Returns:
        List of text strings
    """
    df = pd.read_parquet(parquet_path)
    
    if text_columns is None:
        # Auto-detect text columns (exclude numeric columns)
        text_columns = [col for col in df.columns if df[col].dtype == 'object']
        print(f"Auto-detected text columns: {text_columns}")
    
    texts = []
    for col in text_columns:
        if col in df.columns:
            # Drop NaN values and convert to string
            col_texts = df[col].dropna().astype(str).tolist()
            texts.extend(col_texts)
            print(f"  - Column '{col}': {len(col_texts)} texts")
    
    return texts


def prepare_datasets(output_dir: str = "./data", output_filename: str = "combined_corpus.txt"):
    """
    Prepare synaxarium and canon biblical datasets for training.
    """
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    all_texts = []
    
    # Process Synaxarium dataset
    synaxarium_path = output_path / "synaxarium_dataset.parquet"
    if synaxarium_path.exists():
        print(f"\n=== Processing Synaxarium Dataset ===")
        print(f"File: {synaxarium_path}")
        texts = extract_text_from_parquet(str(synaxarium_path), text_columns=['መጽሃፍ'])
        all_texts.extend(texts)
        print(f"Total texts from Synaxarium: {len(texts)}")
    else:
        print(f"Warning: Synaxarium dataset not found at {synaxarium_path}")
    
    # Process Canon Biblical dataset
    canon_path = output_path / "canon_biblical_am_en.parquet"
    if canon_path.exists():
        print(f"\n=== Processing Canon Biblical Dataset ===")
        print(f"File: {canon_path}")
        # Extract both Amharic and English verses
        texts = extract_text_from_parquet(str(canon_path), text_columns=['verse', 'ጥቅስ'])
        all_texts.extend(texts)
        print(f"Total texts from Canon Biblical: {len(texts)}")
    else:
        print(f"Warning: Canon Biblical dataset not found at {canon_path}")
    
    if not all_texts:
        print("\nError: No texts extracted from datasets!")
        return None
    
    # Write combined corpus
    output_file = output_path / output_filename
    print(f"\n=== Writing Combined Corpus ===")
    print(f"Total texts: {len(all_texts)}")
    
    with open(output_file, 'w', encoding='utf-8') as f:
        for i, text in enumerate(all_texts):
            # Clean text: remove extra whitespace, ensure single line
            cleaned = ' '.join(text.split())
            if cleaned.strip():  # Only write non-empty lines
                f.write(cleaned + '\n')
    
    # Calculate stats
    file_size_mb = output_file.stat().st_size / (1024 * 1024)
    print(f"Output file: {output_file}")
    print(f"File size: {file_size_mb:.2f} MB")
    print(f"Total lines: {len(all_texts)}")
    
    # Sample preview
    print(f"\n=== Sample Preview (first 5 lines) ===")
    with open(output_file, 'r', encoding='utf-8') as f:
        for i, line in enumerate(f):
            if i >= 5:
                break
            preview = line.strip()[:100] + "..." if len(line.strip()) > 100 else line.strip()
            print(f"{i+1}: {preview}")
    
    return output_file


def main():
    parser = argparse.ArgumentParser(description="Prepare datasets for tokenizer training")
    parser.add_argument("--output_dir", type=str, default="./data", 
                       help="Directory containing parquet files and output location")
    parser.add_argument("--output_filename", type=str, default="combined_corpus.txt",
                       help="Name of output combined corpus file")
    
    args = parser.parse_args()
    
    output_file = prepare_datasets(
        output_dir=args.output_dir,
        output_filename=args.output_filename
    )
    
    if output_file:
        print(f"\n✓ Data preparation complete!")
        print(f"Ready for training with: python scripts/train_tokenizer.py --data_dir {args.output_dir}")
    else:
        print("\n✗ Data preparation failed!")
        exit(1)


if __name__ == "__main__":
    main()