main.py

#!/usr/bin/env python3
"""
Text Extraction Tool - CLI Interface

Extract text from multiple file types using OpenAI and Gemini APIs.
Supports: PDF, TXT, DOC, DOCX, XLS, XLSX, MP4, MP3, and more.
"""

import argparse
import sys
from pathlib import Path
from typing import List

from src.text_extractor import TextExtractor

def main():
    """Main CLI function."""
    parser = argparse.ArgumentParser(
        description="Extract text from multiple file types using AI services",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python main.py file1.pdf file2.mp4 file3.docx
  python main.py --directory /path/to/files --recursive
  python main.py *.pdf --output results.txt --parallel --workers 8
  python main.py --info  # Show supported file types
        """
    )

    # File/directory arguments
    parser.add_argument(
        'files',
        nargs='*',
        help='Files to process (supports glob patterns)'
    )

    parser.add_argument(
        '-d', '--directory',
        type=str,
        help='Process all supported files in directory'
    )

    parser.add_argument(
        '-r', '--recursive',
        action='store_true',
        help='Process directory recursively'
    )

    # Processing options
    parser.add_argument(
        '--parallel',
        action='store_true',
        default=True,
        help='Process files in parallel (default: True)'
    )

    parser.add_argument(
        '--no-parallel',
        action='store_true',
        help='Process files sequentially'
    )

    parser.add_argument(
        '--workers',
        type=int,
        default=4,
        help='Number of parallel workers (default: 4)'
    )

    # Output options
    parser.add_argument(
        '-o', '--output',
        type=str,
        help='Save results to file'
    )

    parser.add_argument(
        '--summary-only',
        action='store_true',
        help='Show only summary, not full extracted text'
    )

    # Info options
    parser.add_argument(
        '--info',
        action='store_true',
        help='Show supported file types and exit'
    )

    parser.add_argument(
        '--image-json',
        type=str,
        metavar='IMAGE_FILE',
        help='Transcribe single image and return JSON with title and description'
    )

    args = parser.parse_args()

    # Handle info request
    if args.info:
        show_info()
        return

    # Handle image JSON transcription
    if args.image_json:
        handle_image_json(args.image_json, args.output)
        return

    # Validate arguments
    if not args.files and not args.directory:
        parser.error("Must specify either files or --directory")

    try:
        # Initialize text extractor
        print("Initializing text extractor...")
        extractor = TextExtractor()

        # Show supported extensions
        if args.info:
            show_supported_extensions(extractor)
            return

        # Determine parallel processing
        use_parallel = args.parallel and not args.no_parallel

        # Process files
        results = []

        if args.directory:
            # Process directory
            print(f"Processing directory: {args.directory}")
            results = extractor.extract_from_directory(
                args.directory,
                recursive=args.recursive,
                parallel=use_parallel,
                max_workers=args.workers
            )
        else:
            # Process individual files
            file_paths = []
            for file_pattern in args.files:
                # Handle absolute paths and tilde expansion
                expanded_pattern = Path(file_pattern).expanduser()

                # If it's an absolute path or direct file, check if it exists
                if expanded_pattern.is_absolute() or '/' in file_pattern or '\\' in file_pattern:
                    if expanded_pattern.exists():
                        file_paths.append(expanded_pattern)
                    else:
                        print(f"Warning: File not found: {file_pattern}")
                else:
                    # Try glob pattern matching for relative paths
                    try:
                        paths = list(Path().glob(file_pattern))
                        if paths:
                            file_paths.extend(paths)
                        else:
                            # Try as direct path
                            path = Path(file_pattern)
                            if path.exists():
                                file_paths.append(path)
                            else:
                                print(f"Warning: File not found: {file_pattern}")
                    except ValueError as e:
                        # Fallback for invalid glob patterns
                        path = Path(file_pattern)
                        if path.exists():
                            file_paths.append(path)
                        else:
                            print(f"Warning: File not found: {file_pattern}")

            if not file_paths:
                print("Error: No valid files found")
                sys.exit(1)

            print(f"Processing {len(file_paths)} files...")
            results = extractor.extract_from_files(
                file_paths,
                parallel=use_parallel,
                max_workers=args.workers
            )

        # Display results
        if not args.summary_only:
            print_results(results)

        # Print summary
        extractor.print_summary(results)

        # Save to file if requested
        if args.output:
            extractor.save_results_to_file(results, args.output)

    except KeyboardInterrupt:
        print("\nInterrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

def handle_image_json(image_path: str, output_file: str = None):
    """Handle image transcription to JSON."""
    import json

    # Validate image file
    path = Path(image_path)
    if not path.exists():
        print(f"Error: Image file not found: {image_path}", file=sys.stderr)
        sys.exit(1)

    if not path.is_file():
        print(f"Error: Path is not a file: {image_path}", file=sys.stderr)
        sys.exit(1)

    # Check if file has supported extension
    supported_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.tif', '.webp'}
    if path.suffix.lower() not in supported_extensions:
        print(f"Error: Unsupported image format: {path.suffix}", file=sys.stderr)
        print(f"Supported formats: {', '.join(supported_extensions)}", file=sys.stderr)
        sys.exit(1)

    try:
        # Initialize text extractor
        print("Initializing AI models...", file=sys.stderr)
        extractor = TextExtractor()

        # Transcribe image
        print(f"Processing image: {path.name}...", file=sys.stderr)
        json_result = extractor.transcribe_image_to_json(image_path)

        # Handle output
        if output_file:
            # Save to file
            output_path = Path(output_file)
            with open(output_path, 'w', encoding='utf-8') as f:
                # Re-parse and pretty print
                data = json.loads(json_result)
                json.dump(data, f, indent=2, ensure_ascii=False)
            print(f"Results saved to: {output_path}", file=sys.stderr)
        else:
            # Print to console (pretty formatted)
            data = json.loads(json_result)
            print(json.dumps(data, indent=2, ensure_ascii=False))

    except Exception as e:
        print(f"Error: {str(e)}", file=sys.stderr)
        sys.exit(1)

def show_info():
    """Show general information about the tool."""
    print("""
Text Extraction Tool
====================

This tool extracts text from various file types using AI services:

OpenAI (GPT) - For documents and text files:
  • PDF files (.pdf)
  • Text files (.txt)
  • Word documents (.doc, .docx)
  • Excel spreadsheets (.xls, .xlsx)

Google Gemini - For multimedia files:
  • Video files (.mp4, .avi, .mov, .mkv)
  • Audio files (.mp3, .wav, .m4a, .webm, .ogg)

Vision APIs (OpenAI/Gemini) - For image transcription:
  • Image files (.jpg, .jpeg, .png, .gif, .bmp, .tiff, .tif, .webp)
  • Returns structured JSON with title, description, and extracted text

Setup:
1. Install dependencies: pip install -r requirements.txt
2. Copy env_example.txt to .env
3. Add your API keys to the .env file

Usage examples:
  python main.py document.pdf video.mp4 audio.mp3 audio.webm
  python main.py --directory /path/to/files --recursive
  python main.py *.pdf --output results.txt
  python main.py --image-json photo.jpg  # Transcribe image to JSON
  python main.py --image-json photo.jpg --output result.json
    """)

def show_supported_extensions(extractor: TextExtractor):
    """Show supported file extensions."""
    extensions = extractor.get_supported_extensions()

    print("\nSupported File Extensions:")
    print("=" * 30)

    for processor, exts in extensions.items():
        print(f"\n{processor}:")
        for ext in sorted(exts):
            print(f"  {ext}")

def print_results(results: List[dict]):
    """Print extraction results."""
    print("\n" + "=" * 60)
    print("EXTRACTION RESULTS")
    print("=" * 60)

    for i, result in enumerate(results, 1):
        print(f"\n[{i}] {Path(result['file_path']).name}")
        print("-" * 40)

        if result['success']:
            print(f"Status: SUCCESS")
            print(f"Processor: {result['processor_used']}")
            print(f"Processing time: {result['processing_time']:.2f}s")
            print(f"Text length: {result['text_length']} characters")

            if result['extracted_text']:
                print(f"\nExtracted Text:")
                print("-" * 20)
                # Show first 500 characters
                text = result['extracted_text']
                if len(text) > 500:
                    print(text[:500] + "\n... (truncated)")
                else:
                    print(text)
        else:
            print(f"Status: FAILED")
            print(f"Error: {result['error']}")

if __name__ == "__main__":
    main()