import boto3
import logging
import os
import json
import time
import re
import subprocess
from datetime import datetime
from typing import Optional
from botocore.exceptions import ClientError

# Update imports to use dataset_utils for dataset operations
from utils.rekognition_utils import validate_manifest_file
from utils.logging_utils import setup_logger, reduce_logging_verbosity
from utils.dataset_utils import create_dataset_from_manifest

# Initialize logger
logger = setup_logger()
reduce_logging_verbosity()

# Initialize clients
ssm_client = boto3.client('ssm')
s3_client = boto3.client('s3')

def force_dict(result):
    """Ensure result is a dictionary with success and dataset_arn keys"""
    if isinstance(result, dict):
        return result
    if isinstance(result, str):
        # If result is just the ARN string
        logger.info(f"Converting string result to dict: {result}")
        return {
            "success": True,
            "dataset_arn": result
        }
    # For any other type, return a failed dictionary
    logger.error(f"Unknown result type: {type(result)}, value: {result}")
    return {
        "success": False,
        "error": f"Unknown result type: {type(result)}",
        "value": str(result)
    }

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 lambda_handler(event, context):
    """
    Lambda function to create datasets for AWS Rekognition Custom Labels.
    Focused only on dataset creation - deletion is handled by a separate function.
    
    Args:
        event (dict): Lambda event object containing optional parameters
        context (object): Lambda context
        
    Returns:
        dict: Operation result with status and details
    """
    try:
        logger.info(f"Received event: {json.dumps(event)}")
        
        # Get environment variables
        training_bucket = os.environ.get('TRAINING_BUCKET')
        project_arn = os.environ.get('PROJECT_ARN')
        
        # Try to get from SSM if not in environment
        if not training_bucket:
            training_bucket = get_ssm_parameter('/datafy-rekognition-stack/training-bucket')
            if not training_bucket:
                return {
                    'statusCode': 500,
                    'body': json.dumps({
                        'message': 'TRAINING_BUCKET not found in environment or SSM',
                        'error': 'Missing required configuration'
                    })
                }
            logger.info(f"Got training bucket from SSM: {training_bucket}")
        else:
            logger.info(f"Using TRAINING_BUCKET from environment: {training_bucket}")
        
        if not project_arn:
            project_arn = get_ssm_parameter('/datafy-rekognition-stack/project-name')
            if project_arn:
                # 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_arn}"
                logger.info(f"Constructed project ARN from SSM: {project_arn}")
            else:
                return {
                    'statusCode': 500,
                    'body': json.dumps({
                        'message': 'PROJECT_ARN not found in environment or SSM',
                        'error': 'Missing required configuration'
                    })
                }
        
        # Fix project ARN format
        project_arn = fix_project_arn(project_arn)
        if not project_arn:
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'message': 'Invalid project ARN format',
                    'error': 'Configuration error'
                })
            }
        logger.info(f"Using project ARN: {project_arn}")
        
        # Initialize AWS clients
        rekognition = boto3.client('rekognition')
        
        # For deletion operations, redirect to the specialized dataset deletion function
        if event.get('operation') == 'delete_dataset' or event.get('operation') == 'force_delete_dataset':
            # Call the dedicated deletion function
            lambda_client = boto3.client('lambda')
            
            logger.info(f"Forwarding deletion request to dedicated deletion function")
            
            try:
                deletion_payload = {
                    'project_arn': project_arn,
                    'dataset_type': event.get('dataset_type'),
                    'force': event.get('operation') == 'force_delete_dataset'
                }
                
                # Get function name from environment or use default
                delete_function_name = f"{os.environ.get('PROJECT_NAME', 'datafy-rekognition')}-delete-datasets"
                logger.info(f"Invoking deletion function: {delete_function_name}")
                
                response = lambda_client.invoke(
                    FunctionName=delete_function_name,
                    InvocationType='RequestResponse',
                    Payload=json.dumps(deletion_payload)
                )
                
                # Read and parse the response
                payload = json.loads(response['Payload'].read())
                logger.info(f"Deletion function response: {payload}")
                
                return {
                    'statusCode': 200,
                    'body': json.dumps({
                        'message': 'Dataset deletion request processed',
                        'deletion_result': payload
                    })
                }
            except Exception as e:
                logger.error(f"Error forwarding to deletion function: {str(e)}")
                return {
                    'statusCode': 500,
                    'body': json.dumps({
                        'message': f"Error forwarding deletion request: {str(e)}",
                        'error': 'Internal server error'
                    })
                }
        
        # Use default manifest path for creation
        base_manifest_key = "manifest.jsonl"
        train_manifest_key = "manifest_train.jsonl"
        test_manifest_key = "manifest_test.jsonl"
        
        # Check if split manifest files exist - more robust check
        split_manifests_exist = False
        
        logger.info(f"Checking for split manifest files in bucket {training_bucket}")
        
        # Try listing objects instead of head_object for more reliability
        try:
            # List objects in the bucket with the manifest prefix
            response = s3_client.list_objects_v2(
                Bucket=training_bucket,
                Prefix=""  # Empty prefix to list all objects
            )
            
            if 'Contents' in response:
                # Get the keys of all objects in the bucket
                all_keys = [obj['Key'] for obj in response['Contents']]
                logger.info(f"Found {len(all_keys)} objects in bucket, checking for manifest files")
                logger.info(f"Available files: {', '.join(all_keys)}")
                
                # Check if both manifest files exist
                if train_manifest_key in all_keys and test_manifest_key in all_keys:
                    logger.info(f"✅ Both split manifest files found!")
                    split_manifests_exist = True
                else:
                    if train_manifest_key not in all_keys:
                        logger.warning(f"❌ Train manifest file {train_manifest_key} not found")
                    if test_manifest_key not in all_keys:
                        logger.warning(f"❌ Test manifest file {test_manifest_key} not found")
            else:
                logger.warning(f"No files found in bucket {training_bucket}")
        except Exception as e:
            logger.error(f"Error listing objects in bucket: {str(e)}")
        
        # Set up manifest URIs based on whether split files exist or not
        if split_manifests_exist:
            train_manifest_s3_uri = f"s3://{training_bucket}/{train_manifest_key}"
            test_manifest_s3_uri = f"s3://{training_bucket}/{test_manifest_key}"
            logger.info(f"Using split manifest files: Train: {train_manifest_s3_uri}, Test: {test_manifest_s3_uri}")
        else:
            # Fall back to using the same manifest file for both if split files don't exist
            train_manifest_s3_uri = f"s3://{training_bucket}/{base_manifest_key}"
            test_manifest_s3_uri = train_manifest_s3_uri
            logger.info(f"Using same manifest file for both train and test: {train_manifest_s3_uri}")
        
        # Parse S3 URI for train manifest
        train_bucket_name = train_manifest_s3_uri.split('/')[2]
        train_key = '/'.join(train_manifest_s3_uri.split('/')[3:])
        
        # Validate the train manifest file
        logger.info(f"Validating train manifest file: {train_manifest_s3_uri}")
        if not validate_manifest_file(train_bucket_name, train_key):
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'message': 'Train manifest validation failed',
                    'error': 'Invalid manifest format or content',
                    'manifestFile': train_manifest_s3_uri
                })
            }
        
        logger.info(f"Train manifest file validated successfully: {train_manifest_s3_uri}")
        
        # Create train dataset
        train_result = create_dataset_from_manifest(
            rekognition,
            project_arn,
            'TRAIN',
            train_manifest_s3_uri
        )
        
        # Convert result to dictionary if needed
        train_result = force_dict(train_result)
        
        if not train_result.get('success'):
            logger.error(f"Failed to create TRAIN dataset: {train_result}")
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'message': 'Failed to create TRAIN dataset',
                    'datasetArn': train_result.get('dataset_arn'),
                    'error': train_result.get('error', 'Unknown error'),
                    'manifestFile': train_manifest_s3_uri
                })
            }
        
        # Create test dataset with the test manifest file
        test_result = create_dataset_from_manifest(
            rekognition,
            project_arn,
            'TEST',
            test_manifest_s3_uri
        )
        
        # Convert result to dictionary if needed
        test_result = force_dict(test_result)
        
        # Even if test dataset creation fails, we can still return a success for the train dataset
        test_success = test_result.get('success', False)
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Dataset creation completed',
                'trainDatasetArn': train_result.get('dataset_arn'),
                'testDatasetArn': test_result.get('dataset_arn'),
                'testSuccess': test_success,
                'trainManifestFile': train_manifest_s3_uri,
                'testManifestFile': test_manifest_s3_uri,
                'usedSplitManifests': split_manifests_exist
            })
        }
        
    except Exception as e:
        logger.error(f"Error in lambda_handler: {str(e)}", exc_info=True)
        return {
            'statusCode': 500,
            'body': json.dumps({
                'message': f'Internal error: {str(e)}',
                'error': 'Unexpected exception occurred'
            })
        }