import json
import boto3
import io
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
import matplotlib.pyplot as plt
import pandas as pd
import matplotlib.dates as mdates

# boto3 clients for AWS SES and Lightsail
ses_client = boto3.client('ses')
lightsail_client = boto3.client('lightsail')

# Environment Variables
email_from = os.environ['email_from']
email_to = os.environ['email_to']
email_subject = os.environ['email_subject']
dashboard_name = os.environ['dashboard_name']

def get_lightsail_metrics(instance_name, metric_name, statistics, unit, period=300, start_time=None, end_time=None):
    if start_time is None:
        end_time = datetime.utcnow()
        start_time = end_time - timedelta(hours=24)  # Use last 24 hours for data

    response = lightsail_client.get_instance_metric_data(
        instanceName=instance_name,
        metricName=metric_name,
        period=period,
        startTime=start_time,
        endTime=end_time,
        unit=unit,
        statistics=statistics
    )
    return response['metricData']

def create_metric_plot(metric_data, title, stat_key):
    # Sort data by timestamp
    metric_data = sorted(metric_data, key=lambda x: x['timestamp'])
    # Convert timestamps to datetime objects
    timestamps = [pd.to_datetime(data['timestamp']) for data in metric_data]
    values = [data.get(stat_key, 0) for data in metric_data]
    plt.figure(figsize=(10, 6))
    plt.plot(timestamps, values, linestyle='-')
    plt.title(title)
    plt.xlabel('Time')
    plt.ylabel('Value')
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%m-%d %H:%M'))
    plt.gca().xaxis.set_major_locator(mdates.AutoDateLocator())
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    plt.close()
    return buf.getvalue()

def lambda_handler(event, context):
    # Create a multipart/mixed parent container
    msg = MIMEMultipart('related')
    msg['Subject'] = email_subject
    msg['From'] = email_from
    msg['To'] = email_to
    
    msg_alternative = MIMEMultipart('alternative')
    msg.attach(msg_alternative)
    
    # Get list of Lightsail instances
    instances = lightsail_client.get_instances()['instances']
    
    # Calculate time range for the past 3 days
    end_time = datetime.utcnow()
    start_time = end_time - timedelta(days=3)
    
    # Create email content
    images = f"<h1>Lightsail Metrics Snapshot for {datetime.today().strftime('%Y-%m-%d')}</h1>"
    count = 0
    
    # Metrics to collect for each instance
    metrics = [
        ('CPUUtilization', 'CPU Utilization'),
        ('NetworkIn', 'Network In'),
        ('NetworkOut', 'Network Out'),
        ('StatusCheckFailed', 'Status Check Failed')
    ]
    
    metric_settings = {
        'CPUUtilization': {'statistics': ['Average'], 'stat_key': 'average', 'unit': 'Percent', 'period': 900},
        'NetworkIn': {'statistics': ['Sum'], 'stat_key': 'sum', 'unit': 'Bytes', 'period': 900},
        'NetworkOut': {'statistics': ['Sum'], 'stat_key': 'sum', 'unit': 'Bytes', 'period': 900},
        'StatusCheckFailed': {'statistics': ['Sum'], 'stat_key': 'sum', 'unit': 'Count', 'period': 900}
    }
    
    for instance in instances:
        instance_name = instance['name']
        images += f"<h2>Instance: {instance_name}</h2>"
        for metric_name, metric_title in metrics:
            count += 1
            settings = metric_settings[metric_name]
            metric_data = get_lightsail_metrics(
                instance_name,
                metric_name,
                statistics=settings['statistics'],
                unit=settings['unit'],
                period=settings['period'],
                start_time=start_time,
                end_time=end_time
            )
            image_bytes = create_metric_plot(metric_data, f"{instance_name} - {metric_title}", settings['stat_key'])
            image_name = f'LightsailMetric_{count}_{datetime.today().strftime("%Y-%m-%d")}.png'
            images += f'<img src="cid:{image_name}" alt="{metric_title}"><br/>'
            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 SES
    response = ses_client.send_raw_email(
        Source=email_from,
        Destinations=[email_to],
        RawMessage={
            'Data': msg.as_string().encode('utf-8'),
        }
    )
    return response