main.py

#!/usr/bin/env python3
"""
Simple command-line interface for Arabic text normalization.
Enhanced for massive dataset processing.
"""
import argparse
import sys
from pathlib import Path

from config import Config
from arabic_normalizer import ArabicNormalizer


def main():
    """Main command-line interface."""
    
    parser = argparse.ArgumentParser(
        description="Massive-scale Arabic text normalizer with LLM enhancement",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Basic usage
  python main.py input.txt output.txt
  
  # Scan directory and process all .txt files automatically
  python main.py --scan-dir "C:\\path\\to\\your\\txt\\files" --massive
  
  # Massive file processing with resume capability
  python main.py input.txt output.txt --massive --resume
  
  # Batch process directory (manual input/output)
  python main.py --batch input_dir output_dir --massive
  
  # Estimate processing time for huge file
  python main.py input.txt output.txt --estimate-only
  
  # Ultra-conservative for 3+ billion words
  python main.py input.txt output.txt --massive --ultra-safe
        """
    )
    
    # Input/output
    parser.add_argument("input", nargs="?", help="Input file or directory (optional when using --scan-dir)")
    parser.add_argument("output", nargs="?", help="Output file or directory (optional when using --scan-dir)")
    
    # Configuration
    parser.add_argument("--config", help="Configuration JSON file")
    parser.add_argument("--api-key", help="Gemini API key")
    
    # Processing modes
    parser.add_argument("--batch", action="store_true", 
                       help="Process directory in batch mode")
    parser.add_argument("--scan-dir", metavar="DIR", 
                       help="Scan directory for .txt files and process all (auto-creates output dir)")
    parser.add_argument("--massive", action="store_true",
                       help="Enable massive dataset processing with checkpoints")
    parser.add_argument("--resume", action="store_true",
                       help="Resume from checkpoint (only with --massive)")
    parser.add_argument("--ultra-safe", action="store_true",
                       help="Ultra-conservative settings for massive datasets")
    parser.add_argument("--estimate-only", action="store_true",
                       help="Only estimate processing time, don't process")
    
    # Processing options
    parser.add_argument("--no-llm", action="store_true",
                       help="Skip LLM processing (basic cleaning only)")
    
    # Parameters
    parser.add_argument("--min-length", type=int, default=20,
                       help="Minimum paragraph length (default: 20)")
    parser.add_argument("--max-length", type=int, default=1000,
                       help="Maximum paragraph length (default: 1000)")
    parser.add_argument("--chunk-size", type=int, default=500,
                       help="LLM chunk size (default: 500)")
    parser.add_argument("--requests-per-minute", type=int, default=5,
                       help="API requests per minute (default: 5)")
    parser.add_argument("--api-delay", type=float, default=10.0,
                       help="Delay between API calls in seconds (default: 10)")
    
    args = parser.parse_args()
    
    # Validate arguments
    if not args.scan_dir and (not args.input or not args.output):
        parser.error("input and output are required unless using --scan-dir")
    
    if args.scan_dir and (args.input or args.output):
        parser.error("Cannot use input/output arguments with --scan-dir")
    
    # Load configuration
    if args.config and Path(args.config).exists():
        config = Config.load_from_file(args.config)
        print(f"✓ Loaded config from: {args.config}")
    else:
        config = Config()
        print("✓ Using default configuration")
    
    # Apply ultra-safe settings for massive processing
    if args.ultra_safe:
        config.api_delay = 15.0
        config.requests_per_minute = 3
        config.chunk_size = 300
        config.max_tokens = 800
        config.max_retries = 8
        config.checkpoint_every = 50
        print("✓ Applied ultra-safe settings for massive datasets")
    
    # Apply command line overrides
    if args.min_length:
        config.min_paragraph_length = args.min_length
    if args.max_length:
        config.max_paragraph_length = args.max_length
    if args.chunk_size:
        config.chunk_size = args.chunk_size
    if args.requests_per_minute:
        config.requests_per_minute = args.requests_per_minute
    if args.api_delay:
        config.api_delay = args.api_delay
    
    # Initialize normalizer
    normalizer = ArabicNormalizer(config)
    
    # Estimate processing time if requested
    if args.estimate_only:
        if Path(args.input).is_file():
            estimate = normalizer.estimate_processing_time(args.input)
            print("\n=== Processing Time Estimate ===")
            print(f"File size: {estimate['file_size_gb']:.2f} GB")
            print(f"Estimated chunks: {estimate['estimated_chunks']:,}")
            print(f"Estimated time: {estimate['estimated_days']:.1f} days ({estimate['estimated_hours']:.1f} hours)")
            print(f"Processing rate: {estimate['requests_per_hour']:.1f} requests/hour")
            print("\nNote: This is a rough estimate. Actual time may vary based on:")
            print("- API response times and rate limiting")
            print("- Text complexity and chunk sizes")
            print("- Network conditions and retries")
        else:
            print("✗ Cannot estimate: input must be a file")
        sys.exit(0)
    
    # Setup LLM if not disabled
    if not args.no_llm:
        try:
            normalizer.setup_llm(args.api_key)
            print("✓ LLM setup successful for massive processing")
        except Exception as e:
            print(f"✗ LLM setup failed: {e}")
            print("Continuing with basic preprocessing only...")
            args.no_llm = True
    
    # Process files
    try:
        if args.scan_dir:
            # Directory scanning mode - process all .txt files
            input_dir = Path(args.scan_dir)
            if not input_dir.exists() or not input_dir.is_dir():
                print(f"✗ Directory not found: {args.scan_dir}")
                sys.exit(1)
            
            # Create output directory
            output_dir = input_dir / "normalized_output"
            output_dir.mkdir(exist_ok=True)
            
            # Find all .txt files
            txt_files = list(input_dir.glob("*.txt"))
            if not txt_files:
                print(f"✗ No .txt files found in: {args.scan_dir}")
                sys.exit(1)
            
            print(f"✓ Found {len(txt_files)} .txt files to process")
            print(f"✓ Output directory: {output_dir}")
            
            # Process each file
            successful = 0
            for txt_file in txt_files:
                output_file = output_dir / f"{txt_file.stem}_normalized.txt"
                print(f"\n📄 Processing: {txt_file.name}")
                
                try:
                    if args.massive:
                        success = normalizer.process_file_massive(str(txt_file), str(output_file), args.resume)
                    else:
                        success = normalizer.process_file(str(txt_file), str(output_file))
                    
                    if success:
                        print(f"✓ Completed: {txt_file.name} → {output_file.name}")
                        successful += 1
                    else:
                        print(f"✗ Failed: {txt_file.name}")
                        
                except Exception as e:
                    print(f"✗ Error processing {txt_file.name}: {e}")
            
            print(f"\n🎉 Processing complete: {successful}/{len(txt_files)} files successful")
            
        elif args.batch:
            # Batch processing
            if args.massive:
                count = normalizer.process_batch_massive(args.input, args.output)
            else:
                count = normalizer.process_batch(args.input, args.output)
                
            if count > 0:
                print(f"✓ Successfully processed {count} files")
                sys.exit(0)
            else:
                print("✗ No files were processed successfully")
                sys.exit(1)
        else:
            # Single file processing
            if args.massive:
                success = normalizer.process_file_massive(args.input, args.output, args.resume)
            else:
                success = normalizer.process_file(args.input, args.output)
                
            if success:
                print(f"✓ Successfully processed: {args.input} → {args.output}")
                
                # Show final stats for massive processing
                if args.massive:
                    stats = normalizer.llm.get_stats()
                    print(f"✓ API calls made: {stats['total_requests']}")
                    print(f"✓ Success rate: {stats['success_rate']:.1f}%")
                
                sys.exit(0)
            else:
                print(f"✗ Processing failed: {args.input}")
                sys.exit(1)
                
    except KeyboardInterrupt:
        print("\n✗ Processing interrupted by user")
        print("✓ Progress saved in checkpoints/ directory (use --resume to continue)")
        sys.exit(1)
    except Exception as e:
        print(f"✗ Fatal error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()