#!/usr/bin/env python3
"""
RemoveBG Tool - Command Line & Graphical Interface for Image & Video Background Removal
Powered by rembg and FFmpeg.
"""

import os
import sys
import argparse
from typing import List

# Add parent directory to path if running directly
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from core.processor import BackgroundRemoverEngine, AVAILABLE_MODELS
from core.ffmpeg_utils import check_ffmpeg_installed

def run_cli(args):
    """Execute command-line processing mode."""
    print("=" * 60)
    print("✨ RemoveBG CLI Processor")
    print("=" * 60)

    model_name = args.model
    if model_name not in AVAILABLE_MODELS:
        print(f"Warning: Unknown model '{model_name}'. Available options: {list(AVAILABLE_MODELS.keys())}")
        model_name = "u2net"

    print(f"Loading Model: {model_name} ({AVAILABLE_MODELS.get(model_name, '')})")
    engine = BackgroundRemoverEngine(model_name=model_name)

    bg_color = args.bg_color
    if bg_color:
        if bg_color.lower() == "green":
            bg_color = "#00FF00"
        elif bg_color.lower() == "white":
            bg_color = "#FFFFFF"
        elif bg_color.lower() == "black":
            bg_color = "#000000"

    # 1. Video Processing Mode
    if args.video:
        if not check_ffmpeg_installed():
            print("❌ Error: FFmpeg is not installed or not found in system PATH. Video processing requires FFmpeg.")
            sys.exit(1)

        video_path = args.video
        output_path = args.output
        if not output_path:
            dir_name, base_name = os.path.split(video_path)
            fname, _ = os.path.splitext(base_name)
            output_path = os.path.join(dir_name, f"{fname}_nobg.mov")

        print(f"\n🎬 Input Video:  {video_path}")
        print(f"🎯 Output Destination: {output_path}")
        print(f"⚙️ Export Format: {args.export_format}")
        if bg_color:
            print(f"🎨 Background Color: {bg_color}")
        print(f"✂️ Alpha Matting: {'Enabled' if args.alpha_matting else 'Disabled'}")

        from tqdm import tqdm
        pbar = None

        def pcb(current, total_count, msg):
            nonlocal pbar
            if total_count > 0:
                if pbar is None or pbar.total != total_count:
                    if pbar:
                        pbar.close()
                    pbar = tqdm(total=total_count, desc="Processing Video Frames", unit="frame")
                pbar.n = current
                pbar.refresh()
            else:
                print(f"[Status] {msg}")

        try:
            result = engine.process_video(
                video_path=video_path,
                output_path=output_path,
                fps=args.fps if args.fps > 0 else None,
                export_format=args.export_format,
                bg_color=bg_color,
                alpha_matting=args.alpha_matting,
                progress_callback=pcb
            )
            if pbar:
                pbar.close()
            print(f"\n✅ Video background removal complete! File saved to: {result}")
        except Exception as e:
            if pbar:
                pbar.close()
            print(f"\n❌ Error processing video: {e}")
            sys.exit(1)

    # 2. Batch Directory Mode
    elif args.batch or (args.image and os.path.isdir(args.image[0])):
        batch_dir = args.batch or args.image[0]
        output_dir = args.output or os.path.join(batch_dir, "output_nobg")

        image_files = []
        for root, _, files in os.walk(batch_dir):
            for f in files:
                if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff")):
                    image_files.append(os.path.join(root, f))

        if not image_files:
            print(f"❌ No valid image files found in directory: {batch_dir}")
            sys.exit(1)

        print(f"\n📁 Batch Input Directory: {batch_dir} ({len(image_files)} images found)")
        print(f"🎯 Output Directory: {output_dir}")

        from tqdm import tqdm
        pbar = tqdm(total=len(image_files), desc="Removing Backgrounds", unit="img")

        def pcb(current, total_count, fname):
            pbar.n = current
            pbar.set_postfix_str(fname)
            pbar.refresh()

        results = engine.process_batch_images(
            input_paths=image_files,
            output_dir=output_dir,
            alpha_matting=args.alpha_matting,
            bg_color=bg_color,
            output_format=args.format,
            progress_callback=pcb
        )
        pbar.close()
        print(f"\n✅ Batch processing complete! Saved {len(results)} images to: {output_dir}")

    # 3. Single / Multiple Specific Images
    elif args.image:
        input_paths = args.image
        if len(input_paths) == 1:
            in_file = input_paths[0]
            out_file = args.output
            if not out_file:
                dir_name, base_name = os.path.split(in_file)
                fname, _ = os.path.splitext(base_name)
                out_file = os.path.join(dir_name, f"{fname}_nobg.{args.format.lower()}")

            print(f"\n🖼️ Input Image: {in_file}")
            print(f"🎯 Output File: {out_file}")

            try:
                res = engine.process_single_image(
                    in_file,
                    out_file,
                    alpha_matting=args.alpha_matting,
                    bg_color=bg_color
                )
                print(f"✅ Success! Image saved to: {res}")
            except Exception as e:
                print(f"❌ Error: {e}")
                sys.exit(1)
        else:
            out_dir = args.output or os.path.dirname(input_paths[0])
            print(f"\n🖼️ Processing {len(input_paths)} images...")
            from tqdm import tqdm
            pbar = tqdm(total=len(input_paths), desc="Removing Backgrounds", unit="img")

            def pcb(current, total_count, fname):
                pbar.n = current
                pbar.set_postfix_str(fname)
                pbar.refresh()

            results = engine.process_batch_images(
                input_paths=input_paths,
                output_dir=out_dir,
                alpha_matting=args.alpha_matting,
                bg_color=bg_color,
                output_format=args.format,
                progress_callback=pcb
            )
            pbar.close()
            print(f"✅ Success! Processed {len(results)} images.")

def main():
    parser = argparse.ArgumentParser(
        description="RemoveBG Studio - AI Background Removal for Images & Video Files (FFmpeg)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # For a graphical UI, build and run native/RemoveBGStudio.exe instead (see README.md).

  # Single Image Background Removal:
  python removebg.py -i input.jpg -o output.png

  # Batch Image Directory Background Removal:
  python removebg.py -b ./my_images -o ./no_bg_images --model birefnet-general

  # Video Background Removal with FFmpeg:
  python removebg.py -v input.mp4 -o transparent.mov --export-format mov_transparent

  # Video Background Removal with Green Screen Replacement:
  python removebg.py -v input.mp4 -o greenscreen.mp4 --export-format mp4_solid --bg-color green
"""
    )

    parser.add_argument("-i", "--image", nargs="+", help="Input image file path(s)")
    parser.add_argument("-b", "--batch", help="Input directory containing images for batch processing")
    parser.add_argument("-v", "--video", help="Input video file path to extract and process frames via FFmpeg")
    parser.add_argument("-o", "--output", help="Output file path (single image/video) or output directory (batch)")
    
    parser.add_argument(
        "-m", "--model",
        default="u2net",
        choices=list(AVAILABLE_MODELS.keys()),
        help="AI Model to use for background removal (default: u2net)"
    )
    
    parser.add_argument(
        "-bg", "--bg-color",
        help="Background replacement color (e.g. 'green', 'white', 'black', or Hex '#00FF00'). Default: transparent"
    )
    
    parser.add_argument(
        "-am", "--alpha-matting",
        action="store_true",
        help="Enable alpha matting for finer edge detail (hair, fur)"
    )
    
    parser.add_argument(
        "--format",
        default="png",
        choices=["png", "webp", "jpg"],
        help="Output image format for single/batch processing (default: png)"
    )
    
    parser.add_argument(
        "--fps",
        type=float,
        default=0.0,
        help="Target frame rate for video frame extraction (0 = original video FPS)"
    )
    
    parser.add_argument(
        "--export-format",
        default="mov_transparent",
        choices=["mov_transparent", "webm_transparent", "mp4_solid", "png_sequence"],
        help="Export format for video reassembly (default: mov_transparent)"
    )

    args = parser.parse_args()

    if not args.image and not args.batch and not args.video:
        parser.print_help()
        sys.exit(1)

    run_cli(args)

if __name__ == "__main__":
    main()