import boto3
import logging
import os
import json
import xml.etree.ElementTree as ET
from datetime import datetime
import urllib.parse

from utils.rekognition_utils import validate_manifest_file
from utils.logging_utils import setup_logger, reduce_logging_verbosity
from utils.manifest_utils import convert_xml_to_rekognition_format, validate_manifest_entry

# Initialize logger
logger = setup_logger()
reduce_logging_verbosity()

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

def get_bucket_name(param_name):
    """Get bucket name from SSM parameter store"""
    try:
        response = ssm_client.get_parameter(Name=param_name)
        return response['Parameter']['Value']
    except Exception as e:
        logger.error(f"Error getting parameter {param_name}: {str(e)}")
        return None

def process_xml_file(source_bucket, xml_key):
    """Process XML file and extract annotations"""
    try:
        logger.info(f"Processing XML file: {xml_key}")
        
        # URL decode the XML key
        decoded_xml_key = urllib.parse.unquote_plus(xml_key)
        logger.info(f"Decoded XML file path: {decoded_xml_key}")
        
        # Get the XML content
        response = s3_client.get_object(Bucket=source_bucket, Key=decoded_xml_key)
        xml_content = response['Body'].read()
        
        # Parse XML
        root = ET.fromstring(xml_content)
        
        # Get filename without extension
        base_name = os.path.splitext(os.path.basename(decoded_xml_key))[0]
        
        # Assume jpg file exists with the same name as the XML file but with .jpg extension
        xml_dir = os.path.dirname(decoded_xml_key)
        jpg_filename = f"{base_name}.jpg"
        jpg_key = f"{xml_dir}/{jpg_filename}" if xml_dir else jpg_filename
        
        logger.info(f"Using JPG file at: {jpg_key}")
        
        # Build the entry for later processing
        entry = {
            'xml_content': xml_content,
            'source_bucket': source_bucket,
            'image_key': jpg_key
        }
        
        return entry
        
    except Exception as e:
        logger.error(f"Error processing XML file {xml_key}: {str(e)}")
        return None

def lambda_handler(event, context):
    """Lambda handler for creating and validating the training manifest"""
    try:
        # Log the incoming event
        logger.info(f"Received event: {json.dumps(event)}")
        
        xml_key = None
        source_bucket = None
        
        # Handle S3 event notification
        if 'Records' in event and len(event['Records']) > 0:
            s3_record = event['Records'][0]['s3']
            logger.info(f"Processing S3 event: {json.dumps(s3_record)}")
            
            # Extract bucket and key information
            source_bucket = s3_record['bucket']['name']
            xml_key = s3_record['object']['key']
            
            # Only process XML files
            if not xml_key.lower().endswith('.xml'):
                logger.info(f"Ignoring non-XML file: {xml_key}")
                return {
                    'statusCode': 200,
                    'body': json.dumps({
                        'message': 'Skipped non-XML file'
                    })
                }
            
            logger.info(f"Processing XML file from S3 event: {source_bucket}/{xml_key}")
        else:
            # Get bucket names from parameter store if not in event
            source_bucket = get_bucket_name('/datafy-rekognition-stack/source-bucket')
            
        training_bucket = get_bucket_name('/datafy-rekognition-stack/training-bucket')
        
        if not source_bucket or not training_bucket:
            missing = []
            if not source_bucket: missing.append('source-bucket')
            if not training_bucket: missing.append('training-bucket')
            raise ValueError(f"Missing required parameters: {', '.join(missing)}")
            
        logger.info(f"Using source bucket: {source_bucket}")
        logger.info(f"Using training bucket: {training_bucket}")
        
        # Process the XML file and update manifest
        try:
            manifest_key = "manifest.jsonl"
            
            # If triggered by event, process only the triggering file
            if xml_key:
                # Process the XML file
                entry = process_xml_file(source_bucket, xml_key)
                if not entry:
                    return {
                        'statusCode': 400,
                        'body': json.dumps({
                            'message': f'Failed to process XML file: {xml_key}',
                            'error': 'Invalid XML format'
                        })
                    }
                
                # Convert entry to Rekognition format
                rekognition_entry = convert_xml_to_rekognition_format(
                    entry['xml_content'],
                    entry['source_bucket'],
                    entry['image_key']
                )
                
                if not rekognition_entry:
                    return {
                        'statusCode': 400,
                        'body': json.dumps({
                            'message': f'Failed to convert XML to Rekognition format: {xml_key}',
                            'error': 'Invalid XML format or content'
                        })
                    }
                
                # Check if manifest exists
                try:
                    manifest_exists = True
                    try:
                        s3_client.head_object(Bucket=training_bucket, Key=manifest_key)
                    except:
                        manifest_exists = False
                    
                    # Validate the new entry before attempting to add it to the manifest
                    if not validate_manifest_entry(rekognition_entry):
                        logger.error(f"Entry validation failed for {xml_key}")
                        return {
                            'statusCode': 400,
                            'body': json.dumps({
                                'message': f'Entry validation failed for {xml_key}',
                                'error': 'Invalid entry format or content'
                            })
                        }
                    
                    # Check if source-ref already exists in manifest to avoid duplicates
                    source_ref = rekognition_entry.get('source-ref')
                    has_duplicate = False
                    
                    if manifest_exists:
                        # Read existing manifest
                        response = s3_client.get_object(Bucket=training_bucket, Key=manifest_key)
                        manifest_content = response['Body'].read().decode('utf-8')
                        
                        # Check for duplicates and get count of entries
                        lines = [line for line in manifest_content.strip().split('\n') if line.strip()]
                        existing_entries = []
                        
                        for line in lines:
                            try:
                                entry = json.loads(line)
                                if entry.get('source-ref') == source_ref:
                                    has_duplicate = True
                                    logger.info(f"Found duplicate entry for {source_ref}, will replace it")
                                else:
                                    existing_entries.append(entry)
                            except json.JSONDecodeError:
                                logger.warning(f"Skipping invalid JSON line in manifest")
                        
                        # Add new entry to the list
                        existing_entries.append(rekognition_entry)
                        
                        # Convert to JSONL format
                        manifest_content = '\n'.join(json.dumps(e, separators=(',', ':')) for e in existing_entries)
                        if not manifest_content.endswith('\n'):
                            manifest_content += '\n'
                    else:
                        # Create new manifest with single entry
                        manifest_content = json.dumps(rekognition_entry, separators=(',', ':')) + '\n'
                    
                    # Upload manifest with correct content type
                    s3_client.put_object(
                        Bucket=training_bucket,
                        Key=manifest_key,
                        Body=manifest_content,
                        ContentType='application/x-amazon-s3-object-manifest-jsonl'
                    )
                    
                    # Verify content type is correct
                    response = s3_client.get_object(Bucket=training_bucket, Key=manifest_key)
                    content_type = response.get('ContentType', 'unknown')
                    
                    # If content type is wrong, fix it with a copy operation
                    if content_type != 'application/x-amazon-s3-object-manifest-jsonl':
                        logger.warning(f"Incorrect content type detected: {content_type}, fixing...")
                        s3_client.copy_object(
                            CopySource={'Bucket': training_bucket, 'Key': manifest_key},
                            Bucket=training_bucket,
                            Key=manifest_key,
                            ContentType='application/x-amazon-s3-object-manifest-jsonl',
                            MetadataDirective='REPLACE'
                        )
                    
                    # Count entries in final manifest
                    entry_count = len(manifest_content.strip().split('\n'))
                    
                    action_type = "updated" if has_duplicate else "added to"
                    
                    return {
                        'statusCode': 200,
                        'body': json.dumps({
                            'message': f'Entry {action_type} manifest successfully',
                            'entry_count': entry_count,
                            'entry_source_ref': source_ref,
                            'manifest_file': f"s3://{training_bucket}/{manifest_key}"
                        })
                    }
                    
                except Exception as e:
                    logger.error(f"Error updating manifest: {str(e)}")
                    raise
            else:
                # No specific file triggered - create manifest from all files (full rebuild)
                logger.info("No specific file triggered - rebuilding entire manifest")
                
                # List all XML files in the source bucket
                xml_files = []
                paginator = s3_client.get_paginator('list_objects_v2')
                for page in paginator.paginate(Bucket=source_bucket, Prefix=''):
                    for obj in page.get('Contents', []):
                        if obj['Key'].endswith('.xml'):
                            xml_files.append(obj['Key'])
                
                logger.info(f"Found {len(xml_files)} XML files in bucket")
                
                if not xml_files:
                    return {
                        'statusCode': 400,
                        'body': json.dumps({
                            'message': 'No XML files found',
                            'error': 'No training data available'
                        })
                    }
                
                # Process each XML file
                valid_entries = []
                for xml_key in xml_files:
                    try:
                        entry = process_xml_file(source_bucket, xml_key)
                        if entry:
                            valid_entries.append(entry)
                    except Exception as e:
                        logger.error(f"Error processing XML file {xml_key}: {str(e)}")
                        continue
                
                if not valid_entries:
                    return {
                        'statusCode': 400,
                        'body': json.dumps({
                            'message': 'No valid entries found',
                            'error': 'No valid training data available'
                        })
                    }
                
                # Convert entries to Rekognition format
                rekognition_entries = []
                for entry in valid_entries:
                    try:
                        rekognition_entry = convert_xml_to_rekognition_format(
                            entry['xml_content'],
                            entry['source_bucket'],
                            entry['image_key']
                        )
                        if rekognition_entry:
                            rekognition_entries.append(rekognition_entry)
                    except Exception as e:
                        logger.error(f"Error converting entry to Rekognition format: {str(e)}")
                        continue
                
                if not rekognition_entries:
                    return {
                        'statusCode': 400,
                        'body': json.dumps({
                            'message': 'No valid Rekognition entries',
                            'error': 'Failed to convert entries to Rekognition format'
                        })
                    }
                
                # Create manifest file
                manifest_content = '\n'.join(json.dumps(entry) for entry in rekognition_entries)
                
                # Upload manifest with correct content type
                s3_client.put_object(
                    Bucket=training_bucket,
                    Key=manifest_key,
                    Body=manifest_content,
                    ContentType='application/x-amazon-s3-object-manifest-jsonl'
                )
                
                # Verify manifest is valid
                if not validate_manifest_file(training_bucket, manifest_key):
                    return {
                        'statusCode': 400,
                        'body': json.dumps({
                            'message': 'Manifest validation failed',
                            'error': 'Invalid manifest format or content'
                        })
                    }
                
                return {
                    'statusCode': 200,
                    'body': json.dumps({
                        'message': 'Manifest created successfully',
                        'entry_count': len(rekognition_entries),
                        'manifest_file': f"s3://{training_bucket}/{manifest_key}"
                    })
                }
                
        except Exception as e:
            logger.error(f"Error processing training data: {str(e)}")
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'message': 'Failed to process training data',
                    'error': str(e)
                })
            }
            
    except ValueError as e:
        logger.error(f"Configuration error: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({
                'message': 'Configuration error',
                'error': str(e)
            })
        }