import json
import boto3
from botocore.exceptions import ClientError
import smtplib
from PIL import Image
from datetime import datetime, timedelta
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.image import MIMEImage
import os

# boto3 clients
cloudwatch_client = boto3.client('cloudwatch')
ssm_client = boto3.client('ssm')

# Parameter Store path prefix - can be overridden via environment variable
PARAMETER_PREFIX = os.environ.get('PARAMETER_PREFIX', '/cloudwatch-dashboard-email')

# SMTP port (always uses port 465 with SSL)
smtp_port = 465

def get_parameter(parameter_name, decrypt=False):
    """
    Retrieve a parameter from AWS Systems Manager Parameter Store.
    
    Args:
        parameter_name: Name of the parameter (without prefix)
        decrypt: Whether to decrypt SecureString parameters (default: False)
    
    Returns:
        Parameter value as string
    """
    try:
        full_path = f"{PARAMETER_PREFIX}/{parameter_name}"
        response = ssm_client.get_parameter(Name=full_path, WithDecryption=decrypt)
        param_type = response['Parameter']['Type']
        param_value = response['Parameter']['Value']
        
        # Log parameter type for debugging
        print(f"Retrieved parameter {full_path} (Type: {param_type})")
        
        return param_value
    except ClientError as e:
        error_code = e.response['Error']['Code']
        if error_code == 'ParameterNotFound':
            raise ValueError(f"Parameter not found: {full_path}")
        elif error_code == 'ValidationException' and decrypt:
            # If trying to decrypt a non-SecureString parameter, try without decryption
            print(f"Warning: Parameter {full_path} is not a SecureString, retrieving without decryption")
            try:
                response = ssm_client.get_parameter(Name=full_path, WithDecryption=False)
                return response['Parameter']['Value']
            except Exception as retry_e:
                raise ValueError(f"Failed to retrieve parameter {full_path}: {str(retry_e)}")
        else:
            raise ValueError(f"Failed to retrieve parameter {full_path}: {error_code} - {str(e)}")
    except Exception as e:
        raise ValueError(f"Failed to retrieve parameter {full_path}: {str(e)}")

def get_configuration():
    """
    Retrieve all configuration from Parameter Store.
    Returns a dictionary with all configuration values.
    """
    return {
        'email_from': get_parameter('email_from'),
        'email_to': get_parameter('email_to'),
        'email_subject': get_parameter('email_subject'),
        'dashboard_name': get_parameter('dashboard_name'),
        'smtp_host': get_parameter('smtp_host'),
        'smtp_username': get_parameter('smtp_username'),
        'smtp_password': get_parameter('smtp_password', decrypt=True)  # Decrypt SecureString
    }

def lambda_handler(event, context):
    # Retrieve configuration from Parameter Store
    try:
        config = get_configuration()
    except ValueError as e:
        error_message = f"Configuration error: {str(e)}"
        print(error_message)
        return {
            'statusCode': 500,
            'body': json.dumps({
                'error': error_message
            })
        }
    
    # Create a multipart/mixed parent container
    msg = MIMEMultipart('related')
    msg['Subject'] = config['email_subject']
    msg['From'] = config['email_from']
    msg['To'] = config['email_to']
    
    msg_alternative = MIMEMultipart('alternative')
    msg.attach(msg_alternative)
    
    # Fetch widget images from CloudWatch
    response = cloudwatch_client.get_dashboard(DashboardName=config['dashboard_name'])
    dashboard_body = response['DashboardBody']
    widgets = json.loads(dashboard_body)
    widgets = widgets['widgets']

    # Calculate time range for the past 3 days
    end_time = datetime.utcnow()
    start_time = end_time - timedelta(days=3)
    time_range = {
        "start": start_time.isoformat() + 'Z',
        "end": end_time.isoformat() + 'Z'
    }

    # Create email content by looping through the widgets
    images = f"<h1>CloudWatch Dashboard Snapshot for {datetime.today().strftime('%Y-%m-%d')}</h1>"
    count = 0
    for widget in widgets:
        count += 1
        # Update widget properties with 3-day time range
        widget_properties = widget['properties']
        widget_properties['start'] = time_range['start']
        widget_properties['end'] = time_range['end']
        widget_properties['period'] = 3600  # 1-hour period for aggregation
        metric_widget = json.dumps(widget_properties)
        image_response = cloudwatch_client.get_metric_widget_image(MetricWidget=metric_widget, OutputFormat='png')

        # Get image bytes
        image_bytes = image_response['MetricWidgetImage']
        image_name = f'DashboardImage_{count}_{datetime.today().strftime("%Y-%m-%d")}.png'

        # Create email body (images)
        images += f'<img src="cid:{image_name}" alt="Dashboard Widget Image {count}">'

        # Attach the image to the email
        img = MIMEImage(image_bytes)
        img.add_header('Content-ID', f'<{image_name}>')
        img.add_header('Content-Disposition', 'inline', filename=image_name)
        msg.attach(img)

    msg_alternative.attach(MIMEText(images, 'html'))
    
    # Send the email using SMTP
    try:
        # Validate required SMTP configuration
        if not config['smtp_host']:
            raise ValueError("smtp_host parameter is required")
        if not config['smtp_username'] or not config['smtp_password']:
            raise ValueError("smtp_username and smtp_password parameters are required")
        
        # Create SMTP connection using SSL (port 465)
        # Add timeout to prevent hanging on network issues
        smtp_timeout = 30  # 30 second timeout
        print(f"Attempting to connect to SMTP server {config['smtp_host']}:{smtp_port} (timeout: {smtp_timeout}s)")
        
        try:
            server = smtplib.SMTP_SSL(config['smtp_host'], smtp_port, timeout=smtp_timeout)
            print(f"Successfully connected to SMTP server {config['smtp_host']}:{smtp_port}")
        except OSError as e:
            error_msg = f"Network error connecting to SMTP server {config['smtp_host']}:{smtp_port}: {str(e)}. "
            error_msg += "This may indicate: 1) Lambda is in a VPC without internet access (needs NAT Gateway), "
            error_msg += "2) Security groups blocking outbound traffic, 3) SMTP server is unreachable."
            print(error_msg)
            raise ValueError(error_msg)
        except Exception as e:
            error_msg = f"Failed to connect to SMTP server {config['smtp_host']}:{smtp_port}: {str(e)}"
            print(error_msg)
            raise ValueError(error_msg)
        
        # Authenticate with SMTP server
        try:
            server.login(config['smtp_username'], config['smtp_password'])
            print(f"Successfully authenticated as {config['smtp_username']}")
        except smtplib.SMTPAuthenticationError as e:
            error_msg = f"SMTP authentication failed: {str(e)}"
            print(error_msg)
            server.quit()
            raise ValueError(error_msg)
        
        # Send email and capture server response
        try:
            # sendmail returns a dictionary of failed recipients (empty dict = success)
            failed_recipients = server.sendmail(
                config['email_from'], 
                [config['email_to']], 
                msg.as_string()
            )
            
            if failed_recipients:
                error_msg = f"SMTP server rejected some recipients: {failed_recipients}"
                print(error_msg)
                server.quit()
                raise ValueError(error_msg)
            
            print(f"Email accepted by SMTP server. From: {config['email_from']}, To: {config['email_to']}")
            
        except smtplib.SMTPRecipientsRefused as e:
            error_msg = f"SMTP server refused recipients: {str(e)}"
            print(error_msg)
            server.quit()
            raise ValueError(error_msg)
        except smtplib.SMTPDataError as e:
            error_msg = f"SMTP server rejected email data: {str(e)}"
            print(error_msg)
            server.quit()
            raise ValueError(error_msg)
        
        server.quit()
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Email sent successfully via SMTP (accepted by server)',
                'from': config['email_from'],
                'to': config['email_to'],
                'note': 'Email was accepted by SMTP server. If not received, check: 1) Spam/junk folder, 2) DNS records (SPF/DKIM/DMARC) for overdriveqa.co.za, 3) Recipient server logs'
            })
        }
    except Exception as e:
        # Log error and return failure response
        error_message = f"Failed to send email via SMTP: {str(e)}"
        print(error_message)
        import traceback
        print(traceback.format_exc())  # Print full traceback for debugging
        return {
            'statusCode': 500,
            'body': json.dumps({
                'error': error_message,
                'from': config.get('email_from', 'unknown'),
                'to': config.get('email_to', 'unknown')
            })
        }