import boto3
import logging
import os
import json
from datetime import datetime, timedelta
from typing import Dict, Any
from botocore.exceptions import ClientError

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

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

def get_parameter(name: str) -> str:
    """
    Retrieve a parameter from AWS Parameter Store.
    
    Args:
        name: The parameter name to retrieve
        
    Returns:
        The parameter value
        
    Raises:
        ClientError: If parameter retrieval fails
    """
    try:
        response = ssm.get_parameter(Name=name, WithDecryption=True)
        return response['Parameter']['Value']
    except ClientError as e:
        logger.error(f"Failed to get parameter {name}: {str(e)}")
        raise

def check_manifest_freshness(bucket_name: str, manifest_key: str, max_age_hours: int = 168) -> Dict[str, Any]:
    """
    Check if the manifest file in S3 was updated within the specified timeframe.
    
    Args:
        bucket_name: S3 bucket containing the manifest file
        manifest_key: Key of the manifest file in the bucket
        max_age_hours: Maximum age of the manifest file in hours (default: 168, which is 7 days)
        
    Returns:
        Dictionary with freshness information
    """
    try:
        # Get the object metadata
        response = s3_client.head_object(Bucket=bucket_name, Key=manifest_key)
        
        # Get the last modified timestamp
        last_modified = response['LastModified']
        
        # Calculate age of the manifest file
        current_time = datetime.now(last_modified.tzinfo)
        age = current_time - last_modified
        age_hours = age.total_seconds() / 3600
        age_days = age_hours / 24
        
        # Determine if manifest is fresh enough
        is_fresh = age_hours < max_age_hours
        
        logger.info(f"Manifest file last modified: {last_modified}")
        logger.info(f"Manifest age: {age_hours:.2f} hours ({age_days:.2f} days)")
        logger.info(f"Maximum allowed age: {max_age_hours} hours ({max_age_hours/24:.2f} days)")
        logger.info(f"Manifest is {'fresh' if is_fresh else 'stale'}")
        
        return {
            'is_fresh': is_fresh,
            'last_modified': last_modified.isoformat(),
            'age_hours': age_hours,
            'age_days': age_days,
            'max_age_hours': max_age_hours,
            'max_age_days': max_age_hours / 24
        }
    except ClientError as e:
        logger.error(f"Error checking manifest freshness: {str(e)}")
        # If the file doesn't exist, it's definitely not fresh
        return {
            'is_fresh': False,
            'error': str(e),
            'max_age_hours': max_age_hours,
            'max_age_days': max_age_hours / 24
        }

def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """
    Main Lambda handler function.
    
    Args:
        event: Lambda event
        context: Lambda context
        
    Returns:
        Dictionary with manifest freshness information
    """
    try:
        # Get S3 bucket and manifest key from environment or parameters
        training_bucket = os.environ.get('TRAINING_BUCKET')
        manifest_key = os.environ.get('MANIFEST_KEY', 'manifest.jsonl')
        max_age_hours = int(os.environ.get('MAX_AGE_HOURS', '168'))  # Default to 7 days (168 hours)
        
        # Try to get from SSM if not in environment
        if not training_bucket:
            training_bucket = get_parameter('/datafy-rekognition-stack/training-bucket')
            if not training_bucket:
                raise ValueError("Missing required configuration: TRAINING_BUCKET")
        
        # Check manifest freshness
        freshness_info = check_manifest_freshness(training_bucket, manifest_key, max_age_hours)
        
        # Return the result for the state machine
        return {
            'statusCode': 200,
            'manifest_fresh': freshness_info['is_fresh'],
            'freshness_info': freshness_info
        }
    
    except Exception as e:
        error_msg = str(e)
        logger.error(f"Lambda execution failed: {error_msg}")
        return {
            'statusCode': 500,
            'manifest_fresh': False,  # Default to not fresh on error
            'error': error_msg
        }