cli/main.py

"""
Command-line interface for BackgroundFX Pro.
Integrates with existing app.py infrastructure.
"""

import click
import sys
import os
from pathlib import Path
from typing import Optional, Tuple
import logging
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
from rich.table import Table

# Import the existing application components from app.py
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Import from existing app.py
from app import (
    VideoProcessor,
    processor as app_processor,  # Use the global processor instance
    PROFESSIONAL_BACKGROUNDS,
    TWO_STAGE_AVAILABLE,
    CHROMA_PRESETS
)

console = Console()
logger = logging.getLogger(__name__)


@click.group()
@click.option('--verbose', '-v', is_flag=True, help='Verbose output')
@click.option('--debug', is_flag=True, help='Debug mode')
def cli(verbose: bool, debug: bool):
    """
    BackgroundFX Pro CLI - Professional video background replacement.
    
    Uses the same processing engine as the Gradio UI.
    """
    # Setup logging
    log_level = logging.DEBUG if debug else (logging.INFO if verbose else logging.WARNING)
    logging.basicConfig(
        level=log_level,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )


@cli.command()
@click.option('--force', is_flag=True, help='Force reload models')
def load_models(force: bool):
    """Load AI models for processing."""
    console.print("[bold blue]Loading models...[/bold blue]")
    
    def progress_callback(progress: float, message: str):
        console.print(f"  {int(progress*100)}% - {message}")
    
    # Use existing processor
    if force or not app_processor.models_loaded:
        result = app_processor.load_models(progress_callback)
        console.print(f"[green]✓[/green] {result}")
    else:
        console.print("[yellow]Models already loaded[/yellow]")


@cli.command()
@click.argument('input_video', type=click.Path(exists=True))
@click.argument('output_video', type=click.Path())
@click.option('--background', '-b', 
              type=click.Choice(list(PROFESSIONAL_BACKGROUNDS.keys()) + ['custom']),
              default='blur',
              help='Background type')
@click.option('--background-image', '-i', type=click.Path(exists=True),
              help='Custom background image (when using custom background)')
@click.option('--two-stage', is_flag=True,
              help='Use two-stage processing (cinema quality)')
@click.option('--chroma-preset', 
              type=click.Choice(list(CHROMA_PRESETS.keys()) if TWO_STAGE_AVAILABLE else ['standard']),
              default='standard',
              help='Chroma keying preset for two-stage')
@click.option('--preview-mask', is_flag=True,
              help='Generate mask preview video')
@click.option('--preview-greenscreen', is_flag=True,
              help='Generate greenscreen preview video')
def process(input_video: str, output_video: str, background: str,
           background_image: Optional[str], two_stage: bool,
           chroma_preset: str, preview_mask: bool, preview_greenscreen: bool):
    """Process a video file."""
    
    # Check if models are loaded
    if not app_processor.models_loaded:
        console.print("[yellow]Loading models first...[/yellow]")
        
        def progress_callback(progress: float, message: str):
            console.print(f"  {int(progress*100)}% - {message}")
        
        result = app_processor.load_models(progress_callback)
        console.print(f"[green]✓[/green] {result}")
    
    # Validate custom background
    if background == 'custom' and not background_image:
        console.print("[red]Error: Custom background requires --background-image[/red]")
        sys.exit(1)
    
    console.print(f"[bold blue]Processing video:[/bold blue] {input_video}")
    console.print(f"  Background: {background}")
    console.print(f"  Two-stage: {'Yes' if two_stage else 'No'}")
    
    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        console=console
    ) as progress:
        
        task = progress.add_task("Processing...", total=100)
        
        def progress_callback(value: float, message: str):
            progress.update(task, completed=int(value * 100), description=message)
        
        # Process video using existing processor
        result_path, message = app_processor.process_video(
            video_path=input_video,
            background_choice=background,
            custom_background_path=background_image if background == 'custom' else None,
            progress_callback=progress_callback,
            use_two_stage=two_stage,
            chroma_preset=chroma_preset,
            preview_mask=preview_mask,
            preview_greenscreen=preview_greenscreen
        )
    
    if result_path:
        # Move/copy result to desired output path
        import shutil
        shutil.move(result_path, output_video)
        
        console.print(f"[green]✓ Success![/green]")
        console.print(f"  Output: {output_video}")
        console.print(f"  {message}")
    else:
        console.print(f"[red]✗ Failed:[/red] {message}")
        sys.exit(1)


@cli.command()
def status():
    """Show system and model status."""
    status_info = app_processor.get_status()
    
    # Create status table
    table = Table(title="BackgroundFX Pro Status")
    table.add_column("Component", style="cyan")
    table.add_column("Status", style="green")
    
    table.add_row("Models Loaded", "✓" if status_info['models_loaded'] else "✗")
    table.add_row("Device", str(status_info['device']))
    table.add_row("Two-Stage Available", "✓" if status_info['two_stage_available'] else "✗")
    
    if 'memory_usage' in status_info:
        mem = status_info['memory_usage']
        table.add_row("Memory Usage", f"{mem['percent']:.1f}% ({mem['used_gb']:.1f}/{mem['total_gb']:.1f} GB)")
    
    if 'models' in status_info:
        models = status_info['models']
        table.add_row("SAM2 Predictor", "✓" if models.get('sam2_loaded') else "✗")
        table.add_row("MatAnyone", "✓" if models.get('matanyone_loaded') else "✗")
    
    console.print(table)


@cli.command()
def list_backgrounds():
    """List available background options."""
    table = Table(title="Available Backgrounds")
    table.add_column("ID", style="cyan")
    table.add_column("Description", style="white")
    table.add_column("Type", style="yellow")
    
    for bg_id, bg_info in PROFESSIONAL_BACKGROUNDS.items():
        table.add_row(
            bg_id,
            bg_info.get('description', 'Professional background'),
            bg_info.get('type', 'gradient')
        )
    
    table.add_row("custom", "Use your own image", "image")
    
    console.print(table)


@cli.command()
def cleanup():
    """Clean up resources and cache."""
    console.print("[bold blue]Cleaning up resources...[/bold blue]")
    
    app_processor.cleanup_resources()
    
    # Clean temporary files
    import tempfile
    import shutil
    temp_dir = Path(tempfile.gettempdir())
    
    patterns = ['processed_video_*.mp4', 'mask_preview_*.mp4', 'greenscreen_preview_*.mp4']
    removed = 0
    
    for pattern in patterns:
        for file in temp_dir.glob(pattern):
            try:
                file.unlink()
                removed += 1
            except:
                pass
    
    console.print(f"[green]✓[/green] Cleaned up {removed} temporary files")
    console.print("[green]✓[/green] Memory resources freed")


@cli.command()
@click.argument('input_dir', type=click.Path(exists=True))
@click.argument('output_dir', type=click.Path())
@click.option('--background', '-b', default='blur', help='Background type')
@click.option('--pattern', '-p', default='*.mp4', help='File pattern to match')
@click.option('--two-stage', is_flag=True, help='Use two-stage processing')
def batch(input_dir: str, output_dir: str, background: str, pattern: str, two_stage: bool):
    """Process multiple videos in batch."""
    input_path = Path(input_dir)
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Find videos
    videos = list(input_path.glob(pattern))
    
    if not videos:
        console.print(f"[yellow]No files matching '{pattern}' found in {input_dir}[/yellow]")
        return
    
    console.print(f"[bold blue]Found {len(videos)} videos to process[/bold blue]")
    
    # Ensure models are loaded
    if not app_processor.models_loaded:
        console.print("[yellow]Loading models...[/yellow]")
        app_processor.load_models()
    
    # Process each video
    success_count = 0
    
    for i, video_file in enumerate(videos, 1):
        console.print(f"\n[bold]Processing {i}/{len(videos)}:[/bold] {video_file.name}")
        
        output_file = output_path / f"processed_{video_file.name}"
        
        def progress_callback(value: float, message: str):
            console.print(f"  {int(value*100)}% - {message}", end='\r')
        
        result_path, message = app_processor.process_video(
            video_path=str(video_file),
            background_choice=background,
            custom_background_path=None,
            progress_callback=progress_callback,
            use_two_stage=two_stage,
            chroma_preset='standard'
        )
        
        if result_path:
            import shutil
            shutil.move(result_path, str(output_file))
            console.print(f"  [green]✓[/green] Saved to {output_file.name}")
            success_count += 1
        else:
            console.print(f"  [red]✗[/red] Failed: {message}")
    
    console.print(f"\n[bold]Batch complete:[/bold] {success_count}/{len(videos)} successful")


def main():
    """Main CLI entry point."""
    cli()


if __name__ == '__main__':
    main()