from tb_rest_client.rest_client_pe import *
from tb_rest_client.rest import ApiException
import json, xlsxwriter, boto3, os
from datetime import datetime as dt
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email import encoders
from email.mime.base import MIMEBase

def lambda_handler(event, context):
    #Creating Rest Client for ThingsBoard
    with RestClientPE(base_url="https://hp.henrypump.cloud") as rest_client:
        try:
            rest_client.login(username=os.environ["username"], password=os.environ["password"])
            #Loading Config from file
            with open("./config.json") as f:
                config = json.load(f)

            reportData = {}
            reportToList = {}
            #Loop through each item in config, each item represents a report
            for report in config:
                reportToList[report["name"]] = report["emails"]
                #Each customer becomes it's own xlsx file later
                for customer in report["customers"].keys():
                    #Get all the devices for a given customer
                    devices = rest_client.get_customer_devices(customer_id=customer, page=0, page_size=100)
                    #Filter devices to the desired devices
                    if report["filterDevicesIn"]:
                        devices.data = [device for device in devices.data if device.id.id in report["filterDevicesIn"]]
                    if report["filterDevicesOut"]:
                        devices.data = [device for device in devices.data if device.id.id not in report["filterDevicesOut"]]
                    #Create the report in reportData if needed
                    if not reportData.get(report["name"], None):
                        reportData[report["name"]] = {}
                    #Go through all the devices and add them and their desired data to reportData 
                    for device in devices.data:
                        name = device.name
                        keys = rest_client.get_timeseries_keys_v1(device.id)
                        #Filter keys to desired keys
                        for deviceType in report["customers"][customer]["deviceTypes"]:
                            if device.type == deviceType["deviceType"]:
                                keys = list(filter(lambda x: x in deviceType["dataPoints"], keys))
                        #Create customer if needed
                        if not reportData[report["name"]].get(report["customers"][customer]["name"], None):
                            reportData[report["name"]][report["customers"][customer]["name"]] = {}
                        #Check to make sure the deviceType is desired in the report for the given device
                        if device.type in list(map(lambda x: x["deviceType"], report["customers"][customer]["deviceTypes"])):
                            #Create deviceType if needed
                            if not reportData[report["name"]][report["customers"][customer]["name"]].get(device.type, None):
                                reportData[report["name"]][report["customers"][customer]["name"]][device.type] = {}
                            reportData[report["name"]][report["customers"][customer]["name"]][device.type][device.name] = rest_client.get_latest_timeseries(entity_id=device.id , keys=",".join(keys))
        except ApiException as e:
            print(e)


    # Create an AWS SES client
    ses_client = boto3.client('ses', region_name='us-east-1')


    # Create a workbook for each report
    for report_name, report_data in reportData.items():
        #will generate an email lower down
        spreadsheets = []
        # Create a worksheet for each company
        for company_name, company_data in report_data.items():
            workbook = xlsxwriter.Workbook(f"/tmp/{report_name}-{company_name}-{dt.today().strftime('%Y-%m-%d')}.xlsx")
            bold = workbook.add_format({'bold': True})
            # Create a sheet for each device type
            for device_type, device_data in company_data.items():
                worksheet = workbook.add_worksheet(device_type)
                
                # Set the header row with device types
                device_types = list(device_data.keys())
                worksheet.write_column(1, 0, device_types,bold)
                
                # Write the data to the sheet
                for i, (telemetry_name, telemetry_data) in enumerate(device_data.items()):
                    # Set the header row with telemetry names
                    telemetry_names = list(telemetry_data.keys())
                    worksheet.write_row(0, 1, telemetry_names, bold)
                    for j, (data_name, data) in enumerate(telemetry_data.items()):
                        values = [d["value"] for d in data]
                        worksheet.write_row(i + 1, j+ 1, values)
            worksheet.autofit()
            workbook.close()
            spreadsheets.append(workbook)
        # Create an email message
        msg = MIMEMultipart()
        msg['Subject'] = report_name
        msg['From'] = 'alerts@henry-pump.com'
        msg['To'] = ", ".join(reportToList[report_name])

        # Add a text body to the message (optional)
        body_text = 'Please find the attached spreadsheets.'
        msg.attach(MIMEText(body_text, 'plain'))

        # Attach each workbook in the spreadsheets array
        for spreadsheet in spreadsheets:
            # Attach the file to the email message
            attachment = MIMEBase('application', 'octet-stream')
            attachment.set_payload(open(spreadsheet.filename, "rb").read())
            encoders.encode_base64(attachment)
            attachment.add_header('Content-Disposition', 'attachment', filename=spreadsheet.filename[5:])

            msg.attach(attachment)

        # Send the email using AWS SES
        response = ses_client.send_raw_email(
            
            RawMessage={'Data': msg.as_string()}
        )