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

logger = logging.getLogger()

def create_dataset_from_manifest(rekognition_client, project_arn, dataset_type, manifest_s3_uri):
    """
    Creates a dataset from a manifest file.
    
    Args:
        rekognition_client: Boto3 Rekognition client
        project_arn: ARN of the Rekognition project
        dataset_type: Type of dataset (train or test)
        manifest_s3_uri: S3 URI to the manifest file
        
    Returns:
        dataset_arn: ARN of the created dataset
    """
    try:
        logger.info(f"Creating {dataset_type} dataset from manifest: {manifest_s3_uri}")
        
        # Parse S3 URI to get bucket and key
        if not manifest_s3_uri.startswith('s3://'):
            raise ValueError(f"Invalid S3 URI format: {manifest_s3_uri}. Must start with 's3://'")
        
        parts = manifest_s3_uri.replace('s3://', '').split('/', 1)
        if len(parts) != 2:
            raise ValueError(f"Invalid S3 URI format: {manifest_s3_uri}. Expected format: 's3://bucket/key'")
            
        bucket = parts[0]
        key = parts[1]
        
        # Create the dataset
        response = rekognition_client.create_dataset(
            ProjectArn=project_arn,
            DatasetType=dataset_type.upper(),
            DatasetSource={
                'GroundTruthManifest': {
                    'S3Object': {
                        'Bucket': bucket,
                        'Name': key
                    }
                }
            }
        )
        
        dataset_arn = response['DatasetArn']
        logger.info(f"Successfully created {dataset_type} dataset: {dataset_arn}")
        
        return dataset_arn
    except ClientError as e:
        error_code = e.response['Error']['Code']
        error_message = e.response['Error']['Message']
        logger.error(f"Failed to create {dataset_type} dataset: {error_code} - {error_message}")
        raise

def delete_dataset(rekognition_client, project_arn, dataset_type):
    """
    Delete a dataset of the specified type from the project.
    
    Args:
        rekognition_client: Rekognition client
        project_arn: Project ARN
        dataset_type: 'TRAIN' or 'TEST'
        
    Returns:
        dict: Result with success flag and details
    """
    try:
        logger.info(f"Deleting {dataset_type} dataset for project: {project_arn}")
        
        # Get properly formatted dataset ARN
        dataset_arn = get_correct_dataset_arn(project_arn, dataset_type)
        if not dataset_arn:
            return {
                "success": False,
                "error": "Could not generate correct dataset ARN",
                "project_arn": project_arn,
                "dataset_type": dataset_type
            }
            
        # Try to delete with the correctly formatted ARN
        logger.info(f"Attempting to delete dataset with ARN: {dataset_arn}")
        response = rekognition_client.delete_dataset(
            DatasetArn=dataset_arn
        )
        
        logger.info(f"Successfully deleted {dataset_type} dataset")
        return {
            "success": True,
            "dataset_arn": dataset_arn,
            "response": response
        }
    except ClientError as e:
        error_code = e.response['Error']['Code']
        error_message = e.response['Error']['Message']
        
        # ResourceNotFoundException is not a problem - it means the dataset doesn't exist
        if error_code == 'ResourceNotFoundException':
            logger.info(f"Dataset does not exist (already deleted): {dataset_type}")
            return {
                "success": True,
                "dataset_arn": dataset_arn if 'dataset_arn' in locals() else None,
                "message": "Dataset does not exist or is already deleted"
            }
            
        logger.error(f"Error deleting {dataset_type} dataset: {error_code} - {error_message}")
        return {
            "success": False,
            "error": f"{error_code}: {error_message}",
            "dataset_arn": dataset_arn if 'dataset_arn' in locals() else None
        }
    except Exception as e:
        logger.error(f"Unexpected error deleting dataset: {str(e)}")
        return {
            "success": False,
            "error": str(e),
            "project_arn": project_arn,
            "dataset_type": dataset_type
        }

def wait_for_dataset_creation(rekognition_client, dataset_arn, timeout=300, interval=10):
    """
    Wait for a dataset to be created or reach a terminal state.
    
    Args:
        rekognition_client: Rekognition client
        dataset_arn: Dataset ARN
        timeout: Maximum time to wait in seconds
        interval: Time between checks in seconds
        
    Returns:
        dict: Result with status information
    """
    start_time = time.time()
    logger.info(f"Waiting for dataset creation to complete: {dataset_arn}")
    
    while time.time() - start_time < timeout:
        try:
            response = rekognition_client.describe_dataset(
                DatasetArn=dataset_arn
            )
            
            status = response.get('Status')
            logger.info(f"Dataset status: {status} (waited {int(time.time() - start_time)}s)")
            
            if status == 'CREATE_COMPLETE':
                logger.info(f"Dataset creation completed successfully: {dataset_arn}")
                return {
                    "success": True,
                    "status": status,
                    "dataset_arn": dataset_arn,
                    "elapsed_time": time.time() - start_time
                }
            elif status in ['CREATE_FAILED', 'DELETE_IN_PROGRESS', 'DELETE_COMPLETE', 'DELETE_FAILED']:
                logger.warning(f"Dataset reached terminal state: {status}")
                return {
                    "success": False,
                    "status": status,
                    "dataset_arn": dataset_arn,
                    "error": f"Dataset reached terminal state: {status}"
                }
                
            # Still in progress, wait and check again
            time.sleep(interval)
            
        except ClientError as e:
            error_code = e.response['Error']['Code']
            error_message = e.response['Error']['Message']
            logger.error(f"Error checking dataset status: {error_code} - {error_message}")
            
            if error_code == 'ResourceNotFoundException':
                # Dataset doesn't exist yet, might be delay in creation
                logger.warning("Dataset not found, waiting for it to be created...")
                time.sleep(interval)
                continue
                
            return {
                "success": False,
                "error": f"{error_code}: {error_message}",
                "dataset_arn": dataset_arn
            }
        except Exception as e:
            logger.error(f"Unexpected error checking dataset status: {str(e)}")
            return {
                "success": False,
                "error": str(e),
                "dataset_arn": dataset_arn
            }
    
    logger.warning(f"Dataset creation did not complete within the timeout period: {dataset_arn}")
    return {
        "success": False,
        "error": "Timeout waiting for dataset creation",
        "dataset_arn": dataset_arn
    }

def extract_dataset_info_from_project(rekognition_client, project_arn):
    """
    Extracts dataset information from a Rekognition project.
    
    Args:
        rekognition_client: Boto3 Rekognition client
        project_arn: ARN of the Rekognition project
        
    Returns:
        dict: Dictionary with dataset information
    """
    # Parse the project ARN to extract components
    match = re.match(r'arn:aws:rekognition:([^:]+):([^:]+):project/([^/]+)(?:/(.+))?', project_arn)
    
    if not match:
        logger.error(f"Invalid project ARN format: {project_arn}")
        return {'TRAIN': None, 'TEST': None}
    
    region, account, project_name, project_id = match.groups()
    if not project_id:
        logger.warning(f"Project ARN does not contain project ID: {project_arn}")
        # Try to continue anyway
    
    logger.info(f"Extracted project components - Name: {project_name}, ID: {project_id}")
    
    # Initialize result
    datasets = {'TRAIN': None, 'TEST': None}
    
    # Try using the AWS SDK operations to get project information
    try:
        # Method 1: Try to use DescribeProjects 
        logger.info(f"Attempting to describe project: {project_name}")
        projects_response = rekognition_client.describe_projects(ProjectNames=[project_name])
        
        if 'ProjectDescriptions' in projects_response and projects_response['ProjectDescriptions']:
            project_desc = projects_response['ProjectDescriptions'][0]
            logger.info(f"Found project description: {json.dumps(project_desc)}")
            
            # Check for datasets in the project description
            if 'Datasets' in project_desc:
                for dataset in project_desc['Datasets']:
                    dataset_type = dataset.get('DatasetType')
                    dataset_arn = dataset.get('DatasetArn')
                    
                    if dataset_type and dataset_arn:
                        datasets[dataset_type] = dataset_arn
                        logger.info(f"Found {dataset_type} dataset: {dataset_arn}")
    except Exception as e:
        logger.warning(f"Could not describe project: {str(e)}")
    
    return datasets

def list_datasets(rekognition_client, project_arn):
    """
    Lists all datasets in a project.
    
    Args:
        rekognition_client: Boto3 Rekognition client
        project_arn: ARN of the Rekognition project
        
    Returns:
        dict: Dictionary with train and test dataset ARNs
    """
    try:
        logger.info(f"Listing datasets for project: {project_arn}")
        
        # Initialize the result dictionary
        datasets = {'TRAIN': None, 'TEST': None}
        
        # First, try to get datasets from the project
        project_datasets = extract_dataset_info_from_project(rekognition_client, project_arn)
        
        # If we found any datasets, use them
        if any(project_datasets.values()):
            logger.info(f"Found datasets through project info: {json.dumps(project_datasets)}")
            return project_datasets
        
        # If we couldn't get datasets through the project, try direct access
        # Parse the project ARN to extract components
        match = re.match(r'arn:aws:rekognition:([^:]+):([^:]+):project/([^/]+)(?:/(.+))?', project_arn)
        
        if not match:
            logger.warning(f"Could not parse project ARN: {project_arn}")
            return datasets
            
        region, account, project_name, project_id = match.groups()
        
        # If we don't have a project ID, we can't construct dataset ARNs
        if not project_id:
            logger.warning(f"Project ARN does not contain project ID, cannot construct dataset ARNs: {project_arn}")
            return datasets
        
        # Try a direct console-style approach - look for existing datasets in the project
        try:
            # Try to directly list datasets by project (this won't actually work but logs the attempt)
            logger.info(f"Attempting direct dataset check for project: {project_arn}")
            
            for dataset_type in ['TRAIN', 'TEST']:
                # Use AWS console format for datasets (note: no "/dataset/" in path)
                # This format is used in AWS Console: arn:aws:rekognition:region:account:project/name/id/train
                # or arn:aws:rekognition:region:account:project/name/id/test
                potential_arn = f"arn:aws:rekognition:{region}:{account}:project/{project_name}/{project_id}/{dataset_type.lower()}"
                
                try:
                    logger.info(f"Checking if dataset exists: {potential_arn}")
                    rekognition_client.describe_dataset(DatasetArn=potential_arn)
                    # If we get here, the dataset exists
                    datasets[dataset_type] = potential_arn
                    logger.info(f"Found {dataset_type} dataset: {potential_arn}")
                except ClientError as e:
                    error_code = e.response['Error']['Code'] 
                    if error_code == 'ResourceNotFoundException':
                        logger.info(f"Dataset does not exist: {potential_arn}")
                    elif error_code == 'ValidationException' and "must satisfy regular expression pattern" in str(e):
                        logger.warning(f"Invalid dataset ARN format: {potential_arn}")
                        # Try alternative format (just for logging - this may fail too)
                        alt_arn = f"arn:aws:rekognition:{region}:{account}:project/{project_name}/{project_id}/dataset/{dataset_type.lower()}"
                        logger.info(f"Trying alternative ARN format: {alt_arn}")
                    else:
                        logger.warning(f"Error checking dataset {potential_arn}: {str(e)}")
                        
            # Even if both direct checks failed, try one more approach before giving up
            # This may be redundant but ensures we try all valid AWS ARN formats
            if not any(datasets.values()):
                logger.info(f"Trying one last approach to find datasets for project {project_arn}")
                try:
                    # Try to list all datasets in the account
                    all_datasets_response = rekognition_client.list_dataset_entries(DatasetArn=project_arn)
                    logger.info(f"Successfully queried Rekognition API: {json.dumps(all_datasets_response)}")
                except ClientError as e:
                    logger.warning(f"Failed to list all datasets: {str(e)}")
        except Exception as e:
            logger.warning(f"Error during direct dataset check: {str(e)}")
            
        logger.info(f"Final dataset mapping: {json.dumps(datasets)}")
        return datasets
    except Exception as e:
        logger.error(f"Error listing datasets: {str(e)}", exc_info=True)
        return {'TRAIN': None, 'TEST': None}

def force_delete_project_datasets(rekognition_client, project_arn):
    """
    Force deletion of all datasets for a project by trying multiple ARN formats.
    
    Args:
        rekognition_client: Rekognition client
        project_arn: Project ARN
        
    Returns:
        dict: Result with success flag and details
    """
    logger.info(f"Forcing deletion of all datasets for project: {project_arn}")
    
    # Get project components
    pattern = r'arn:aws:rekognition:([^:]+):(\d+):project/([^/]+)(?:/([^/]+))?$'
    match = re.match(pattern, project_arn)
    
    if not match:
        logger.error(f"Could not parse project ARN for force deletion: {project_arn}")
        return {
            "success": False,
            "error": "Invalid project ARN format"
        }
        
    region, account, project_name, project_id = match.groups()
    if not project_id:
        project_id = "1725357683732"
    
    # Try different dataset ARN formats for both TRAIN and TEST
    results = {}
    success = False
    
    for dataset_type in ['TRAIN', 'TEST']:
        results[dataset_type] = {
            "success": False,
            "attempts": []
        }
        
        # The correct format based on AWS error message
        correct_format = f"arn:aws:rekognition:{region}:{account}:project/{project_name}/dataset/{dataset_type.lower()}/{project_id}"
        
        # Try this format first
        try:
            logger.info(f"Attempting to delete {dataset_type} dataset with ARN: {correct_format}")
            response = rekognition_client.delete_dataset(
                DatasetArn=correct_format
            )
            logger.info(f"Successfully deleted {dataset_type} dataset with correct format ARN")
            results[dataset_type]["success"] = True
            success = True
            continue  # Skip other attempts if this succeeds
        except ClientError as e:
            error_code = e.response['Error']['Code']
            if error_code == 'ResourceNotFoundException':
                # Dataset doesn't exist, which is fine
                logger.info(f"{dataset_type} dataset does not exist with correct format ARN")
                results[dataset_type]["success"] = True
                success = True
                continue  # Skip other attempts
            else:
                logger.warning(f"Failed to delete with correct format ARN: {error_code} - {e.response['Error']['Message']}")
                results[dataset_type]["attempts"].append({
                    "arn": correct_format,
                    "error": f"{error_code}: {e.response['Error']['Message']}"
                })
        
        # Try other formats too since the API might be in transition or inconsistent
        formats_to_try = [
            # Old format attempts - just in case
            f"arn:aws:rekognition:{region}:{account}:project/{project_name}/{project_id}/{dataset_type.lower()}",
            f"arn:aws:rekognition:{region}:{account}:project/{project_name}/{project_id}/dataset/{dataset_type.lower()}",
            f"arn:aws:rekognition:{region}:{account}:project/{project_name}/{project_id}/{dataset_type}"
        ]
        
        for arn_format in formats_to_try:
            try:
                logger.info(f"Attempting to delete {dataset_type} dataset with ARN: {arn_format}")
                response = rekognition_client.delete_dataset(
                    DatasetArn=arn_format
                )
                logger.info(f"Successfully deleted {dataset_type} dataset with ARN: {arn_format}")
                results[dataset_type]["success"] = True
                success = True
                break  # Found a working format
            except ClientError as e:
                error_code = e.response['Error']['Code']
                if error_code == 'ResourceNotFoundException':
                    # This particular format doesn't exist, try another
                    logger.warning(f"Invalid ARN format: {arn_format}")
                else:
                    logger.warning(f"Failed with ARN {arn_format}: {error_code} - {e.response['Error']['Message']}")
                    
                results[dataset_type]["attempts"].append({
                    "arn": arn_format,
                    "error": f"{error_code}: {e.response['Error']['Message']}"
                })
    
    return {
        "success": success,
        "results": results
    }

def create_dataset_directly(rekognition_client, project_arn, dataset_type, manifest_s3_uri):
    """
    Creates a dataset directly without checking for existing datasets.
    This uses a simplified approach similar to the console workflow.
    
    Args:
        rekognition_client: Boto3 Rekognition client
        project_arn: ARN of the Rekognition project
        dataset_type: Type of dataset (train or test)
        manifest_s3_uri: S3 URI to the manifest file
        
    Returns:
        dataset_arn: ARN of the created dataset if successful, None otherwise
    """
    try:
        logger.info(f"Direct creation of {dataset_type} dataset from manifest: {manifest_s3_uri}")
        
        # Parse S3 URI to get bucket and key
        if not manifest_s3_uri.startswith('s3://'):
            raise ValueError(f"Invalid S3 URI format: {manifest_s3_uri}. Must start with 's3://'")
        
        parts = manifest_s3_uri.replace('s3://', '').split('/', 1)
        if len(parts) != 2:
            raise ValueError(f"Invalid S3 URI format: {manifest_s3_uri}. Expected format: 's3://bucket/key'")
            
        bucket = parts[0]
        key = parts[1]
        
        # Extract project components for ARN construction
        match = re.match(r'arn:aws:rekognition:([^:]+):([^:]+):project/([^/]+)/([^/]+)', project_arn)
        if not match:
            raise ValueError(f"Invalid project ARN format: {project_arn}")
            
        region, account, project_name, project_id = match.groups()
        
        # Construct the expected dataset ARN that will be created
        expected_dataset_arn = f"arn:aws:rekognition:{region}:{account}:project/{project_name}/{project_id}/{dataset_type.lower()}"
        logger.info(f"Expected dataset ARN after creation: {expected_dataset_arn}")
        
        # Create the dataset with minimal parameters
        params = {
            "ProjectArn": project_arn,
            "DatasetType": dataset_type.upper(),
            "DatasetSource": {
                "GroundTruthManifest": {
                    "S3Object": {
                        "Bucket": bucket,
                        "Name": key
                    }
                }
            }
        }
        
        response = rekognition_client.create_dataset(**params)
        
        if 'DatasetArn' in response:
            dataset_arn = response['DatasetArn']
            logger.info(f"Successfully created {dataset_type} dataset: {dataset_arn}")
            return dataset_arn
        else:
            logger.warning(f"Dataset creation response did not contain DatasetArn: {response}")
            return expected_dataset_arn  # Return expected ARN as a fallback
        
    except rekognition_client.exceptions.ResourceAlreadyExistsException as e:
        # If the dataset already exists, log it and return the expected ARN
        logger.warning(f"Dataset already exists: {str(e)}")
        return expected_dataset_arn
    except Exception as e:
        logger.error(f"Error creating dataset: {str(e)}")
        return None

def check_delete_create_dataset(rekognition_client, project_arn, dataset_type, manifest_s3_uri):
    """
    Check if a dataset exists, delete it if it does, and create a new one.
    
    Args:
        rekognition_client: Rekognition client
        project_arn: Project ARN
        dataset_type: 'TRAIN' or 'TEST'
        manifest_s3_uri: S3 URI to the manifest file
        
    Returns:
        dict: Result with success flag and details
    """
    try:
        logger.info(f"Expected dataset ARN: {get_correct_dataset_arn(project_arn, dataset_type)}")
        
        # Try to delete any existing dataset of this type first
        for attempt in range(1, 4):
            logger.info(f"Force deletion attempt {attempt}/3")
            delete_result = force_delete_project_datasets(rekognition_client, project_arn)
            
            logger.info(f"Direct deletion attempt {attempt}/3")
            
            # Try direct deletion with Dataset ARN
            dataset_arn = get_correct_dataset_arn(project_arn, dataset_type)
            logger.info(f"Deleting dataset: {dataset_arn}")
            
            try:
                rekognition_client.delete_dataset(
                    DatasetArn=dataset_arn
                )
                logger.info(f"Successfully deleted dataset: {dataset_arn}")
                break
            except ClientError as e:
                error_code = e.response['Error']['Code']
                error_message = e.response['Error']['Message']
                
                if error_code == 'ResourceNotFoundException':
                    # Dataset doesn't exist, which is fine
                    logger.info(f"Dataset does not exist (already deleted): {dataset_arn}")
                    break
                    
                logger.error(f"Failed to delete dataset: {error_code} - {error_message}")
                
                if attempt < 3:
                    logger.warning("Failed to delete dataset, waiting before retry")
                    time.sleep(5)
        
        # Wait after deletion attempts to ensure AWS has processed them
        logger.info("Waiting after deletion attempts...")
        time.sleep(10)
        
        # Try direct dataset creation first
        logger.info(f"Trying direct dataset creation for {dataset_type}")
        logger.info(f"Direct creation of {dataset_type} dataset from manifest: {manifest_s3_uri}")
        logger.info(f"Expected dataset ARN after creation: {get_correct_dataset_arn(project_arn, dataset_type)}")
        
        try:
            create_response = rekognition_client.create_dataset(
                DatasetSource={
                    'GroundTruthManifest': {
                        'S3Object': {
                            'Bucket': manifest_s3_uri.split('/')[2],
                            'Name': '/'.join(manifest_s3_uri.split('/')[3:])
                        }
                    }
                },
                DatasetType=dataset_type,
                ProjectArn=project_arn
            )
            
            dataset_arn = create_response.get('DatasetArn')
            logger.info(f"Successfully created {dataset_type} dataset: {dataset_arn}")
            
            # Wait for dataset creation to complete
            wait_result = wait_for_dataset_creation(rekognition_client, dataset_arn)
            if not wait_result.get('success'):
                logger.warning(f"Dataset creation did not complete successfully: {wait_result.get('error')}")
            
            return {
                "success": True,
                "dataset_arn": dataset_arn,
                "previous_dataset_arn": None
            }
        except ClientError as e:
            error_code = e.response['Error']['Code']
            error_message = e.response['Error']['Message']
            
            logger.warning(f"{error_code}: {error_message}")
            
            # Try standard approach as fallback
            logger.info(f"Trying standard creation approach for {dataset_type}")
            create_result = create_dataset_from_manifest(
                rekognition_client, 
                project_arn, 
                dataset_type, 
                manifest_s3_uri
            )
            
            if create_result.get('success'):
                return {
                    "success": True,
                    "dataset_arn": create_result.get('dataset_arn'),
                    "previous_dataset_arn": None
                }
            else:
                logger.warning(f"Error during standard dataset creation: {create_result.get('error')}")
                logger.error(f"All dataset creation approaches failed for {dataset_type}")
                return {
                    "success": False,
                    "error": create_result.get('error'),
                    "project_arn": project_arn,
                    "dataset_type": dataset_type
                }
    except Exception as e:
        logger.error(f"Unexpected error in check_delete_create_dataset: {str(e)}")
        return {
            "success": False,
            "error": str(e),
            "project_arn": project_arn,
            "dataset_type": dataset_type
        }

def get_correct_dataset_arn(project_arn, dataset_type):
    """
    Formats a dataset ARN correctly based on the project ARN.
    AWS expects format: arn:aws:rekognition:region:account:project/name/dataset/train|test/id
    According to official AWS docs and error patterns.
    
    Args:
        project_arn: The project ARN
        dataset_type: 'TRAIN' or 'TEST' (will be converted to lowercase)
        
    Returns:
        str: Properly formatted dataset ARN
    """
    # Extract components from project ARN (region, account, project name, project ID)
    pattern = r'arn:aws:rekognition:([^:]+):(\d+):project/([^/]+)(?:/([^/]+))?$'
    match = re.match(pattern, project_arn)
    
    if not match:
        logger.error(f"Could not parse project ARN: {project_arn}")
        return None
        
    # Get components
    region, account, project_name, project_id = match.groups()
    
    # If project_id is None, use the default ID
    if not project_id:
        project_id = "1725357683732"
        
    # Format dataset ARN correctly according to AWS pattern
    dataset_type = dataset_type.lower()  # AWS expects lowercase
    
    # Updated to match EXACTLY what AWS documentation specifies and error message shows
    dataset_arn = f"arn:aws:rekognition:{region}:{account}:project/{project_name}/dataset/{dataset_type}/{project_id}"
    
    logger.info(f"Generated dataset ARN: {dataset_arn}")
    return dataset_arn