import boto3
import logging
from typing import Dict, Any, Optional
from supabase import create_client, Client
from botocore.exceptions import ClientError

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Initialize clients
ssm = boto3.client('ssm')

# Initialize global variables
SUPABASE_URL: Optional[str] = None
SUPABASE_KEY: Optional[str] = None
supabase: Optional[Client] = None

def get_parameter(name: str) -> str:
    """
    Retrieve a parameter from AWS Parameter Store.
    
    Args:
        name: The parameter name to retrieve
        
    Returns:
        The parameter value
        
    Raises:
        ClientError: If parameter retrieval fails
    """
    try:
        response = ssm.get_parameter(Name=name, WithDecryption=True)
        return response['Parameter']['Value']
    except ClientError as e:
        logger.error(f"Failed to get parameter {name}: {str(e)}")
        raise

def initialize_parameters() -> None:
    """Initialize global parameters and clients."""
    global SUPABASE_URL, SUPABASE_KEY, supabase
    
    if SUPABASE_URL is None or SUPABASE_KEY is None:
        try:
            SUPABASE_URL = get_parameter('/supabase/url')
            SUPABASE_KEY = get_parameter('/supabase/anon')
            supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
            logger.info("Initialized Supabase client successfully")
        except Exception as e:
            logger.error(f"Failed to initialize parameters: {str(e)}")
            raise

def check_pending_images() -> Dict[str, Any]:
    """
    Check if there are any unprocessed images in the product_images table.
    
    Returns:
        Dictionary with count of pending images and whether to proceed
    """
    try:
        # Query for unprocessed images
        response = supabase.table('product_images') \
            .select('count', count='exact') \
            .eq('processed', False) \
            .execute()
        
        # Get the count from response
        count = 0
        if hasattr(response, 'count') and response.count is not None:
            count = response.count
        elif response.data and isinstance(response.data, list) and len(response.data) > 0:
            count = len(response.data)
            
        logger.info(f"Found {count} unprocessed images")
        
        # Determine if we should proceed with starting the model
        should_proceed = count > 0
        
        return {
            'pending_images_count': count,
            'should_start_model': should_proceed
        }
        
    except Exception as e:
        logger.error(f"Error checking for pending images: {str(e)}")
        # In case of error, default to proceeding (safer option)
        return {
            'pending_images_count': 0,
            'should_start_model': False,
            'error': str(e)
        }

def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    Main Lambda handler function.
    
    Args:
        event: Lambda event
        context: Lambda context
        
    Returns:
        Dictionary with pending images count and should_start_model flag
    """
    try:
        initialize_parameters()
        
        # Check for pending images
        result = check_pending_images()
        
        # Add a message for clearer logs
        message = f"Found {result['pending_images_count']} pending images for processing"
        if not result['should_start_model']:
            message += ". Skipping model start as there are no images to process."
        
        return {
            'statusCode': 200,
            'body': {
                'message': message,
                'pending_images_count': result['pending_images_count']
            },
            'should_start_model': result['should_start_model']
        }
        
    except Exception as e:
        error_msg = str(e)
        logger.error(f"Lambda execution failed: {error_msg}")
        return {
            'statusCode': 500,
            'body': {
                'message': f"Error checking for pending images: {error_msg}"
            },
            'should_start_model': False
        }