import boto3
import logging
import os
import json
from datetime import datetime
from typing import Dict, Any, Optional

# Initialize logger
logger = logging.getLogger(__name__)

# Initialize clients
rekognition_client = boto3.client('rekognition')
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 start_rekognition_training(
    training_bucket: str,
    version_name: Optional[str] = None,
    project_arn: Optional[str] = None
) -> Dict[str, Any]:
    """
    Start a Rekognition Custom Labels training job using the minimal parameters approach.
    
    This function will use the datasets already created in the AWS Rekognition Console
    by providing only the ProjectArn, VersionName and OutputConfig parameters, which is
    the approach that successfully works with console-created datasets.
    """
    try:
        # Get project ARN from environment if not provided
        if not project_arn:
            project_arn = os.environ.get('PROJECT_ARN')
            if not project_arn:
                raise ValueError("PROJECT_ARN not provided or found in environment")
                
        # Generate version name if not provided
        if not version_name:
            version_name = f"version-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
            
        logger.info(f"Starting training for project: {project_arn}")
        logger.info(f"Version name: {version_name}")
        logger.info(f"Training bucket: {training_bucket}")
        
        # Use the minimal parameters approach that we know works
        minimal_params = {
            "ProjectArn": project_arn,
            "VersionName": version_name,
            "OutputConfig": {
                "S3Bucket": training_bucket,
                "S3KeyPrefix": f"training_output/{version_name}"
            }
        }
        
        logger.info("Starting training with minimal parameters (no dataset references)")
        response = rekognition_client.create_project_version(**minimal_params)
        logger.info(f"Training job started successfully with minimal parameters: {json.dumps(response, default=str)}")
        return response
    
    except rekognition_client.exceptions.InvalidParameterException as e:
        error_message = str(e)
        # Log the full error message to help with debugging
        logger.error(f"InvalidParameterException: {error_message}")
        
        if "project already has associated datasets" in error_message.lower():
            logger.error("Project has datasets but the API call failed. Please ensure your project is properly set up in the console.")
            raise ValueError("Project configuration error. Please check that your training dataset is properly set up in the AWS console.")
        elif "test dataset doesn't exist" in error_message.lower():
            logger.error("Failed to create test dataset automatically. Please set up a test dataset in the AWS console.")
            raise ValueError("Test dataset creation failed. Please go to the AWS Rekognition Console and either upload a test dataset or split your training dataset.")
        else:
            # Re-raise the original exception for other errors
            logger.error(f"Error starting training: {str(e)}")
            raise
    except Exception as e:
        logger.error(f"Error starting training: {str(e)}")
        raise

def check_training_job_status(project_arn: str, version_name: str) -> Dict[str, Any]:
    """
    Check the status of a training job.
    
    Args:
        project_arn: The project ARN
        version_name: The version name to check
        
    Returns:
        Dictionary containing status information
    """
    try:
        # Get version details
        response = rekognition_client.describe_project_versions(
            ProjectArn=project_arn,
            VersionNames=[version_name]
        )
        
        if not response['ProjectVersionDescriptions']:
            raise ValueError(f"Version {version_name} not found")
            
        version_info = response['ProjectVersionDescriptions'][0]
        
        # Extract relevant information
        status_info = {
            'version_name': version_name,
            'status': version_info.get('Status'),
            'status_message': version_info.get('StatusMessage'),
            'created': version_info.get('CreationTimestamp'),
            'evaluation': version_info.get('EvaluationResults', {})
        }
        
        logger.info(f"Training status for {version_name}: {status_info['status']}")
        if status_info['status_message']:
            logger.info(f"Status message: {status_info['status_message']}")
            
        return status_info
        
    except Exception as e:
        logger.error(f"Error checking training status: {str(e)}")
        raise