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

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

# Initialize rekognition client
rekognition = boto3.client('rekognition')

def get_project_details(project_arn):
    """
    Get project details including dataset ARNs directly from AWS.
    
    Args:
        project_arn: ARN of the project
        
    Returns:
        dict: Project details including datasets
    """
    try:
        # Extract project name from ARN
        match = re.match(r'arn:aws:rekognition:[^:]+:[^:]+:project/([^/]+)', project_arn)
        if not match:
            logger.error(f"Could not extract project name from ARN: {project_arn}")
            return None
            
        project_name = match.group(1)
        logger.info(f"Looking up project details for: {project_name}")
        
        # Use AWS's recommended approach to get dataset information
        response = rekognition.describe_projects(ProjectNames=[project_name])
        
        if 'ProjectDescriptions' in response and response['ProjectDescriptions']:
            project = response['ProjectDescriptions'][0]
            return project
        else:
            logger.warning(f"No project found with name: {project_name}")
            return None
    except Exception as e:
        logger.error(f"Error getting project details: {str(e)}")
        return None

def delete_dataset(dataset_arn, max_retries=3):
    """
    Delete a dataset using its ARN with retries.
    
    Args:
        dataset_arn: ARN of the dataset to delete
        max_retries: Maximum number of deletion attempts
        
    Returns:
        dict: Result with status and message
    """
    for attempt in range(max_retries):
        try:
            logger.info(f"Attempting to delete dataset (attempt {attempt+1}/{max_retries}): {dataset_arn}")
            
            # Log ARN format for debugging
            pattern = r'^arn:aws:rekognition:[^:]+:\d+:project/[^/]+/dataset/(train|test)/\d+$'
            if re.match(pattern, dataset_arn):
                logger.info(f"Dataset ARN matches expected pattern: {dataset_arn}")
            else:
                logger.warning(f"Dataset ARN does not match expected pattern: {dataset_arn}")
                logger.warning(f"Expected pattern: {pattern}")
                # Try to extract components and rebuild the ARN
                components = re.match(r'arn:aws:rekognition:([^:]+):(\d+):project/([^/]+)(?:/([^/]+))?', dataset_arn)
                if components:
                    region, account, project, project_id = components.groups()
                    dataset_type = "train" if "/train" in dataset_arn.lower() else "test"
                    if project_id is None:
                        project_id = "1725357683732"
                    corrected_arn = f"arn:aws:rekognition:{region}:{account}:project/{project}/dataset/{dataset_type}/{project_id}"
                    logger.info(f"Rebuilt ARN: {corrected_arn}")
                    dataset_arn = corrected_arn
            
            response = rekognition.delete_dataset(DatasetArn=dataset_arn)
            logger.info(f"Successfully deleted dataset: {dataset_arn}")
            
            return {
                "success": True,
                "message": f"Successfully deleted dataset: {dataset_arn}",
                "response": response
            }
        except ClientError as e:
            error_code = e.response['Error']['Code']
            error_message = e.response['Error']['Message']
            
            # If dataset doesn't exist, consider it a success
            if error_code == 'ResourceNotFoundException':
                logger.info(f"Dataset doesn't exist (already deleted): {dataset_arn}")
                return {
                    "success": True,
                    "message": f"Dataset already deleted: {dataset_arn}"
                }
            
            # If dataset is in use, log the error but can't proceed
            if "Cannot delete dataset when in use" in error_message:
                logger.error(f"Dataset is in use and cannot be deleted: {dataset_arn}")
                return {
                    "success": False,
                    "message": f"Cannot delete dataset because it's in use: {dataset_arn}",
                    "error": error_message
                }
                
            # Wait and retry for other errors
            logger.warning(f"Error deleting dataset (attempt {attempt+1}): {error_code} - {error_message}")
            if attempt < max_retries - 1:
                wait_time = 2 ** attempt  # Exponential backoff
                logger.info(f"Waiting {wait_time} seconds before retry...")
                time.sleep(wait_time)
            else:
                return {
                    "success": False,
                    "message": f"Failed to delete dataset after {max_retries} attempts: {dataset_arn}",
                    "error": f"{error_code}: {error_message}"
                }
    
    return {
        "success": False,
        "message": f"Failed to delete dataset: {dataset_arn}"
    }

def wait_for_dataset_deletion(dataset_arn, timeout=60, check_interval=5):
    """
    Wait for a dataset to be fully deleted.
    
    Args:
        dataset_arn: ARN of the dataset
        timeout: Maximum time to wait in seconds
        check_interval: Time between checks in seconds
        
    Returns:
        bool: True if deleted, False if timed out
    """
    start_time = time.time()
    logger.info(f"Waiting for dataset deletion to complete: {dataset_arn}")
    
    while time.time() - start_time < timeout:
        try:
            response = rekognition.describe_dataset(DatasetArn=dataset_arn)
            status = response.get('Status')
            
            if status == 'DELETING':
                logger.info(f"Dataset deletion in progress. Waiting {check_interval} seconds...")
                time.sleep(check_interval)
                continue
            else:
                logger.warning(f"Dataset not in DELETING state, current status: {status}")
                return False
                
        except ClientError as e:
            if e.response['Error']['Code'] == 'ResourceNotFoundException':
                logger.info(f"Dataset successfully deleted: {dataset_arn}")
                return True
            else:
                logger.error(f"Error checking dataset status: {str(e)}")
                return False
    
    logger.warning(f"Timed out waiting for dataset deletion: {dataset_arn}")
    return False

def fix_project_arn(project_arn):
    """
    Ensure project ARN includes the project ID.
    
    Args:
        project_arn: Project ARN to fix
        
    Returns:
        str: Fixed project ARN
    """
    if not project_arn:
        return None
        
    # Use the known project ID
    project_id = "1725357683732"
    
    # Check if ARN already ends with project ID
    if re.search(r'/[0-9]+$', project_arn):
        if not project_arn.endswith(project_id):
            # Replace existing ID with correct one
            base_arn = project_arn.rsplit('/', 1)[0]
            return f"{base_arn}/{project_id}"
        return project_arn
    
    # Add project ID if not present
    if project_arn.endswith('/'):
        return f"{project_arn}{project_id}"
    else:
        return f"{project_arn}/{project_id}"

def lambda_handler(event, context):
    """
    Lambda function to delete Rekognition datasets.
    Simple dedicated function that only handles dataset deletion.
    
    Args:
        event (dict): Lambda event with dataset info
        context: Lambda context
        
    Returns:
        dict: Result of the deletion operation
    """
    try:
        logger.info(f"Received event: {json.dumps(event)}")
        
        # Get project ARN from event or environment
        project_arn = event.get('project_arn') or os.environ.get('PROJECT_ARN')
        if not project_arn:
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'message': 'PROJECT_ARN is required',
                    'error': 'Missing required parameter'
                })
            }
            
        # Ensure project ARN is properly formatted
        project_arn = fix_project_arn(project_arn)
        logger.info(f"Using project ARN: {project_arn}")
        
        # Get dataset type(s) to delete
        dataset_types = event.get('dataset_types')
        specific_type = event.get('dataset_type')
        
        # Allow single type or list of types
        if specific_type and not dataset_types:
            dataset_types = [specific_type]
        elif not dataset_types:
            dataset_types = ['TRAIN', 'TEST']  # Default to both types
            
        logger.info(f"Deleting dataset types: {dataset_types}")
        
        # Handle force operation (last resort)
        if event.get('force'):
            logger.info("Force flag is set, will attempt more aggressive deletion if regular methods fail")
            
        # Get project details to find actual dataset ARNs
        project_details = get_project_details(project_arn)
        if not project_details:
            # If project details not found, try console approach as a last resort
            if event.get('force'):
                logger.info("Project details not found, attempting direct deletion using project name")
                
                # Extract project components
                match = re.match(r'arn:aws:rekognition:([^:]+):([^:]+):project/([^/]+)', project_arn)
                if match:
                    region, account, project_name = match.groups()
                    
                    # Log how to manually delete through console
                    console_url = f"https://{region}.console.aws.amazon.com/rekognition/custom-labels?region={region}#/projects/{project_name}/dataset"
                    logger.info(f"Console URL for manual deletion if needed: {console_url}")
                    
                    # Dump raw deletion attempt for emergency purposes
                    for dataset_type in dataset_types:
                        ds_type = dataset_type.lower()
                        try:
                            # Direct emergency approach - handles edge cases
                            logger.info(f"Attempting emergency deletion for {ds_type} dataset")
                            # Try various deletion methods 
                            for attempt_method in ["deleteDataset", "DeleteDataset"]:
                                try:
                                    # Handle different Rekognition API versions
                                    response = rekognition._make_api_call(
                                        attempt_method,
                                        {
                                            "ProjectArn": project_arn,
                                            "DatasetType": ds_type.upper()
                                        }
                                    )
                                    logger.info(f"Emergency deletion successful: {response}")
                                    break
                                except Exception as method_e:
                                    logger.warning(f"Emergency method {attempt_method} failed: {str(method_e)}")
                        except Exception as e:
                            logger.error(f"Emergency deletion failed: {str(e)}")
                    
                    # Return partial success
                    return {
                        'statusCode': 202,  # Accepted
                        'body': json.dumps({
                            'message': 'Attempted emergency dataset deletion',
                            'project_arn': project_arn,
                            'console_url': console_url
                        })
                    }
            
            return {
                'statusCode': 404,
                'body': json.dumps({
                    'message': 'Project not found',
                    'error': 'Project does not exist or could not be accessed',
                    'project_arn': project_arn
                })
            }
            
        logger.info(f"Found project: {project_details.get('ProjectArn')}")
        
        # Process delete operations for each dataset
        results = {}
        
        # Check if project has datasets
        if 'Datasets' not in project_details or not project_details['Datasets']:
            logger.info("No datasets found in project")
            return {
                'statusCode': 200,
                'body': json.dumps({
                    'message': 'No datasets found to delete',
                    'project_arn': project_arn
                })
            }
            
        # Get dataset ARNs from project details
        datasets = project_details['Datasets']
        logger.info(f"Found {len(datasets)} datasets in project")
        
        # Delete each requested dataset type
        for dataset in datasets:
            dataset_type = dataset.get('DatasetType')
            dataset_arn = dataset.get('DatasetArn')
            status = dataset.get('Status', 'UNKNOWN')
            
            if not dataset_type or not dataset_arn:
                continue
                
            if dataset_type in dataset_types:
                logger.info(f"Processing {dataset_type} dataset: {dataset_arn} (Status: {status})")
                
                # Handle datasets in transitional states
                if status in ['CREATING', 'UPDATING', 'DELETING']:
                    logger.info(f"Dataset in {status} state, waiting before deletion...")
                    time.sleep(10)  # Wait for state to stabilize
                
                # Delete the dataset
                deletion_result = delete_dataset(dataset_arn)
                results[dataset_type] = deletion_result
                
                # Wait for deletion to complete
                if deletion_result.get('success'):
                    wait_for_dataset_deletion(dataset_arn)
        
        # Check if we actually deleted anything
        if not results:
            return {
                'statusCode': 404,
                'body': json.dumps({
                    'message': 'No matching datasets found to delete',
                    'project_arn': project_arn,
                    'requested_types': dataset_types
                })
            }
            
        # Determine overall success
        all_success = all(result.get('success', False) for result in results.values())
        status_code = 200 if all_success else 500
        
        return {
            'statusCode': status_code,
            'body': json.dumps({
                'message': 'Dataset deletion completed',
                'all_success': all_success,
                'results': results,
                'project_arn': project_arn
            })
        }
    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: {str(e)}',
                'error': 'Internal server error'
            })
        }