import boto3
import logging
import os
import json
import time
import re
from datetime import datetime
from botocore.exceptions import ClientError
import uuid

from utils.logging_utils import setup_logger, reduce_logging_verbosity

# Initialize logger
logger = setup_logger()
reduce_logging_verbosity()

# Initialize clients
rekognition_client = boto3.client('rekognition')
ssm_client = boto3.client('ssm')
supabase_client = None

def get_ssm_parameter(param_name):
    """Get a parameter from SSM Parameter Store"""
    try:
        response = ssm_client.get_parameter(Name=param_name, WithDecryption=True)
        return response['Parameter']['Value']
    except Exception as e:
        logger.warning(f"Failed to get SSM parameter {param_name}: {str(e)}")
        return None

def fix_project_arn(project_arn):
    """
    Fix the project ARN format by ensuring it ends with the correct project ID.
    AWS Rekognition project ARNs must use the correct project ID: 1725357683732
    """
    if not project_arn:
        logger.error("No project ARN provided")
        return None
        
    logger.info(f"Checking project ARN format: {project_arn}")
    
    # Use the known working project ID 
    project_id = "1725357683732"
        
    # Check if the ARN already has a version number
    pattern = r'project\/[a-zA-Z0-9_\.-]{1,255}\/[0-9]+$'
    if re.search(pattern, project_arn):
        # If it matches the pattern but has a different ID, replace it
        if not project_arn.endswith(f"/{project_id}"):
            # Extract everything up to the last slash
            base_arn = '/'.join(project_arn.split('/')[:-1])
            fixed_arn = f"{base_arn}/{project_id}"
            logger.info(f"Updated project ID: {project_arn} -> {fixed_arn}")
            return fixed_arn
        else:
            logger.info(f"Project ARN format is already correct: {project_arn}")
            return project_arn
    
    # Get just the project name without any trailing slashes
    if '/project/' in project_arn:
        base_arn = project_arn.rstrip('/')
        # Standard case: append the correct project ID
        fixed_arn = f"{base_arn}/{project_id}"
        logger.info(f"Fixed ARN by adding correct project ID: {project_arn} -> {fixed_arn}")
        return fixed_arn
    
    logger.warning(f"Cannot determine correct ARN format from: {project_arn}")
    
    # Fall back to simple approach if more sophisticated extraction fails
    if project_arn.endswith('/'):
        fixed_arn = f"{project_arn}{project_id}"
    else:
        fixed_arn = f"{project_arn}/{project_id}"
        
    logger.info(f"Applied fallback ARN fix: {project_arn} -> {fixed_arn}")
    return fixed_arn

def get_supabase_client():
    """Initialize the Supabase client"""
    global supabase_client
    
    if supabase_client is not None:
        return supabase_client
        
    try:
        from supabase import create_client

        # Get Supabase URL and key from environment variables or SSM
        supabase_url = os.environ.get('SUPABASE_URL')
        supabase_key = os.environ.get('SUPABASE_KEY')
        
        # If not in env vars, try SSM
        if not supabase_url:
            supabase_url = get_ssm_parameter('/datafy-rekognition-stack/supabase-url')
        
        if not supabase_key:
            supabase_key = get_ssm_parameter('/datafy-rekognition-stack/supabase-anon-key')
            
        if not supabase_url or not supabase_key:
            logger.error("Supabase credentials not found")
            return None
            
        # Initialize Supabase client
        supabase_client = create_client(supabase_url, supabase_key)
        logger.info("Supabase client initialized")
        return supabase_client
        
    except Exception as e:
        logger.error(f"Failed to initialize Supabase client: {str(e)}")
        return None

def get_model_status(project_version_arn):
    """
    Get the current status of a model version
    
    Args:
        project_version_arn: ARN of the model version
        
    Returns:
        str: Status of the model
    """
    try:
        response = rekognition_client.describe_project_versions(
            ProjectArn='/'.join(project_version_arn.split('/')[:-2]),
            VersionNames=[project_version_arn.split('/')[-1]]
        )
        
        if not response.get('ProjectVersionDescriptions'):
            logger.warning(f"No project version found for ARN: {project_version_arn}")
            return None
            
        status = response['ProjectVersionDescriptions'][0]['Status']
        logger.info(f"Model status: {status}")
        return status
        
    except Exception as e:
        logger.error(f"Error getting model status: {str(e)}")
        return None

def update_supabase_model_status(model_id, status, error_message=None):
    """
    Update the model status in Supabase
    
    Args:
        model_id: ID of the model in Supabase
        status: New status for the model
        error_message: Optional error message
        
    Returns:
        bool: Success or failure
    """
    try:
        supabase = get_supabase_client()
        if not supabase:
            logger.error("No Supabase client available")
            return False
            
        # Prepare update data
        update_data = {
            'status': status,
            'updated_at': datetime.utcnow().isoformat()
        }
        
        # Add error message if provided
        if error_message:
            update_data['error_message'] = error_message
            
        # Update the model in Supabase
        result = supabase.table('models').update(update_data).eq('id', model_id).execute()
        
        if result.data:
            logger.info(f"Updated model {model_id} status to {status}")
            return True
        else:
            logger.warning(f"No model updated, check if model {model_id} exists")
            return False
            
    except Exception as e:
        logger.error(f"Error updating model status in Supabase: {str(e)}")
        return False

def update_model_status(model_arn, model_id):
    """
    Check and update the status of a model
    
    Args:
        model_arn: ARN of the model version
        model_id: ID of the model in Supabase
        
    Returns:
        dict: Information about the updated model
    """
    try:
        # Get current status from Rekognition
        current_status = get_model_status(model_arn)
        
        if not current_status:
            logger.error(f"Could not get status for model ARN: {model_arn}")
            return {
                "success": False,
                "status": "ERROR",
                "message": "Could not get model status from Rekognition"
            }
            
        # Status transitions to handle
        status_mapping = {
            "TRAINING_IN_PROGRESS": "TRAINING",
            "TRAINING_COMPLETED": "RUNNING",
            "TRAINING_FAILED": "FAILED",
            "RUNNING": "RUNNING",
            "STARTING_HOSTING": "RUNNING", 
            "STOPPING_HOSTING": "STOPPING",
            "STOPPED": "STOPPED",
            "DELETING": "DELETING",
            "FAILED": "FAILED"
        }
        
        # Map Rekognition status to Supabase status
        supabase_status = status_mapping.get(current_status, current_status)
        
        # Update status in Supabase
        update_success = update_supabase_model_status(model_id, supabase_status)
        
        if not update_success:
            logger.warning(f"Failed to update model {model_id} status in Supabase")
            
        return {
            "success": update_success,
            "status": supabase_status,
            "rekognition_status": current_status,
            "model_arn": model_arn,
            "model_id": model_id
        }
        
    except Exception as e:
        logger.error(f"Error in update_model_status: {str(e)}")
        return {
            "success": False,
            "status": "ERROR",
            "message": str(e)
        }

def lambda_handler(event, context):
    """
    Lambda handler to check and update model status
    
    Args:
        event: Lambda event
        context: Lambda context
        
    Returns:
        dict: Response with update results
    """
    try:
        logger.info(f"Starting check and update model: {json.dumps(event, default=str)}")
        
        # Extract model information from the event
        model_arn = event.get('model_arn')
        model_id = event.get('model_id')
        project_name = event.get('project_name')
        
        # Validate inputs
        if not model_arn and not model_id:
            logger.error("No model_arn or model_id provided in event")
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'message': 'Missing required parameters',
                    'error': 'Both model_arn and model_id are required'
                })
            }
            
        # If we have a model_id but no model_arn, try to get it from Supabase
        if model_id and not model_arn:
            try:
                supabase = get_supabase_client()
                if supabase:
                    response = supabase.table('models').select('arn').eq('id', model_id).execute()
                    if response.data and len(response.data) > 0:
                        model_arn = response.data[0].get('arn')
                        logger.info(f"Retrieved model ARN from Supabase: {model_arn}")
            except Exception as e:
                logger.error(f"Error retrieving model ARN from Supabase: {str(e)}")
        
        # If we have no model_arn but a project_name, try to find the latest model
        if not model_arn and project_name:
            # Get PROJECT_NAME from environment variables with fallbacks
            if not project_name:
                project_name = os.environ.get('PROJECT_NAME')
                logger.info(f"Using project name from environment: {project_name}")
            
            # Try to get from SSM if not in environment
            if not project_name:
                project_name = get_ssm_parameter('/datafy-rekognition-stack/project-name')
                logger.info(f"Got project name from SSM: {project_name}")
            
            if project_name:
                try:
                    # Construct the ARN using the project name
                    aws_region = os.environ.get('AWS_REGION') or 'eu-west-1'
                    aws_account_id = os.environ.get('AWS_ACCOUNT_ID') or '587594388832'
                    project_arn = f"arn:aws:rekognition:{aws_region}:{aws_account_id}:project/{project_name}"
                    logger.info(f"Constructed project ARN: {project_arn}")
                    
                    # Fix project ARN format using the function
                    project_arn = fix_project_arn(project_arn)
                    if project_arn:
                        logger.info(f"Using fixed project ARN: {project_arn}")
                        
                        # Try to find the most recent model version
                        response = rekognition_client.describe_project_versions(
                            ProjectArn=project_arn
                        )
                        
                        if response.get('ProjectVersionDescriptions'):
                            # Sort by creation timestamp (newest first)
                            versions = sorted(
                                response['ProjectVersionDescriptions'],
                                key=lambda x: x.get('CreationTimestamp', 0),
                                reverse=True
                            )
                            
                            if versions:
                                model_arn = versions[0]['ProjectVersionArn']
                                logger.info(f"Found latest model ARN: {model_arn}")
                                
                                # If no model_id is provided, create a UUID for tracking
                                if not model_id:
                                    model_id = str(uuid.uuid4())
                                    logger.info(f"Generated temporary model_id: {model_id}")
                except Exception as e:
                    logger.error(f"Error finding latest model: {str(e)}")
                
        if not model_arn:
            logger.error("No model ARN found or provided")
            return {
                'statusCode': 404,
                'body': json.dumps({
                    'message': 'No model ARN found',
                    'error': 'Model not found'
                })
            }
            
        # Check and update the model status
        update_result = update_model_status(model_arn, model_id)
        
        logger.info(f"Model update result: {update_result}")
        
        return {
            'statusCode': 200 if update_result.get('success', False) else 500,
            'body': json.dumps(update_result),
            'update_result': update_result
        }
        
    except Exception as e:
        logger.error(f"Error in lambda_handler: {str(e)}", exc_info=True)
        return {
            'statusCode': 500,
            'body': json.dumps({
                'message': f'Error checking and updating model: {str(e)}',
                'error': 'Internal server error'
            })
        }