import boto3
import io
import logging
import os
from typing import Dict, List, Optional, Any
import requests
from PIL import Image, ImageDraw, ImageFont
from supabase import create_client, Client
from botocore.exceptions import ClientError
import time

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

# Initialize clients and session
ssm = boto3.client('ssm')
rekognition_client = boto3.client('rekognition')
http_session = requests.Session()

# Initialize global variables
SUPABASE_URL: Optional[str] = None
SUPABASE_KEY: Optional[str] = None
supabase: Optional[Client] = None
PROJECT_ARN: Optional[str] = None
MIN_CONFIDENCE: float = float(os.environ.get('MIN_CONFIDENCE', '50.0'))

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 get_active_model_version() -> Optional[Dict[str, Any]]:
    """
    Get the currently active model version from Supabase.
    
    Returns:
        Dictionary containing model version information or None if not found
    """
    try:
        # First try to find models with RUNNING in model_status field (new approach)
        try:
            response = supabase.table('model_versions') \
                .select('*') \
                .eq('model_status', 'RUNNING') \
                .order('training_timestamp', desc=True) \
                .limit(1) \
                .execute()
            
            if response.data and len(response.data) > 0:
                logger.info("Found running model using model_status field")
                return response.data[0]
        except Exception as inner_e:
            logger.warning(f"Error querying with model_status field: {str(inner_e)}")
        
        # If no model found or error occurred, try with original status field (legacy approach)
        response = supabase.table('model_versions') \
            .select('*') \
            .eq('status', 'RUNNING') \
            .order('training_timestamp', desc=True) \
            .limit(1) \
            .execute()
        
        if response.data and len(response.data) > 0:
            logger.info("Found running model using status field")
            return response.data[0]
        
        # If no running model, try to get the latest trained model (from either field)
        try:
            # Try model_status field first (new approach)
            response = supabase.table('model_versions') \
                .select('*') \
                .eq('model_status', 'TRAINING_COMPLETED') \
                .order('training_timestamp', desc=True) \
                .limit(1) \
                .execute()
                
            if response.data and len(response.data) > 0:
                logger.info("Found latest trained model using model_status field")
                return response.data[0]
        except Exception as inner_e:
            logger.warning(f"Error querying trained models with model_status field: {str(inner_e)}")
        
        # Try status field (legacy approach)
        response = supabase.table('model_versions') \
            .select('*') \
            .eq('status', 'TRAINING_COMPLETED') \
            .order('training_timestamp', desc=True) \
            .limit(1) \
            .execute()
            
        if response.data and len(response.data) > 0:
            logger.info("Found latest trained model using status field")
            return response.data[0]
            
        logger.warning("No running or trained models found in database")
        return None
        
    except Exception as e:
        logger.error(f"Failed to get active model version: {str(e)}")
        return None

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

def process_single_image(item: Dict[str, Any], model_version: Dict[str, Any]) -> None:
    """
    Process a single image from the product_images table.
    
    Args:
        item: Dictionary containing image information
        model_version: Dictionary containing model version information
    """
    record_id = item['id']
    product_id = item['product_id']
    image_path = item['image_path']
    error_message = None

    try:
        rekognition_response = show_custom_labels(image_path, model_version['project_version_arn'])
        logger.info(f"Custom labels detected for product {product_id}: {rekognition_response}")

        if not rekognition_response.get('CustomLabels'):
            logger.warning(f"No custom labels detected for product {product_id}")
            error_message = "No custom labels detected"
            # Update the processed field and error message
            supabase.table('product_images').update({
                'processed': True,
                'error': error_message,
                'project_version_arn': model_version['project_version_arn']
            }).eq('id', record_id).execute()
            return

        new_image_url = save_image_to_supabase(rekognition_response, image_path)
        if not new_image_url:
            error_message = "Failed to save processed image"
            logger.error(f"Failed to save processed image for product {product_id}")
            # Update the processed field and error message
            supabase.table('product_images').update({
                'processed': True,
                'error': error_message,
                'project_version_arn': model_version['project_version_arn']
            }).eq('id', record_id).execute()
            return

        # Check if a record already exists in product_ml table
        try:
            existing_record = supabase.table('product_ml') \
                .select('id') \
                .eq('product_images_id', record_id) \
                .execute()
            
            if existing_record.data and len(existing_record.data) > 0:
                # Update existing record
                logger.info(f"Updating existing product_ml record for product_images_id: {record_id}")
                supabase.table('product_ml') \
                    .update({
                        'ml_image_path': new_image_url,
                        'ml_result': rekognition_response,
                        'project_version_arn': model_version['project_version_arn']
                    }) \
                    .eq('product_images_id', record_id) \
                    .execute()
            else:
                # Insert new record
                logger.info(f"Inserting new product_ml record for product_images_id: {record_id}")
                supabase.table('product_ml').insert({
                    'product_id': product_id,
                    'ml_image_path': new_image_url,
                    'ml_result': rekognition_response,
                    'product_images_id': record_id,
                    'project_version_arn': model_version['project_version_arn']
                }).execute()
        except Exception as e:
            logger.error(f"Error checking/updating product_ml table: {str(e)}")
            # Continue with insert as fallback
            logger.info(f"Falling back to insert for product_ml record")
            supabase.table('product_ml').insert({
                'product_id': product_id,
                'ml_image_path': new_image_url,
                'ml_result': rekognition_response,
                'product_images_id': record_id,
                'project_version_arn': model_version['project_version_arn']
            }).execute()

        # Update the processed field with success
        supabase.table('product_images').update({
            'processed': True,
            'error': None,
            'project_version_arn': model_version['project_version_arn']
        }).eq('id', record_id).execute()

    except Exception as e:
        error_message = str(e)
        logger.error(f"Error processing image {image_path} for product {product_id}: {error_message}")
        # Update the processed field and error message
        supabase.table('product_images').update({
            'processed': True,
            'error': error_message,
            'project_version_arn': model_version['project_version_arn']
        }).eq('id', record_id).execute()
        # Don't raise the exception as we want to continue processing other images

def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    Main Lambda handler function.
    
    Args:
        event: Lambda event
        context: Lambda context
        
    Returns:
        API Gateway response
    """
    try:
        initialize_parameters()
        
        # Get active model version
        model_version = get_active_model_version()
        if not model_version:
            logger.error("No active model version found")
            return {
                'statusCode': 503,
                'body': "No active model version available for processing."
            }
        
        # The model's running state is now checked by check_model_availability function
        # before this function is called in the state machine, so we don't need to check again
        
        response = supabase.table('product_images').select('id, product_id, image_path').eq('processed', False).execute()
        
        if not response.data:
            logger.info("No unprocessed images found")
            return {
                'statusCode': 404,
                'body': "No unprocessed images found in the product_images table."
            }

        for item in response.data:
            process_single_image(item, model_version)

        return {
            'statusCode': 200,
            'body': f"Processed images successfully using model version {model_version['project_version_arn']}."
        }

    except Exception as e:
        logger.error(f"Lambda execution failed: {str(e)}")
        return {
            'statusCode': 500,
            'body': f"Internal server error: {str(e)}"
        }

def display_image(image_url: str, response: Dict[str, Any]) -> Optional[Image.Image]:
    """
    Create an annotated image with bounding boxes and labels.
    
    Args:
        image_url: URL of the image to process
        response: Rekognition response containing labels
        
    Returns:
        Modified PIL Image or None if processing fails
    """
    try:
        img_response = http_session.get(image_url)
        img_response.raise_for_status()

        stream = io.BytesIO(img_response.content)
        image = Image.open(stream)
        draw_image = ImageDraw.Draw(image, 'RGBA')  # Use RGBA mode to support transparency
        imgWidth, imgHeight = image.size

        logger.info(f'Processing custom labels for {image_url}')
        
        for customLabel in response['CustomLabels']:
            logger.debug(f"Label: {customLabel['Name']}, Confidence: {customLabel['Confidence']}")
            
            if 'Geometry' in customLabel:
                box = customLabel['Geometry']['BoundingBox']
                left = imgWidth * box['Left']
                top = imgHeight * box['Top']
                width = imgWidth * box['Width']
                height = imgHeight * box['Height']

                # Text drawing code commented out due to display issues
                # But kept for future reference if needed
                """
                # Calculate text dimensions
                fnt = ImageFont.load_default()
                text = customLabel['Name']
                
                # Calculate text size in a way that works with newer Pillow versions
                try:
                    # For newer Pillow versions (9.2.0+)
                    # Use the newer getbbox method and calculate dimensions
                    text_bbox = fnt.getbbox(text)
                    text_width = text_bbox[2] - text_bbox[0]
                    text_height = text_bbox[3] - text_bbox[1]
                except AttributeError:
                    try:
                        # Try older method (Pillow < 9.0)
                        text_width, text_height = draw_image.textsize(text, font=fnt)
                    except AttributeError:
                        # Last resort - use a fixed size
                        text_width = len(text) * 8  # Approximate width
                        text_height = 12            # Approximate height
                
                # Draw bounding rectangle for the text (semi-transparent black background)
                text_bg_coords = [
                    (left, top),
                    (left + text_width + 10, top + text_height + 4)  # Add padding
                ]
                draw_image.rectangle(text_bg_coords, fill=(0, 0, 0, 128))  # Semi-transparent black (RGBA)
                
                # Draw text in white on top of the transparent background
                draw_image.text((left + 5, top + 2), text, fill=(255, 255, 255), font=fnt)  # White text
                """

                # Draw bounding box for the object with reduced thickness (2 instead of 5)
                points = (
                    (left, top),
                    (left + width, top),
                    (left + width, top + height),
                    (left, top + height),
                    (left, top)
                )
                draw_image.line(points, fill='#000000', width=2)  # Reduced thickness from 5 to 2

        return image

    except Exception as e:
        logger.error(f"Error processing image {image_url}: {str(e)}")
        return None

def show_custom_labels(image_url: str, model_arn: str) -> Dict[str, Any]:
    """
    Detect custom labels in an image using AWS Rekognition.
    
    Args:
        image_url: URL of the image to analyze
        model_arn: ARN of the model version to use
        
    Returns:
        Rekognition response or empty dict if processing fails
    """
    try:
        img_response = http_session.get(image_url)
        img_response.raise_for_status()
        
        img_bytes = io.BytesIO(img_response.content)

        response = rekognition_client.detect_custom_labels(
            Image={'Bytes': img_bytes.getvalue()},
            MinConfidence=MIN_CONFIDENCE,
            ProjectVersionArn=model_arn
        )
        return response

    except Exception as e:
        logger.error(f"Error detecting custom labels for {image_url}: {str(e)}")
        return {'CustomLabels': []}

def save_image_to_supabase(rekognition_response: Dict[str, Any], image_path: str) -> Optional[str]:
    """
    Save the processed image to Supabase storage.
    
    Args:
        rekognition_response: Rekognition API response
        image_path: Original image URL
        
    Returns:
        Public URL of saved image or None if saving fails
    """
    try:
        modified_image = display_image(image_path, rekognition_response)
        if not modified_image:
            return None
        
        output_stream = io.BytesIO()
        
        if modified_image.mode == 'RGBA':
            modified_image = modified_image.convert('RGB')
            
        modified_image.save(output_stream, format='JPEG')
        output_stream.seek(0)
        
        # Extract the original filename
        original_filename = image_path.split('/')[-1]
        file_path = f"product/{original_filename.split('.')[0]}.jpg"
        bucket_name = 'uploads/ml'  # This is the correct bucket
        
        # Instead of trying to delete first, use upsert to overwrite
        logger.info(f"Uploading file to {bucket_name}/{file_path} with upsert")
        upload_response = supabase.storage.from_(bucket_name).upload(
            file_path,
            output_stream.getvalue(),
            {"contentType": "image/jpeg", "upsert": "true"}  # Use upsert: "true" as a string, not a boolean
        )
        logger.info(f"Upload successful")
        
        # Get the public URL
        public_url = supabase.storage.from_(bucket_name).get_public_url(file_path)
        logger.info(f"Public URL: {public_url}")
        
        return public_url.rstrip('?')

    except Exception as e:
        logger.error(f"Error saving image to Supabase: {str(e)}")
        return None