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

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')

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_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_client.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 get_dataset_status(project_arn, dataset_type):
    """
    Check the status of a dataset
    
    Args:
        project_arn: ARN of the project
        dataset_type: 'TRAIN' or 'TEST'
        
    Returns:
        dict: Information about the dataset status
    """
    try:
        # Standardize dataset type to uppercase for comparison
        dataset_type_upper = dataset_type.upper()
        
        # Get project details including datasets
        project_details = get_project_details(project_arn)
        
        if not project_details or 'Datasets' not in project_details or not project_details['Datasets']:
            logger.info(f"No datasets found in project: {project_arn}")
            return {
                "exists": False,
                "status": "NOT_FOUND",
                "is_ready": False
            }
        
        # Find the matching dataset by type
        matching_datasets = [d for d in project_details['Datasets'] 
                            if d.get('DatasetType') == dataset_type_upper]
        
        if not matching_datasets:
            logger.info(f"{dataset_type} dataset does not exist in project")
            return {
                "exists": False,
                "status": "NOT_FOUND",
                "is_ready": False
            }
        
        # Use the actual dataset information returned by the API
        dataset = matching_datasets[0]
        dataset_arn = dataset.get('DatasetArn')
        status = dataset.get('Status', 'UNKNOWN')
        
        logger.info(f"Found {dataset_type} dataset ARN: {dataset_arn}")
        logger.info(f"{dataset_type} dataset status: {status}")
        
        return {
            "exists": True,
            "status": status,
            "dataset_arn": dataset_arn,
            "is_ready": status == 'CREATE_COMPLETE'
        }
        
    except ClientError as e:
        error_code = e.response['Error']['Code']
        error_message = e.response['Error']['Message']
        
        if error_code == 'ResourceNotFoundException':
            logger.info(f"{dataset_type} dataset does not exist")
            return {
                "exists": False,
                "status": "NOT_FOUND",
                "is_ready": False
            }
            
        logger.error(f"Error checking {dataset_type} dataset: {error_code} - {error_message}")
        return {
            "exists": False,
            "status": "ERROR",
            "error": f"{error_code}: {error_message}",
            "is_ready": False
        }
        
    except Exception as e:
        logger.error(f"Unexpected error checking {dataset_type} dataset: {str(e)}")
        return {
            "exists": False,
            "status": "ERROR",
            "error": str(e),
            "is_ready": False
        }

def lambda_handler(event, context):
    """
    Lambda handler to check if datasets have been created
    
    Args:
        event: Lambda event
        context: Lambda context
        
    Returns:
        dict: Response with datasets_ready flag for Step Functions
    """
    try:
        logger.info(f"Starting check of dataset status")
        
        # Get environment variables
        project_arn = os.environ.get('PROJECT_ARN')
        
        if not project_arn:
            missing = ['PROJECT_ARN']
            
            # Try to get from SSM if not in environment
            if 'PROJECT_ARN' in missing:
                project_name = get_ssm_parameter('/datafy-rekognition-stack/project-name')
                if project_name:
                    # Construct the ARN using the project name with fallback values if needed
                    aws_region = os.environ.get('AWS_REGION', 'eu-west-1')
                    aws_account_id = os.environ.get('AWS_ACCOUNT_ID', '587594388832')
                    project_arn = f"arn:aws:rekognition:{aws_region}:{aws_account_id}:project/{project_name}"
                    logger.info(f"Constructed project ARN from SSM: {project_arn}")
                    missing.remove('PROJECT_ARN')
            
            if missing:
                logger.warning(f"Missing required configuration: {', '.join(missing)}")
        
        # Fix project ARN format
        project_arn = fix_project_arn(project_arn)
        if not project_arn:
            logger.warning("Failed to fix project ARN, will try alternatives")
            project_arn = None
        else:
            logger.info(f"Using fixed project ARN: {project_arn}")
        
        # If we still don't have a working ARN, try using the describe_projects method
        if not project_arn:
            # Try different variations of the project name if provided
            project_names_to_try = []
            
            if project_name:
                # Original project name
                project_names_to_try.append(project_name)
                
                # Try without hyphens
                if '-' in project_name:
                    project_names_to_try.append(project_name.replace('-', ''))
                
                # Try with "Datafy" appended if not already present
                if not project_name.endswith('Datafy'):
                    project_names_to_try.append(f"{project_name}Datafy")
                    project_names_to_try.append("Datafy")  # Just try "Datafy" alone
            else:
                # Fallbacks if no project name provided
                project_names_to_try = ["Datafy", "datafy-rekognition", "datafyrekognition"]
                logger.warning(f"No project name available, will try these names: {project_names_to_try}")
            
            # Try to find the project using the names
            project_found = False
            
            for name in project_names_to_try:
                try:
                    logger.info(f"Trying to find project with name: {name}")
                    project_response = rekognition_client.describe_projects(ProjectNames=[name])
                    
                    if project_response.get('ProjectDescriptions'):
                        raw_project_arn = project_response['ProjectDescriptions'][0]['ProjectArn']
                        logger.info(f"Found project ARN: {raw_project_arn}")
                        
                        # Fix the ARN format
                        project_arn = fix_project_arn(raw_project_arn)
                        if project_arn:
                            logger.info(f"Using fixed project ARN: {project_arn}")
                            project_found = True
                            break
                except Exception as e:
                    logger.warning(f"Could not find project with name '{name}': {str(e)}")
                    continue
            
            # If no project found through names, try listing all projects
            if not project_found:
                try:
                    logger.info("No project found by name, listing all projects")
                    list_response = rekognition_client.describe_projects()
                    
                    if list_response.get('ProjectDescriptions'):
                        # Log all available projects to help with debugging
                        for idx, project in enumerate(list_response['ProjectDescriptions']):
                            logger.info(f"Available project {idx+1}: {project.get('ProjectName')} - {project.get('ProjectArn')}")
                        
                        # Use the first project found
                        if list_response['ProjectDescriptions']:
                            raw_project_arn = list_response['ProjectDescriptions'][0]['ProjectArn']
                            logger.info(f"Using first available project ARN: {raw_project_arn}")
                            
                            # Fix the ARN format
                            project_arn = fix_project_arn(raw_project_arn)
                            if project_arn:
                                logger.info(f"Using fixed project ARN: {project_arn}")
                                project_found = True
                except Exception as e:
                    logger.error(f"Error listing projects: {str(e)}")
        
        if not project_arn:
            logger.error("Could not find any Rekognition projects")
            return {
                'statusCode': 404,
                'body': json.dumps({
                    'message': 'No Rekognition projects found',
                    'error': 'Project not found'
                }),
                'datasets_ready': False
            }
        
        # Check status of both datasets
        train_dataset = get_dataset_status(project_arn, 'TRAIN')
        test_dataset = get_dataset_status(project_arn, 'TEST')
        
        # Both datasets must exist and be in CREATE_COMPLETE status
        datasets_ready = train_dataset.get('is_ready') and test_dataset.get('is_ready')
        
        logger.info(f"Datasets ready: {datasets_ready}")
        logger.info(f"Train dataset: {train_dataset}")
        logger.info(f"Test dataset: {test_dataset}")
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Dataset status check completed',
                'train_dataset': train_dataset,
                'test_dataset': test_dataset
            }),
            'datasets_ready': datasets_ready
        }
        
    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 dataset status: {str(e)}',
                'error': 'Internal server error'
            }),
            'datasets_ready': False
        }