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

logger = logging.getLogger(__name__)

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

def validate_manifest_file(bucket: str, key: str) -> bool:
    """
    Validate manifest file format and content.
    Returns True if valid, False if not.
    """
    try:
        # Check if manifest exists
        try:
            response = s3_client.get_object(Bucket=bucket, Key=key)
        except Exception as e:
            logger.error(f"Manifest file not found: {str(e)}")
            return False
            
        content = response['Body'].read().decode('utf-8')
        
        # Check content type
        content_type = response.get('ContentType', '')
        expected_type = 'application/x-amazon-s3-object-manifest-jsonl'
        
        if content_type != expected_type:
            logger.warning(f"Incorrect content type: {content_type}, updating to {expected_type}")
            # Update the content type
            s3_client.copy_object(
                CopySource={'Bucket': bucket, 'Key': key},
                Bucket=bucket,
                Key=key,
                ContentType=expected_type,
                MetadataDirective='REPLACE'
            )
            logger.info(f"Updated content type to {expected_type}")
        
        # Validate each line is valid JSON and has required fields
        valid_count = 0
        invalid_count = 0
        lines = content.splitlines()
        
        if not lines:
            logger.error("Manifest file is empty")
            return False
            
        # Log first few entries for debugging
        logger.info(f"First manifest entry sample:")
        try:
            first_entry = json.loads(lines[0])
            logger.info(json.dumps(first_entry, indent=2))
        except Exception as e:
            logger.error(f"Error parsing first entry: {str(e)}")
            
        for line in lines:
            if not line.strip():
                continue
                
            try:
                entry = json.loads(line)
                if all(k in entry for k in ['source-ref', 'bounding-box', 'bounding-box-metadata']):
                    valid_count += 1
                else:
                    invalid_count += 1
                    logger.warning(f"Invalid entry missing required fields: {line[:200]}...")
            except Exception as e:
                invalid_count += 1
                logger.warning(f"Invalid JSON in line: {str(e)}")
                
        logger.info(f"Manifest validation results:")
        logger.info(f"Total lines: {len(lines)}")
        logger.info(f"Valid entries: {valid_count}")
        logger.info(f"Invalid entries: {invalid_count}")
        
        return valid_count > 0 and invalid_count == 0
        
    except Exception as e:
        logger.error(f"Error validating manifest: {str(e)}")
        return False

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]
        
        # Validate the manifest file
        if not validate_manifest_file(bucket, key):
            raise ValueError(f"Manifest file is not valid: {manifest_s3_uri}")
        
        # 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, dataset_arn):
    """
    Deletes a dataset.
    
    Args:
        rekognition_client: Boto3 Rekognition client
        dataset_arn: ARN of the dataset to delete
        
    Returns:
        bool: True if deletion was successful
    """
    try:
        logger.info(f"Deleting dataset: {dataset_arn}")
        rekognition_client.delete_dataset(DatasetArn=dataset_arn)
        logger.info(f"Successfully deleted dataset: {dataset_arn}")
        return True
    except ClientError as e:
        error_code = e.response['Error']['Code']
        error_message = e.response['Error']['Message']
        
        # Don't fail if the dataset doesn't exist
        if error_code == 'ResourceNotFoundException':
            logger.warning(f"Dataset {dataset_arn} not found, continuing")
            return True
        
        logger.error(f"Failed to delete dataset: {error_code} - {error_message}")
        raise