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

logger = logging.getLogger(__name__)

def get_rekognition_client():
    """Get Rekognition client"""
    return boto3.client('rekognition')

def find_project_arn(project_arn: str) -> str:
    """
    Find the actual project ARN based on the provided ARN.
    
    Args:
        project_arn: The project ARN to find
        
    Returns:
        The actual project ARN
        
    Raises:
        ValueError: If project is not found
    """
    try:
        # Get Rekognition client
        rekognition_client = get_rekognition_client()
        
        # Extract project name from ARN
        project_name = project_arn.split('/')[-1]
        logger.info(f"Looking for project: {project_name}")
        
        # Get actual project ARN
        projects = rekognition_client.describe_projects()
        actual_project_arn = None
        
        for project in projects.get('ProjectDescriptions', []):
            if project_name in project['ProjectArn']:
                actual_project_arn = project['ProjectArn']
                break
        
        if not actual_project_arn:
            raise ValueError(f"Project not found: {project_name}")
            
        logger.info(f"Found project ARN: {actual_project_arn}")
        return actual_project_arn
        
    except Exception as e:
        logger.error(f"Error finding project ARN: {str(e)}")
        raise

def create_dataset(project_arn: str, training_bucket: str, manifest_key: str, dataset_type: str = 'TRAIN') -> str:
    """
    Create a dataset in the Rekognition Custom Labels project using SageMaker Ground Truth manifest.
    
    Args:
        project_arn: The project ARN
        training_bucket: The S3 bucket containing the manifest
        manifest_key: The manifest file key
        dataset_type: Type of dataset ('TRAIN' or 'TEST')
        
    Returns:
        The dataset ARN
    """
    try:
        # Create Rekognition client
        rekognition_client = get_rekognition_client()
        
        # Generate a dataset name with timestamp and type
        timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        dataset_name = f"dataset-{dataset_type.lower()}-{timestamp}"
        
        logger.info(f"Creating {dataset_type} dataset: {dataset_name}")
        
        # Create the dataset using SageMaker Ground Truth format
        response = rekognition_client.create_dataset(
            ProjectArn=project_arn,
            DatasetType=dataset_type,
            DatasetSource={
                'GroundTruthManifest': {
                    'S3Object': {
                        'Bucket': training_bucket,
                        'Name': manifest_key
                    }
                }
            }
        )
        
        dataset_arn = response['DatasetArn']
        logger.info(f"Created {dataset_type} dataset with ARN: {dataset_arn}")
        
        # Wait for dataset creation to complete
        while True:
            try:
                status_response = rekognition_client.describe_dataset(
                    DatasetArn=dataset_arn
                )
                
                status = status_response['DatasetDescription']['Status']
                logger.info(f"Dataset creation status: {status}")
                
                if status == 'CREATE_COMPLETE':
                    break
                elif status == 'CREATE_FAILED':
                    error_message = status_response['DatasetDescription'].get('StatusMessage', 'Unknown error')
                    raise ValueError(f"Dataset creation failed: {error_message}")
                elif status == 'CREATE_IN_PROGRESS':
                    time.sleep(5)  # Wait 5 seconds before checking again
                else:
                    raise ValueError(f"Unexpected dataset status: {status}")
            except rekognition_client.exceptions.ResourceNotFoundException:
                logger.warning("Dataset not found yet, waiting...")
                time.sleep(5)  # Wait before retrying
            except Exception as e:
                if 'ValidationException' in str(e):
                    logger.error(f"Validation error during dataset status check: {str(e)}")
                    time.sleep(5)
                    continue
                raise
        
        return dataset_arn
        
    except Exception as e:
        logger.error(f"Error creating dataset: {str(e)}")
        raise

def evaluate_model(project_arn: str, version_name: str, min_f1_score: float = 0.85) -> bool:
    """
    Evaluate the model using F1 score.
    
    Args:
        project_arn: The project ARN
        version_name: The version name to evaluate
        min_f1_score: Minimum required F1 score (default 0.85)
        
    Returns:
        True if model meets evaluation criteria, False otherwise
    """
    try:
        status_info = check_training_job_status(project_arn, version_name)
        
        if 'evaluation' in status_info:
            f1_score = status_info['evaluation'].get('F1Score', 0.0)
            logger.info(f"Model F1 score: {f1_score}")
            return f1_score >= min_f1_score
        
        return False
        
    except Exception as e:
        logger.error(f"Error evaluating model: {str(e)}")
        return False

def start_model_inference(project_version_arn: str, min_inference_units: int = 1) -> None:
    """
    Start the model for inference.
    
    Args:
        project_version_arn: The model version ARN to start
        min_inference_units: Minimum inference units (default 1)
    """
    try:
        rekognition_client = get_rekognition_client()
        
        # Start the model
        rekognition_client.start_project_version(
            ProjectVersionArn=project_version_arn,
            MinInferenceUnits=min_inference_units
        )
        
        logger.info(f"Started model version: {project_version_arn}")
        
        # Wait for model to be running
        while True:
            response = rekognition_client.describe_project_versions(
                ProjectArn=project_version_arn.split('/version/')[0],
                VersionNames=[project_version_arn.split('/')[-1]]
            )
            
            if response['ProjectVersionDescriptions']:
                status = response['ProjectVersionDescriptions'][0]['Status']
                logger.info(f"Model status: {status}")
                
                if status == 'RUNNING':
                    break
                elif status == 'FAILED':
                    raise ValueError("Model failed to start")
                
            time.sleep(10)  # Check every 10 seconds
            
    except Exception as e:
        logger.error(f"Error starting model: {str(e)}")
        raise

def validate_manifest_file(bucket: str, key: str) -> bool:
    """
    Validate a manifest file in S3.
    Returns True if valid, False if not.
    """
    try:
        s3_client = boto3.client('s3')
        
        # Get the manifest file
        response = s3_client.get_object(Bucket=bucket, Key=key)
        content = response['Body'].read().decode('utf-8')
        
        # Log first few lines for debugging
        lines = content.splitlines()
        if lines:
            logger.info("First manifest entry sample:")
            try:
                first_entry = json.loads(lines[0])
                logger.info(json.dumps(first_entry, indent=2))
                
                # Log specific details about the first entry
                if 'source-ref' in first_entry:
                    logger.info(f"Source-ref format: {first_entry['source-ref']}")
                if 'bounding-box' in first_entry:
                    bb = first_entry['bounding-box']
                    logger.info(f"Number of annotations: {len(bb.get('annotations', []))}")
                    if 'annotations' in bb and bb['annotations']:
                        logger.info(f"First annotation sample: {json.dumps(bb['annotations'][0], indent=2)}")
                if 'bounding-box-metadata' in first_entry:
                    bbm = first_entry['bounding-box-metadata']
                    logger.info(f"Class map: {json.dumps(bbm.get('class-map', {}), indent=2)}")
            except Exception as e:
                logger.error(f"Error parsing first entry: {str(e)}")
        
        # Check if content type needs to be updated
        try:
            current_content_type = response.get('ContentType', 'unknown')
            logger.info(f"Current manifest content type: {current_content_type}")
            
            # Update content type if needed
            if current_content_type != 'application/x-amazon-s3-object-manifest-jsonl':
                logger.warning(f"Updating content type from {current_content_type} to application/x-amazon-s3-object-manifest-jsonl")
                s3_client.copy_object(
                    CopySource={'Bucket': bucket, 'Key': key},
                    Bucket=bucket,
                    Key=key,
                    ContentType='application/x-amazon-s3-object-manifest-jsonl',
                    MetadataDirective='REPLACE'
                )
                
                # Verify content type was updated
                verify_response = s3_client.get_object(Bucket=bucket, Key=key)
                updated_content_type = verify_response.get('ContentType', 'unknown')
                logger.info(f"Updated manifest content type: {updated_content_type}")
        except Exception as e:
            logger.warning(f"Error updating content type: {str(e)}")
        
        # Track statistics
        line_count = 0
        valid_count = 0
        invalid_lines = []
        
        for i, line in enumerate(content.splitlines()):
            line_count += 1
            line = line.strip()
            if not line:
                continue  # Skip empty lines
                
            try:
                entry = json.loads(line)
                
                # Log sample entry for debugging
                if i == 0:
                    logger.info("Sample manifest entry:")
                    logger.info(json.dumps(entry, indent=2))
                
                # Validate source-ref
                if 'source-ref' not in entry:
                    invalid_lines.append((i+1, "Missing source-ref field"))
                    continue
                    
                if not isinstance(entry.get('source-ref'), str) or not entry['source-ref'].startswith('s3://'):
                    invalid_lines.append((i+1, f"Invalid source-ref format: {entry.get('source-ref')}"))
                    continue
                
                # Verify file exists in S3 at the exact path specified
                source_ref = entry['source-ref']
                file_bucket = source_ref.split('/')[2]
                file_key = '/'.join(source_ref.split('/')[3:])
                
                try:
                    s3_client.head_object(Bucket=file_bucket, Key=file_key)
                    logger.info(f"Verified file exists: {file_key}")
                except Exception as e:
                    invalid_lines.append((i+1, f"File not found in S3: {source_ref}"))
                    logger.warning(f"File not found: {source_ref}")
                    continue
                    
                if 'bounding-box' not in entry or 'bounding-box-metadata' not in entry:
                    invalid_lines.append((i+1, "Missing bounding-box or bounding-box-metadata"))
                    continue
                    
                # Validate bounding-box structure
                bb = entry['bounding-box']
                if not isinstance(bb, dict):
                    invalid_lines.append((i+1, "bounding-box must be a dictionary"))
                    continue
                    
                if 'annotations' not in bb or not isinstance(bb['annotations'], list):
                    invalid_lines.append((i+1, "Missing or invalid annotations array"))
                    continue
                    
                if 'image_size' not in bb or not isinstance(bb['image_size'], list):
                    invalid_lines.append((i+1, "Missing or invalid image_size array"))
                    continue
                    
                # Validate annotations
                for j, annotation in enumerate(bb['annotations']):
                    if not isinstance(annotation, dict):
                        invalid_lines.append((i+1, f"Invalid annotation format at index {j}"))
                        continue
                        
                    required_fields = ['class_id', 'left', 'top', 'width', 'height']
                    missing_fields = [f for f in required_fields if f not in annotation]
                    if missing_fields:
                        invalid_lines.append((i+1, f"Missing required annotation fields: {', '.join(missing_fields)}"))
                        continue
                        
                    # Get image dimensions from image_size
                    try:
                        img_width = bb['image_size'][0]['width']
                        img_height = bb['image_size'][0]['height']
                    except (KeyError, IndexError):
                        invalid_lines.append((i+1, "Invalid image_size structure"))
                        continue
                        
                    # Validate coordinate ranges - now checking pixel values against image dimensions
                    try:
                        # For left and top, must be within image dimensions
                        if not 0 <= annotation['left'] < img_width:
                            invalid_lines.append((i+1, f"Invalid left value: {annotation['left']} (must be between 0 and {img_width-1})"))
                            continue
                            
                        if not 0 <= annotation['top'] < img_height:
                            invalid_lines.append((i+1, f"Invalid top value: {annotation['top']} (must be between 0 and {img_height-1})"))
                            continue
                            
                        # For width and height, must be positive and not extend beyond image
                        if annotation['width'] <= 0:
                            invalid_lines.append((i+1, f"Invalid width value: {annotation['width']} (must be positive)"))
                            continue
                            
                        if annotation['height'] <= 0:
                            invalid_lines.append((i+1, f"Invalid height value: {annotation['height']} (must be positive)"))
                            continue
                            
                        # Check if bounding box extends beyond image boundaries
                        if annotation['left'] + annotation['width'] > img_width:
                            invalid_lines.append((i+1, f"Bounding box extends beyond right edge: left({annotation['left']}) + width({annotation['width']}) > image_width({img_width})"))
                            continue
                            
                        if annotation['top'] + annotation['height'] > img_height:
                            invalid_lines.append((i+1, f"Bounding box extends beyond bottom edge: top({annotation['top']}) + height({annotation['height']}) > image_height({img_height})"))
                            continue
                    except (TypeError, ValueError):
                        invalid_lines.append((i+1, f"Coordinate values must be integers or numbers"))
                        continue
                
                # Validate bounding-box-metadata structure
                bbm = entry['bounding-box-metadata']
                required_meta = ['objects', 'class-map', 'type', 'human-annotated', 'creation-date', 'job-name']
                missing_fields = [f for f in required_meta if f not in bbm]
                if missing_fields:
                    invalid_lines.append((i+1, f"Missing required metadata fields: {', '.join(missing_fields)}"))
                    continue
                    
                if bbm['type'] != 'groundtruth/object-detection':
                    invalid_lines.append((i+1, f"Invalid type in metadata: {bbm['type']}"))
                    continue
                
                # If we get here, the entry is valid
                valid_count += 1
                
            except json.JSONDecodeError as e:
                invalid_lines.append((i+1, f"Invalid JSON: {str(e)}"))
                continue
            except Exception as e:
                invalid_lines.append((i+1, f"Validation error: {str(e)}"))
                continue
        
        # Log validation results
        logger.info(f"Manifest validation results:")
        logger.info(f"Total lines: {line_count}")
        logger.info(f"Valid entries: {valid_count}")
        logger.info(f"Invalid entries: {len(invalid_lines)}")
        
        if invalid_lines:
            logger.warning("Invalid entries found:")
            for line_num, error in invalid_lines:
                logger.warning(f"Line {line_num}: {error}")
        
        # If we have any valid entries, rebuild the manifest with only valid entries
        if valid_count > 0:
            # Read the file again to get valid entries
            valid_lines = []
            content = s3_client.get_object(Bucket=bucket, Key=key)['Body'].read().decode('utf-8')
            
            # Log manifest statistics before rebuilding
            logger.info(f"Rebuilding manifest with {valid_count} entries")
            
            for line in content.splitlines():
                if not line.strip():
                    continue
                try:
                    entry = json.loads(line)
                    # Basic validation
                    if all(k in entry for k in ['source-ref', 'bounding-box', 'bounding-box-metadata']):
                        # Ensure each line ends with exactly one newline
                        valid_lines.append(json.dumps(entry, separators=(',', ':')))
                except Exception as e:
                    logger.warning(f"Error parsing line during rebuild: {str(e)}")
                    continue
            
            # Write back only valid entries
            if valid_lines:
                # Ensure proper JSONL format with exactly one newline between entries and at end
                manifest_content = '\n'.join(valid_lines) + '\n'
                
                # Log a sample of the final content
                logger.info(f"Final manifest sample (first 500 chars):\n{manifest_content[:500]}")
                
                s3_client.put_object(
                    Bucket=bucket,
                    Key=key,
                    Body=manifest_content,
                    ContentType='application/x-amazon-s3-object-manifest-jsonl'
                )
                
                # Verify the upload
                verify_response = s3_client.get_object(Bucket=bucket, Key=key)
                logger.info(f"Final manifest content type: {verify_response.get('ContentType')}")
                logger.info(f"Final manifest size: {len(manifest_content)} bytes")
                return True
        
        return valid_count > 0
        
    except Exception as e:
        logger.error(f"Error validating manifest file: {str(e)}")
        return False

def start_rekognition_training(
    project_arn: str,
    training_bucket: str,
    manifest_key: str = "manifest.jsonl",
    test_manifest_key: str = None,
    version_name: str = None,
    output_bucket: str = None,
    output_folder: str = None
) -> Dict[str, Any]:
    """
    Start training a Rekognition Custom Labels model using external manifest files.
    """
    try:
        rekognition_client = boto3.client('rekognition')
        
        # Generate version name if not provided
        if not version_name:
            version_name = f"version-{datetime.utcnow().strftime('%Y%m%d-%H%M%S')}"
        logger.info(f"Starting training job: {version_name}")

        # Use training bucket as output bucket if not specified
        if not output_bucket:
            output_bucket = training_bucket
            
        # Use version name as output folder if not specified
        if not output_folder:
            output_folder = f"training_output/{version_name}"
            
        # Validate the manifest file
        logger.info(f"Validating manifest file: s3://{training_bucket}/{manifest_key}")
        if not validate_manifest_file(training_bucket, manifest_key):
            raise ValueError(f"Manifest file is not valid: s3://{training_bucket}/{manifest_key}")
            
        # Validate test manifest if provided
        if test_manifest_key:
            logger.info(f"Validating test manifest file: s3://{training_bucket}/{test_manifest_key}")
            if not validate_manifest_file(training_bucket, test_manifest_key):
                raise ValueError(f"Test manifest file is not valid: s3://{training_bucket}/{test_manifest_key}")

        # Prepare training parameters
        training_params = {
            'ProjectArn': project_arn,
            'VersionName': version_name,
            'OutputConfig': {
                'S3Bucket': output_bucket,
                'S3KeyPrefix': output_folder
            },
            'TrainingData': {
                'Assets': [
                    {
                        'GroundTruthManifest': {
                            'S3Object': {
                                'Bucket': training_bucket,
                                'Name': manifest_key
                            }
                        }
                    }
                ]
            },
            'TestingData': {
                'AutoCreate': True  # Will split training data if no test data provided
            }
        }

        # Add test data if provided
        if test_manifest_key:
            training_params['TestingData'] = {
                'Assets': [
                    {
                        'GroundTruthManifest': {
                            'S3Object': {
                                'Bucket': training_bucket,
                                'Name': test_manifest_key
                            }
                        }
                    }
                ]
            }
        
        # Start training
        logger.info("Starting model training with parameters:")
        logger.info(json.dumps(training_params, indent=2))
        response = rekognition_client.create_project_version(**training_params)
        
        logger.info(f"Training started successfully: {response['ProjectVersionArn']}")
        
        # Wait for training to complete
        logger.info("Waiting for training to complete...")
        waiter = rekognition_client.get_waiter('project_version_training_completed')
        waiter.wait(
            ProjectArn=project_arn,
            VersionNames=[version_name],
            WaiterConfig={
                'Delay': 30,
                'MaxAttempts': 48  # Up to 24 hours (30 sec * 48 = 24 min, will check status every 24 min)
            }
        )
        
        # Get final training status
        status_info = check_training_job_status(project_arn, version_name)
        logger.info(f"Training completed with status: {status_info['status']}")
        if status_info['status_message']:
            logger.info(f"Status message: {status_info['status_message']}")
            
        return response
        
    except Exception as e:
        logger.error(f"Error in training workflow: {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 Rekognition client
        rekognition_client = get_rekognition_client()
        
        # Get project versions
        response = rekognition_client.describe_project_versions(
            ProjectArn=project_arn,
            VersionNames=[version_name]
        )
        
        if not response['ProjectVersionDescriptions']:
            raise ValueError(f"Version not found: {version_name}")
            
        version = response['ProjectVersionDescriptions'][0]
        status = version['Status']
        
        # Log status and any status message
        logger.info(f"Training status for {version_name}: {status}")
        if 'StatusMessage' in version:
            logger.info(f"Status message: {version['StatusMessage']}")
            
        return {
            'status': status,
            'version_name': version_name,
            'project_version_arn': version.get('ProjectVersionArn'),
            'status_message': version.get('StatusMessage', ''),
            'evaluation': version.get('EvaluationResults', {})
        }
        
    except Exception as e:
        logger.error(f"Error checking training status: {str(e)}")
        raise

def list_model_versions(project_arn: str) -> List[Dict[str, Any]]:
    """
    List all model versions for a project.
    
    Args:
        project_arn: The project ARN
        
    Returns:
        List of dictionaries containing version information
    """
    try:
        # Get Rekognition client
        rekognition_client = get_rekognition_client()
        
        # Get all versions
        response = rekognition_client.describe_project_versions(
            ProjectArn=project_arn
        )
        
        versions = []
        for version in response.get('ProjectVersionDescriptions', []):
            versions.append({
                'version_name': version.get('VersionName'),
                'arn': version.get('ProjectVersionArn'),
                'status': version.get('Status'),
                'created': version.get('CreationTimestamp'),
                'evaluation': version.get('EvaluationResults', {})
            })
            
        logger.info(f"Found {len(versions)} model versions")
        return versions
        
    except Exception as e:
        logger.error(f"Error listing model versions: {str(e)}")
        raise

def delete_model_version(version_arn: str) -> None:
    """
    Delete a model version.
    
    Args:
        version_arn: The version ARN to delete
    """
    try:
        # Get Rekognition client
        rekognition_client = get_rekognition_client()
        
        # Delete the version
        rekognition_client.delete_project_version(
            ProjectVersionArn=version_arn
        )
        
        logger.info(f"Deleted model version: {version_arn}")
        
    except Exception as e:
        logger.error(f"Error deleting model version: {str(e)}")
        raise

def create_dataset_from_manifest(
    project_arn: str,
    bucket: str,
    manifest_key: str = "manifest.jsonl",
    dataset_type: str = "TRAIN"
) -> Dict[str, Any]:
    """
    Creates a new dataset from a manifest file.
    
    Args:
        project_arn: The ARN of the project
        bucket: The S3 bucket containing the manifest file
        manifest_key: The key of the manifest file
        dataset_type: The type of dataset to create ('TRAIN' or 'TEST')
        
    Returns:
        Dictionary containing information about the created dataset
    """
    try:
        # Initialize rekognition client
        rekognition_client = boto3.client('rekognition')
        
        # Validate dataset_type
        if dataset_type not in ['TRAIN', 'TEST']:
            raise ValueError(f"Invalid dataset_type: {dataset_type}. Must be 'TRAIN' or 'TEST'")
            
        # Verify manifest file exists
        s3_client = boto3.client('s3')
        try:
            s3_client.head_object(Bucket=bucket, Key=manifest_key)
        except Exception as e:
            logger.error(f"Manifest file not found: {str(e)}")
            raise ValueError(f"Manifest file s3://{bucket}/{manifest_key} does not exist")
            
        # Check content type
        response = s3_client.get_object(Bucket=bucket, Key=manifest_key)
        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': manifest_key},
                Bucket=bucket,
                Key=manifest_key,
                ContentType=expected_type,
                MetadataDirective='REPLACE'
            )
            logger.info(f"Updated content type to {expected_type}")
            
        # Create dataset
        logger.info(f"Creating {dataset_type} dataset from manifest s3://{bucket}/{manifest_key}")
        
        dataset_source = {
            "GroundTruthManifest": {
                "S3Object": {
                    "Bucket": bucket,
                    "Name": manifest_key
                }
            }
        }
        
        response = rekognition_client.create_dataset(
            ProjectArn=project_arn,
            DatasetType=dataset_type,
            DatasetSource=dataset_source
        )
        
        logger.info(f"Successfully created {dataset_type} dataset: {response['DatasetArn']}")
        return response
        
    except Exception as e:
        logger.error(f"Error creating dataset from manifest: {str(e)}")
        raise

def delete_dataset(project_arn: str, dataset_type: str = "TRAIN") -> bool:
    """
    Deletes a dataset from a project.
    
    Args:
        project_arn: The ARN of the project
        dataset_type: The type of dataset to delete ('TRAIN' or 'TEST')
        
    Returns:
        True if successful, False otherwise
    """
    try:
        # Initialize rekognition client
        rekognition_client = boto3.client('rekognition')
        
        # Validate dataset_type
        if dataset_type not in ['TRAIN', 'TEST']:
            raise ValueError(f"Invalid dataset_type: {dataset_type}. Must be 'TRAIN' or 'TEST'")
            
        # Find the dataset
        response = rekognition_client.describe_projects()
        
        dataset_arn = None
        for project in response.get('ProjectDescriptions', []):
            if project.get('ProjectArn') == project_arn:
                for dataset in project.get('Datasets', []):
                    if dataset.get('DatasetType') == dataset_type:
                        dataset_arn = dataset.get('DatasetArn')
                        break
                break
                
        if not dataset_arn:
            logger.warning(f"No {dataset_type} dataset found for project: {project_arn}")
            return False
            
        # Delete the dataset
        logger.info(f"Deleting {dataset_type} dataset: {dataset_arn}")
        rekognition_client.delete_dataset(
            DatasetArn=dataset_arn
        )
        
        logger.info(f"Successfully deleted {dataset_type} dataset")
        return True
        
    except Exception as e:
        logger.error(f"Error deleting dataset: {str(e)}")
        return False