tam.py

#!/usr/bin/env python3

import argparse
import sys
import os
import json
from datetime import datetime
from typing import Dict, Any, List

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.database import DatabaseManager, SQLiteAdapter, PostgreSQLAdapter, create_default_schema
from src.metadata_extractor import ImageMetadataExtractor
from src.query_engine import QueryEngine


class TAMCLI:
    """TAM (Image Metadata) Command Line Interface."""
    
    def __init__(self):
        self.db_manager = DatabaseManager()
        self.query_engine = None
        self.config = {}
        self.setup_adapters()
    
    def setup_adapters(self):
        """Setup database adapters."""
        self.db_manager.register_adapter('sqlite', SQLiteAdapter())
        self.db_manager.register_adapter('postgresql', PostgreSQLAdapter())
    
    def load_config(self, config_path: str = None):
        """Load configuration from file."""
        if config_path and os.path.exists(config_path):
            with open(config_path, 'r') as f:
                self.config = json.load(f)
        else:
            # Default configuration
            self.config = {
                'default_database': {
                    'adapter': 'sqlite',
                    'connection_string': 'images.db',
                    'table_name': 'images'
                }
            }
    
    def init_database(self, args):
        """Initialize database with default schema."""
        adapter_type = args.adapter or 'sqlite'
        connection_string = args.connection or 'images.db'
        table_name = args.table or 'images'
        
        try:
            self.db_manager.connect_to_database(adapter_type, connection_string, 'default')
            adapter = self.db_manager.get_adapter('default')
            schema = create_default_schema()
            adapter.create_table(table_name, schema)
            print(f"Database initialized successfully: {connection_string} (table: {table_name})")
        except Exception as e:
            print(f"Error initializing database: {e}")
            return False
        
        self.query_engine = QueryEngine(self.db_manager)
        return True
    
    def index_images(self, args):
        """Index images in directory or single file."""
        if not self.query_engine:
            print("Please initialize database first using 'init' command")
            return False
        
        extractor = ImageMetadataExtractor()
        
        if args.directory:
            metadata_list = extractor.batch_extract(args.directory)
            print(f"Found {len(metadata_list)} images to index")
        elif args.file:
            try:
                metadata = extractor.extract_metadata(args.file)
                metadata_list = [metadata]
                print(f"Indexing single file: {args.file}")
            except Exception as e:
                print(f"Error extracting metadata: {e}")
                return False
        else:
            print("Please specify either --directory or --file")
            return False
        
        # Store metadata in database
        adapter = self.db_manager.get_adapter('default')
        table_name = self.config.get('default_database', {}).get('table_name', 'images')
        
        indexed_count = 0
        for metadata in metadata_list:
            try:
                adapter.insert_metadata(table_name, metadata)
                indexed_count += 1
            except Exception as e:
                print(f"Error indexing {metadata.get('file_name', 'unknown')}: {e}")
        
        print(f"Successfully indexed {indexed_count} images")
        return True
    
    def search(self, args):
        """Search for images based on criteria."""
        if not self.query_engine:
            print("Please initialize database first using 'init' command")
            return False
        
        table_name = self.config.get('default_database', {}).get('table_name', 'images')
        results = []
        
        # Build search query based on provided arguments
        query_params = {}
        
        if args.location:
            loc_params = {}
            if args.latitude and args.longitude:
                loc_params['latitude'] = args.latitude
                loc_params['longitude'] = args.longitude
                if args.radius:
                    loc_params['radius_km'] = args.radius
            query_params['location'] = loc_params
        
        if args.start_time or args.end_time:
            time_params = {}
            if args.start_time:
                time_params['start_time'] = args.start_time
            if args.end_time:
                time_params['end_time'] = args.end_time
            query_params['time'] = time_params
        
        if args.width_min is not None or args.width_max is not None:
            width = (args.width_min or 0, args.width_max or 99999)
            if 'properties' not in query_params:
                query_params['properties'] = {}
            query_params['properties']['width'] = width
        
        if args.height_min is not None or args.height_max is not None:
            height = (args.height_min or 0, args.height_max or 99999)
            if 'properties' not in query_params:
                query_params['properties'] = {}
            query_params['properties']['height'] = height
        
        if args.format:
            if 'properties' not in query_params:
                query_params['properties'] = {}
            query_params['properties']['format_type'] = args.format
        
        if args.pattern:
            query_params['filename_pattern'] = args.pattern
        
        if args.camera_make or args.camera_model:
            exif_params = {}
            if args.camera_make:
                exif_params['camera_make'] = args.camera_make
            if args.camera_model:
                exif_params['camera_model'] = args.camera_model
            query_params['exif'] = exif_params
        
        # Execute search
        try:
            if query_params:
                results = self.query_engine.advanced_search('default', table_name, query_params)
            else:
                # If no criteria specified, return all images
                adapter = self.db_manager.get_adapter('default')
                results = adapter.query_metadata(table_name, {})
        except Exception as e:
            print(f"Error during search: {e}")
            return False
        
        # Display results
        print(f"Found {len(results)} matching images:")
        for i, result in enumerate(results[:args.limit], 1):
            print(f"\n{i}. {result['file_name']}")
            print(f"   Path: {result['file_path']}")
            print(f"   Size: {result.get('file_size', 0):,} bytes")
            print(f"   Dimensions: {result.get('width', 'N/A')}x{result.get('height', 'N/A')}")
            print(f"   Format: {result.get('format', 'N/A')}")
            print(f"   Created: {result.get('created_time', 'N/A')}")
            
            # Show distance if available (from location search)
            if 'distance_km' in result:
                print(f"   Distance: {result['distance_km']:.2f} km")
        
        if len(results) > args.limit:
            print(f"\n... and {len(results) - args.limit} more results (use --limit to see more)")
        
        return True
    
    def list_databases(self, args):
        """List active database connections."""
        connections = self.db_manager.list_active_connections()
        print(f"Active database connections: {len(connections)}")
        for conn in connections:
            print(f"  - {conn}")
    
    def export_results(self, args):
        """Export search results to file."""
        if not self.query_engine:
            print("Please initialize database first using 'init' command")
            return False
        
        # For simplicity, export all images
        table_name = self.config.get('default_database', {}).get('table_name', 'images')
        adapter = self.db_manager.get_adapter('default')
        results = adapter.query_metadata(table_name, {})
        
        export_format = args.format.lower()
        
        if export_format == 'json':
            with open(args.output, 'w') as f:
                json.dump(results, f, indent=2, default=str)
        elif export_format == 'csv':
            import csv
            if results:
                with open(args.output, 'w', newline='') as f:
                    writer = csv.DictWriter(f, fieldnames=results[0].keys())
                    writer.writeheader()
                    writer.writerows(results)
        else:
            print(f"Unsupported export format: {export_format}")
            return False
        
        print(f"Exported {len(results)} records to {args.output}")
        return True
    
    def run(self):
        """Main CLI entry point."""
        parser = argparse.ArgumentParser(description='TAM - Image Metadata Search Tool')
        parser.add_argument('--config', help='Configuration file path')
        
        subparsers = parser.add_subparsers(dest='command', help='Available commands')
        
        # Initialize command
        init_parser = subparsers.add_parser('init', help='Initialize database')
        init_parser.add_argument('--adapter', choices=['sqlite', 'postgresql'], default='sqlite')
        init_parser.add_argument('--connection', default='images.db')
        init_parser.add_argument('--table', default='images')
        
        # Index command
        index_parser = subparsers.add_parser('index', help='Index images')
        index_group = index_parser.add_mutually_exclusive_group(required=True)
        index_group.add_argument('--directory', help='Directory to scan for images')
        index_group.add_argument('--file', help='Single image file to index')
        
        # Search command
        search_parser = subparsers.add_parser('search', help='Search for images')
        search_parser.add_argument('--latitude', type=float, help='Latitude for location search')
        search_parser.add_argument('--longitude', type=float, help='Longitude for location search')
        search_parser.add_argument('--radius', type=float, help='Radius in km for location search')
        search_parser.add_argument('--start-time', help='Start time (ISO format)')
        search_parser.add_argument('--end-time', help='End time (ISO format)')
        search_parser.add_argument('--width-min', type=int, help='Minimum width')
        search_parser.add_argument('--width-max', type=int, help='Maximum width')
        search_parser.add_argument('--height-min', type=int, help='Minimum height')
        search_parser.add_argument('--height-max', type=int, help='Maximum height')
        search_parser.add_argument('--format', help='Image format (JPEG, PNG, etc.)')
        search_parser.add_argument('--pattern', help='Filename pattern (supports wildcards)')
        search_parser.add_argument('--camera-make', help='Camera make')
        search_parser.add_argument('--camera-model', help='Camera model')
        search_parser.add_argument('--limit', type=int, default=20, help='Maximum number of results')
        
        # List command
        list_parser = subparsers.add_parser('list', help='List database connections')
        
        # Export command
        export_parser = subparsers.add_parser('export', help='Export results')
        export_parser.add_argument('--output', required=True, help='Output file path')
        export_parser.add_argument('--format', choices=['json', 'csv'], default='json', help='Export format')
        
        args = parser.parse_args()
        
        if not args.command:
            parser.print_help()
            return
        
        # Load configuration
        self.load_config(args.config)
        
        # Initialize default database connection if specified
        if args.command != 'init':
            default_db = self.config.get('default_database', {})
            adapter_type = default_db.get('adapter', 'sqlite')
            connection_string = default_db.get('connection_string', 'images.db')
            
            try:
                self.db_manager.connect_to_database(adapter_type, connection_string, 'default')
                self.query_engine = QueryEngine(self.db_manager)
            except Exception as e:
                print(f"Error connecting to database: {e}")
                return
        
        # Execute command
        if args.command == 'init':
            self.init_database(args)
        elif args.command == 'index':
            self.index_images(args)
        elif args.command == 'search':
            self.search(args)
        elif args.command == 'list':
            self.list_databases(args)
        elif args.command == 'export':
            self.export_results(args)
        
        # Cleanup
        self.db_manager.disconnect_all()


def main():
    """Main entry point."""
    cli = TAMCLI()
    cli.run()


if __name__ == '__main__':
    main()