import boto3
import logging
import os
import json
import time
import random
from collections import defaultdict
from botocore.exceptions import ClientError

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

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

# Get training bucket from environment variable
TRAINING_BUCKET = os.environ.get('TRAINING_BUCKET')
MANIFEST_PREFIX = ''  # Empty prefix to look in the root of the bucket
MANIFEST_FILENAME = 'manifest.jsonl'  # Exact filename to look for
TRAIN_RATIO = 0.8  # Fixed 80/20 split ratio

def get_latest_manifest():
    """
    Find the manifest file in the training bucket.
    
    Returns:
        tuple: (bucket, key) of the manifest file
    """
    try:
        if not TRAINING_BUCKET:
            logger.error("TRAINING_BUCKET environment variable not set")
            raise ValueError("Missing TRAINING_BUCKET environment variable")
            
        logger.info(f"Looking for manifest file in s3://{TRAINING_BUCKET}/")
        
        # Check if the manifest file exists directly
        manifest_key = MANIFEST_FILENAME
        try:
            # Try to check if the file exists
            s3_client.head_object(Bucket=TRAINING_BUCKET, Key=manifest_key)
            logger.info(f"Found manifest file: s3://{TRAINING_BUCKET}/{manifest_key}")
            return TRAINING_BUCKET, manifest_key
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                logger.warning(f"Manifest file {manifest_key} not found in root, trying general search")
            else:
                # For other errors, reraise
                raise
        
        # If the exact file wasn't found, try listing objects to find manifest files
        response = s3_client.list_objects_v2(
            Bucket=TRAINING_BUCKET,
            Prefix=MANIFEST_PREFIX
        )
        
        if 'Contents' not in response or not response['Contents']:
            logger.error(f"No files found in s3://{TRAINING_BUCKET}/")
            raise FileNotFoundError(f"No files found in bucket")
            
        # Find the latest .jsonl file (sorting by LastModified)
        manifest_files = [
            obj for obj in response['Contents'] 
            if obj['Key'].endswith('.jsonl') and not (obj['Key'].endswith('_train.jsonl') or obj['Key'].endswith('_test.jsonl'))
        ]
        
        if not manifest_files:
            logger.error(f"No .jsonl manifest files found in s3://{TRAINING_BUCKET}/")
            raise FileNotFoundError(f"No .jsonl manifest files found")
            
        # Sort by last modified (newest first)
        latest_manifest = sorted(manifest_files, key=lambda x: x['LastModified'], reverse=True)[0]
        manifest_key = latest_manifest['Key']
        
        logger.info(f"Found manifest file: s3://{TRAINING_BUCKET}/{manifest_key}")
        return TRAINING_BUCKET, manifest_key
    
    except Exception as e:
        logger.error(f"Error finding manifest file: {str(e)}")
        raise

def split_manifest_file(bucket: str, key: str, train_ratio: float = 0.8) -> dict:
    """
    Split the manifest file into train and test datasets.
    The split is done by folder, with train_ratio percent of images from each folder
    going to the training dataset and the rest to the test dataset.
    
    Args:
        bucket: S3 bucket containing the manifest file
        key: S3 key for the manifest file
        train_ratio: Ratio of data to use for training (default 0.8 for 80/20 split)
        
    Returns:
        dict: Dictionary with keys 'train_key' and 'test_key' pointing to the created manifest files
    """
    try:
        logger.info(f"Splitting manifest file s3://{bucket}/{key} with {train_ratio:.0%}/{1-train_ratio:.0%} split")
        
        # Get the manifest file
        response = s3_client.get_object(Bucket=bucket, Key=key)
        content = response['Body'].read().decode('utf-8')
        lines = content.splitlines()
        
        # Group entries by folder
        folder_entries = defaultdict(list)
        for line in lines:
            if not line.strip():
                continue
                
            try:
                entry = json.loads(line)
                source_ref = entry.get('source-ref', '')
                
                # Extract folder path from S3 URL
                if source_ref.startswith('s3://'):
                    # Remove s3:// and bucket, then split by /
                    path_parts = source_ref.replace('s3://', '').split('/', 1)[1].split('/')
                    # Use parent folder as grouping key
                    if len(path_parts) > 1:
                        folder = '/'.join(path_parts[:-1])
                    else:
                        folder = 'root'  # Default for files at the bucket root
                    
                    folder_entries[folder].append(line)
            except Exception as e:
                logger.warning(f"Error parsing manifest entry: {str(e)}")
        
        # Initialize train and test lists
        train_entries = []
        test_entries = []
        
        # Split each folder's entries
        for folder, entries in folder_entries.items():
            # Shuffle to ensure random selection
            random.shuffle(entries)
            
            # Calculate split point
            split_idx = int(len(entries) * train_ratio)
            
            # Split into train and test
            folder_train = entries[:split_idx]
            folder_test = entries[split_idx:]
            
            logger.info(f"Folder '{folder}': {len(folder_train)} train, {len(folder_test)} test")
            
            # Add to respective lists
            train_entries.extend(folder_train)
            test_entries.extend(folder_test)
        
        # Generate the new file names
        base_name = os.path.splitext(key)[0]
        train_key = f"{base_name}_train.jsonl"
        test_key = f"{base_name}_test.jsonl"
        
        # Create the new manifest files
        content_type = 'application/x-amazon-s3-object-manifest-jsonl'
        
        # Write train manifest
        s3_client.put_object(
            Bucket=bucket,
            Key=train_key,
            Body='\n'.join(train_entries),
            ContentType=content_type
        )
        
        # Write test manifest
        s3_client.put_object(
            Bucket=bucket,
            Key=test_key,
            Body='\n'.join(test_entries),
            ContentType=content_type
        )
        
        logger.info(f"Created train manifest s3://{bucket}/{train_key} with {len(train_entries)} entries")
        logger.info(f"Created test manifest s3://{bucket}/{test_key} with {len(test_entries)} entries")
        
        return {
            'train_key': train_key,
            'test_key': test_key,
            'train_count': len(train_entries),
            'test_count': len(test_entries)
        }
    except Exception as e:
        logger.error(f"Error splitting manifest file: {str(e)}")
        raise

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 lambda_handler(event, context):
    """
    Lambda handler for splitting manifest files.
    
    Automatically finds the latest manifest file in the training bucket
    and splits it into train and test datasets with 80/20 split ratio.
    
    Returns:
        dict: Information about the created train and test manifest files
    """
    try:
        logger.info(f"Received event: {json.dumps(event)}")
        
        # Find the latest manifest file in the training bucket
        try:
            bucket, key = get_latest_manifest()
        except Exception as e:
            logger.error(f"Failed to find latest manifest file: {str(e)}")
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'message': 'Failed to find latest manifest file',
                    'error': str(e)
                })
            }
        
        # Validate the manifest file
        if not validate_manifest_file(bucket, key):
            logger.error(f"Manifest validation failed for s3://{bucket}/{key}")
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'message': 'Manifest validation failed',
                    'error': 'Invalid manifest format or content',
                    'manifestFile': f"s3://{bucket}/{key}"
                })
            }
            
        # Split the manifest file (using fixed 0.8 train ratio)
        result = split_manifest_file(bucket, key, TRAIN_RATIO)
        
        # Generate S3 URIs for the created files
        train_s3_uri = f"s3://{bucket}/{result['train_key']}"
        test_s3_uri = f"s3://{bucket}/{result['test_key']}"
        
        logger.info(f"Split completed: {result['train_count']} train, {result['test_count']} test")
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Manifest file successfully split',
                'trainManifest': train_s3_uri,
                'testManifest': test_s3_uri,
                'trainCount': result['train_count'],
                'testCount': result['test_count'],
                'originalManifest': f"s3://{bucket}/{key}"
            })
        }
    
    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'
            })
        }