#!/usr/bin/python
import sys
import os
import boto3
from pprint import pprint
from functools import reduce

# export AWS_ACCESS_KEY_ID=
# export AWS_SECRET_ACCESS_KEY=
# export AWS_DEFAULT_REGION=

try:
    AWS_ACCESS_KEY_ID = os.environ['AWS_ACCESS_KEY_ID']
    AWS_SECRET_ACCESS_KEY = os.environ['AWS_SECRET_ACCESS_KEY']
except KeyError:
    print("Please set env variable")
    sys.exit(1)


ec2 = boto3.client('ec2')
result = ec2.describe_instances()
reservations = result["Reservations"]

running_instances = {}
spot_instances = {}
state_name = []
for reservation in reservations:
    for instance in reservation["Instances"]:
        if instance["State"]["Name"] not in state_name:
            state_name.append(instance["State"]["Name"])
        if instance["State"]["Name"] != "running":
            sys.stderr.write(
                "Disqualifying instance %s: not running\n" % (instance["InstanceId"]))
        elif "SpotInstanceRequestId" in instance:
            instance_type = instance["InstanceType"]
            spot_instances[(instance_type)] = spot_instances.get(
                (instance_type), 0) + 1
            # sys.stderr.write(
            #     "Disqualifying instance %s: spot\n" % (instance["InstanceId"]))
        else:
            instance_type = instance["InstanceType"]
            running_instances[(instance_type)] = running_instances.get(
                (instance_type), 0) + 1


# pprint(state_name)
# pprint(running_instances)


reserved_instances = {}
result = ec2.describe_reserved_instances()
for reserved_instance in result["ReservedInstances"]:
    if reserved_instance["State"] != "active":
        sys.stderr.write(
            "Excluding reserved instances %s: no longer active\n" % (reserved_instance["ReservedInstancesId"]))
    else:
        instance_type = reserved_instance["InstanceType"]
        reserved_instances[(instance_type)] = reserved_instances.get(
            (instance_type), 0) + reserved_instance["InstanceCount"]

# pprint(reserved_instances)


# this dict will have a positive number if there are unused reservations
# and negative number if an instance is on demand
instance_diff = dict([(x, reserved_instances[x] -
                       running_instances.get(x, 0)) for x in reserved_instances])

# instance_diff only has the keys that were present in reserved_instances. There's probably a cooler way to add a filtered dict here
for placement_key in running_instances:
    if not placement_key in reserved_instances:
        instance_diff[placement_key] = -running_instances[placement_key]

# pprint(instance_diff)

print("")

unused_reservations = dict((key, value)
                           for key, value in instance_diff.items() if value > 0)
if unused_reservations == {}:
    print("Congratulations, you have no unused reservations")
else:
    for unused_reservation in unused_reservations:
        print("UNUSED RESERVATION!\t(%s)\t%s" % (
            unused_reservations[unused_reservation], unused_reservation))

print("")

unreserved_instances = dict((key, -value)
                            for key, value in instance_diff.items() if value < 0)
if unreserved_instances == {}:
    print("Congratulations, you have no unreserved instances")
else:
    for unreserved_instance in unreserved_instances:
        print("Instance not reserved:\t(%s)\t%s" % (
            unreserved_instances[unreserved_instance], unreserved_instance))

qty_running_instances = reduce(
    lambda x, y: x+y, running_instances.values()) if len(running_instances) > 0 else 0
qty_reserved_instances = reduce(
    lambda x, y: x+y, reserved_instances.values()) if len(reserved_instances) > 0 else 0
qty_spot_instances = reduce(
    lambda x, y: x+y, spot_instances.values()) if len(spot_instances) > 0 else 0

print("")
print("(%s) running non-spot instances" % (qty_running_instances))
print("(%s) reservations" % (qty_reserved_instances))
print("(%s) spots" % (qty_spot_instances))