#!/usr/bin/env python3
"""
Batch processing wrapper for Market-1501 dataset.
Processes multiple images in parallel to maximize GPU utilization.

Usage:
    # Persons-only mode (default, faster for person re-id datasets):
    python batch_process.py --input data/Market-1501-v15.09.15/bounding_box_test \
                            --batch-size 8 --workers 8
    
    # Full pipeline mode (all PII types - persons, license plates, traffic signs):
    python batch_process.py --input test_data/other \
                            --batch-size 4 --workers 4 --full-pipeline
"""
import argparse
from pathlib import Path
from queue import Queue, Empty
from threading import Thread, Lock
import subprocess
import sys
import time
import os
from typing import List, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed

def find_images(input_dir: Path, resume: bool = False) -> List[Path]:
    """
    Find all image files in the input directory.
    If resume=True, only return images that have "fail" in their log file.
    """
    extensions = ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG']
    images = []
    print("   Scanning for images", end='', flush=True)
    for ext in extensions:
        print('.', end='', flush=True)
        images.extend(input_dir.rglob(f'*{ext}'))
    print(' Done!', flush=True)
    
    if resume:
        print("   Checking log files for failures", end='', flush=True)
        # Check which images have "fail" in their log file
        # Log files mirror the input structure: artifacts/data/<input_path>/<filename>.log
        # e.g., input: data/detected/bounding_box_test/img.png
        #       log:   artifacts/data/detected/bounding_box_test/img.log
        artifacts_base = Path('artifacts/data')
        failed_images = []
        
        for img in images:
            try:
                # Get path relative to 'data/' directory
                # If input is data/detected/bounding_box_test, we want detected/bounding_box_test
                data_dir = Path('data')
                if data_dir in img.parents:
                    rel_to_data = img.relative_to(data_dir)
                    log_name = img.stem + '.log'
                    log_path = artifacts_base / rel_to_data.parent / log_name
                else:
                    # Fallback: can't determine relative path
                    failed_images.append(img)
                    continue
                
                # Check if log exists and contains "fail" (case-insensitive)
                if log_path.exists():
                    try:
                        with open(log_path, 'r') as f:
                            log_content = f.read()
                            if 'fail' in log_content.lower():
                                failed_images.append(img)
                                print('.', end='', flush=True)
                    except Exception as e:
                        # If we can't read the log, include the image to be safe
                        failed_images.append(img)
                else:
                    # If no log exists, the image needs processing
                    failed_images.append(img)
                    print('.', end='', flush=True)
            except (ValueError, Exception) as e:
                # If we can't compute relative path, include the image
                failed_images.append(img)
        
        print(' Done!', flush=True)
        print(f"   Found {len(images)} total images, {len(failed_images)} with failures in logs")
        return sorted(failed_images)
    
    return sorted(images)

def process_image(image_path: Path, gpu_id: int = 0, persons_only: bool = True) -> Tuple[Path, int, float]:
    """
    Process a single image using the main pipeline.
    Returns: (image_path, return_code, processing_time)
    
    Note: Output is NOT captured to avoid BrokenPipeError from subprocess pipes.
    The subprocess writes its own log file which can be consulted if needed.
    """
    start_time = time.time()
    
    # Choose pipeline module based on persons_only flag
    pipeline_module = 'mcp.clients.agents.main_persons_only' if persons_only else 'mcp.clients.agents.main'
    
    # Call the main pipeline with specific GPU
    cmd = [
        'conda', 'run', '-n', 'beyond-faces-and-license-plates',
        '--no-capture-output', 'python', '-m', pipeline_module,
        str(image_path)
    ]
    
    # Set CUDA_VISIBLE_DEVICES to use specific GPU
    env = os.environ.copy()
    env['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
    if persons_only:
        env['SKIP_PHASE2'] = '1'  # Ensure persons-only mode
    
    try:
        # Don't capture output to avoid pipe buffer issues (BrokenPipeError)
        # The subprocess writes its own log file at artifacts/{path}/{stem}.log
        # which contains all the detailed output
        result = subprocess.run(
            cmd,
            stdout=subprocess.DEVNULL,  # Discard stdout to prevent pipe issues
            stderr=subprocess.DEVNULL,  # Discard stderr to prevent pipe issues
            timeout=600,  # 10 minute timeout per image
            env=env
        )
        return_code = result.returncode
        
        # Determine expected output path based on input path
        # The pipeline uses: artifacts/{relative_parent}/{input_stem}{input_suffix}
        src = Path(image_path)
        try:
            src_absolute = src.resolve()
            relative_parent = src_absolute.parent.relative_to(Path.cwd())
        except ValueError:
            relative_parent = src.parent.name if src.parent.name else "."
        
        output_dir = Path("artifacts") / relative_parent
        expected_output = output_dir / f"{src.stem}{src.suffix if src.suffix else '.jpeg'}"
        
        if return_code == 0 and expected_output.exists():
            print(f"✓ {expected_output}")
        elif return_code == 0:
            print(f"⚠ {image_path} (completed but output not found at expected location)")
        else:
            print(f"✗ {image_path} (exit code: {return_code})", file=sys.stderr)
                
    except subprocess.TimeoutExpired:
        print(f"✗ TIMEOUT: {image_path}", file=sys.stderr)
        return_code = -1
    except Exception as e:
        print(f"✗ ERROR: {image_path}: {e}", file=sys.stderr)
        return_code = -1
    
    elapsed = time.time() - start_time
    return (image_path, return_code, elapsed)

def batch_process(
    input_dir: Path,
    batch_size: int = 8,
    num_workers: int = 8,
    gpu_id: int = 0,
    persons_only: bool = True,
    resume: bool = False
):
    """
    Process images in batches using multiple worker threads.
    
    Args:
        input_dir: Directory containing images
        batch_size: Number of images to process concurrently
        num_workers: Number of worker threads
        gpu_id: GPU device ID to use (default: 0)
        persons_only: If True, only process persons (skip Phase 2). Default: True
        resume: If True, skip already processed images. Default: False
    """
    # Find all images
    print(f"🔍 Scanning {input_dir}...")
    images = find_images(input_dir, resume=resume)
    total_images = len(images)
    
    if resume:
        print(f"🔄 Resume mode: Processing {total_images} remaining images")
    else:
        print(f"📸 Found {total_images} images")
    
    print(f"⚙️  Batch size: {batch_size}, Workers: {num_workers}")
    print(f"🎯 Pipeline mode: {'Persons-only' if persons_only else 'Full (all PII types)'}")
    print()
    
    if total_images == 0:
        print("No images found!")
        return
    
    # Monitor progress
    start_time = time.time()
    processed = 0
    failed = 0
    total_processing_time = 0.0
    progress_lock = Lock()
    
    def update_progress(result):
        """Update progress counters thread-safely."""
        nonlocal processed, failed, total_processing_time
        
        image_path, return_code, elapsed = result
        
        with progress_lock:
            total_processing_time += elapsed
            
            if return_code == 0:
                processed += 1
                status_icon = "✓"
            else:
                failed += 1
                status_icon = "✗"
            
            # Print progress with enhanced display
            elapsed_total = time.time() - start_time
            completed = processed + failed
            rate = completed / elapsed_total if elapsed_total > 0 else 0
            remaining = total_images - completed
            eta = remaining / rate if rate > 0 else 0
            
            # Calculate percentage
            percent = (100.0 * processed / total_images) if total_images > 0 else 0
            
            # Create progress bar
            bar_width = 30
            filled = int(percent * bar_width / 100)
            bar = '█' * filled + '░' * (bar_width - filled)
            
            # Format ETA
            eta_hours = int(eta // 3600)
            eta_mins = int((eta % 3600) // 60)
            eta_secs = int(eta % 60)
            if eta_hours > 0:
                eta_str = f"{eta_hours}h {eta_mins}m"
            elif eta_mins > 0:
                eta_str = f"{eta_mins}m {eta_secs}s"
            else:
                eta_str = f"{eta_secs}s"
            
            # Format elapsed time
            elapsed_hours = int(elapsed_total // 3600)
            elapsed_mins = int((elapsed_total % 3600) // 60)
            if elapsed_hours > 0:
                elapsed_str = f"{elapsed_hours}h {elapsed_mins}m"
            else:
                elapsed_str = f"{elapsed_mins}m"
            
            # Print compact progress line
            print(f"\r{status_icon} [{bar}] {percent:.1f}% | "
                  f"{processed}/{total_images} | "
                  f"Failed: {failed} | "
                  f"Rate: {rate:.2f} img/s | "
                  f"Elapsed: {elapsed_str} | "
                  f"ETA: {eta_str}     ", 
                  end='', flush=True)
    
    # Process images using ThreadPoolExecutor
    print(f"🚀 Starting parallel processing on GPU {gpu_id}...")
    print()
    try:
        with ThreadPoolExecutor(max_workers=num_workers) as executor:
            # Submit all tasks with GPU ID and persons_only flag
            futures = {executor.submit(process_image, img, gpu_id, persons_only): img for img in images}
            
            # Process results as they complete
            for future in as_completed(futures):
                try:
                    result = future.result()
                    update_progress(result)
                except Exception as e:
                    img = futures[future]
                    print(f"\n✗ ERROR processing {img}: {e}", file=sys.stderr)
                    with progress_lock:
                        failed += 1
                        
    except KeyboardInterrupt:
        print("\n\n⚠️  Interrupted by user. Cleaning up...")
        executor.shutdown(wait=False, cancel_futures=True)
    
    print("\n")
    
    # Print statistics
    total_time = time.time() - start_time
    hours = int(total_time // 3600)
    minutes = int((total_time % 3600) // 60)
    seconds = int(total_time % 60)
    
    avg_time = total_processing_time / processed if processed > 0 else 0
    
    print("=" * 60)
    print("📊 Final Statistics:")
    print(f"   Total images: {total_images}")
    print(f"   Processed: {processed}")
    print(f"   Failed: {failed}")
    print(f"   Total time: {hours}h {minutes}m {seconds}s")
    print(f"   Average time per image: {avg_time:.2f}s")
    print(f"   Throughput: {processed/total_time:.2f} images/second")
    print("=" * 60)

def batch_process_from_list(
    image_paths: List[Path],
    batch_size: int = 8,
    num_workers: int = 8, 
    gpu_id: int = 0,
    persons_only: bool = True,
    resume: bool = False
):
    """
    Process a specific list of images in parallel.
    
    Args:
        image_paths: List of image paths to process
        batch_size: Number of concurrent images to process
        num_workers: Number of parallel worker threads
        gpu_id: GPU device ID to use (default: 0)
        persons_only: If True, only process persons (skip Phase 2). Default: True
        resume: If True, skip already processed images. Default: False
    """
    total_images = len(image_paths)
    
    print(f"📸 Processing {total_images} images from list")
    print(f"⚙️  Batch size: {batch_size}, Workers: {num_workers}")
    print(f"🎯 Pipeline mode: {'Persons-only' if persons_only else 'Full (all PII types)'}")
    print()
    
    if total_images == 0:
        print("No images in list!")
        return
    
    # Monitor progress
    start_time = time.time()
    processed = 0
    failed = 0
    total_processing_time = 0.0
    progress_lock = Lock()
    
    def update_progress(result):
        """Update progress counters thread-safely."""
        nonlocal processed, failed, total_processing_time
        
        image_path, return_code, elapsed = result
        
        with progress_lock:
            total_processing_time += elapsed
            
            if return_code == 0:
                processed += 1
                status_icon = "✓"
            else:
                failed += 1
                status_icon = "✗"
            
            # Print progress with enhanced display
            elapsed_total = time.time() - start_time
            completed = processed + failed
            rate = completed / elapsed_total if elapsed_total > 0 else 0
            remaining = total_images - completed
            eta = remaining / rate if rate > 0 else 0
            
            # Calculate percentage
            percent = (100.0 * processed / total_images) if total_images > 0 else 0
            
            # Create progress bar
            bar_width = 30
            filled = int(percent * bar_width / 100)
            bar = '█' * filled + '░' * (bar_width - filled)
            
            # Format ETA
            eta_hours = int(eta // 3600)
            eta_mins = int((eta % 3600) // 60)
            eta_secs = int(eta % 60)
            if eta_hours > 0:
                eta_str = f"{eta_hours}h {eta_mins}m"
            elif eta_mins > 0:
                eta_str = f"{eta_mins}m {eta_secs}s"
            else:
                eta_str = f"{eta_secs}s"
            
            # Format elapsed time
            elapsed_hours = int(elapsed_total // 3600)
            elapsed_mins = int((elapsed_total % 3600) // 60)
            if elapsed_hours > 0:
                elapsed_str = f"{elapsed_hours}h {elapsed_mins}m"
            else:
                elapsed_str = f"{elapsed_mins}m"
            
            # Print compact progress line
            print(f"\r{status_icon} [{bar}] {percent:.1f}% | "
                  f"{processed}/{total_images} | "
                  f"Failed: {failed} | "
                  f"Rate: {rate:.2f} img/s | "
                  f"Elapsed: {elapsed_str} | "
                  f"ETA: {eta_str}     ", 
                  end='', flush=True)
    
    # Process images using ThreadPoolExecutor
    print(f"🚀 Starting parallel processing on GPU {gpu_id}...")
    print()
    try:
        with ThreadPoolExecutor(max_workers=num_workers) as executor:
            # Submit all tasks with GPU ID and persons_only flag
            futures = {executor.submit(process_image, img, gpu_id, persons_only): img for img in image_paths}
            
            # Process results as they complete
            for future in as_completed(futures):
                try:
                    result = future.result()
                    update_progress(result)
                except Exception as e:
                    img = futures[future]
                    print(f"\n✗ ERROR processing {img}: {e}", file=sys.stderr)
                    with progress_lock:
                        failed += 1
                        
    except KeyboardInterrupt:
        print("\n\n⚠️  Interrupted by user. Cleaning up...")
        executor.shutdown(wait=False, cancel_futures=True)
    
    print("\n")
    
    # Print statistics
    total_time = time.time() - start_time
    print("=" * 60)
    print("📊 BATCH PROCESSING SUMMARY")
    print("=" * 60)
    print(f"   Total images: {total_images}")
    print(f"   Successfully processed: {processed}")
    print(f"   Failed: {failed}")
    print(f"   Total wall time: {total_time:.1f}s ({total_time/60:.1f}m)")
    if processed > 0:
        print(f"   Average per image: {total_processing_time/processed:.1f}s")
    print(f"   Throughput: {processed/total_time:.2f} images/second")
    print("=" * 60)

def main():
    parser = argparse.ArgumentParser(
        description='Batch process images with parallel workers'
    )
    parser.add_argument(
        '--input',
        type=Path,
        required=False,
        help='Input directory containing images'
    )
    parser.add_argument(
        '--file-list',
        type=Path,
        required=False,
        help='Text file containing list of image paths (one per line)'
    )
    parser.add_argument(
        '--batch-size',
        type=int,
        default=8,
        help='Number of images to process concurrently (default: 8)'
    )
    parser.add_argument(
        '--workers',
        type=int,
        default=8,
        help='Number of worker threads (default: 8)'
    )
    parser.add_argument(
        '--gpu',
        type=int,
        default=0,
        help='GPU device ID to use (default: 0)'
    )
    parser.add_argument(
        '--full-pipeline',
        action='store_true',
        help='Use full pipeline (all PII types). Default is persons-only mode.'
    )
    parser.add_argument(
        '--resume',
        action='store_true',
        help='Resume mode: Skip already processed images and only process failed/remaining ones.'
    )
    
    args = parser.parse_args()
    
    # Validate arguments
    if not args.input and not args.file_list:
        print("Error: Must specify either --input or --file-list!")
        sys.exit(1)
    
    if args.input and args.file_list:
        print("Error: Cannot specify both --input and --file-list!")
        sys.exit(1)
    
    # Handle --file-list mode
    if args.file_list:
        if not args.file_list.exists():
            print(f"Error: File list {args.file_list} does not exist!")
            sys.exit(1)
        
        # Read file list
        print(f"📂 Reading file list from: {args.file_list}")
        with open(args.file_list, 'r') as f:
            image_paths = [Path(line.strip()) for line in f if line.strip()]
        
        # Validate all files exist
        missing = [p for p in image_paths if not p.exists()]
        if missing:
            print(f"⚠️  Warning: {len(missing)} files from list do not exist (will skip)")
            image_paths = [p for p in image_paths if p.exists()]
        
        if not image_paths:
            print("Error: No valid images found in file list!")
            sys.exit(1)
        
        print(f"📊 Processing {len(image_paths)} images from file list")
        
        # Process directly without calling find_images
        batch_process_from_list(
            image_paths=image_paths,
            batch_size=args.batch_size,
            num_workers=args.workers,
            gpu_id=args.gpu,
            persons_only=not args.full_pipeline,
            resume=args.resume
        )
    else:
        # Original --input mode
        if not args.input.exists():
            print(f"Error: Input directory {args.input} does not exist!")
            sys.exit(1)
        
        batch_process(
            input_dir=args.input,
            batch_size=args.batch_size,
            num_workers=args.workers,
            gpu_id=args.gpu,
            persons_only=not args.full_pipeline,  # Invert the flag: --full-pipeline means persons_only=False
            resume=args.resume
        )

if __name__ == '__main__':
    main()
