from pcluster.api.errors import NotFoundException, BadRequestException
from wrapper.ec2_wrapper import EC2Wrapper
import pcluster.lib as pc
import pprint
import yaml

pp = pprint.PrettyPrinter()
ec2 = EC2Wrapper()


class ParallelClusterWrapper:
    # Source: https://docs.aws.amazon.com/parallelcluster/latest/ug/supported-regions-v3.html
    supported_regions_v3 = ['us-east-2', 'us-east-1', 'us-west-1', 'us-west-2', 'af-south-1', 'ap-east-1',
                            'ap-south-1', 'ap-northeast-2', 'ap-southeast-1', 'ap-southeast-2', 'ap-northeast-1', 'ca-central-1', 'cn-north-1',
                            'cn-northwest-1', 'eu-central-1', 'eu-west-1', 'eu-west-2', 'eu-south-1', 'eu-west-3', 'eu-north-1', 'me-south-1', 'sa-east-1']

    def get_supported_regions(self):
        # Convert the lists to sets
        pc_supported_set = set(self.supported_regions_v3)
        user_supported_set = set(ec2.list_all_regions())

        # Find the common regions by taking the intersection of the sets
        common_regions = list(
            pc_supported_set.intersection(user_supported_set))
        return sorted(common_regions)

    def get_cluster_head_node(self, cluster_name):
        try:
            head_node = pc.describe_cluster_instances(
                cluster_name=cluster_name, node_type='HeadNode')
            if not head_node['instances']:
                pp.pprint("No head node available for this cluster!")
            else:
                pp.pprint(head_node)
        except Exception as e:
            print(f"ERROR: An error occurred while creating the cluster: {e}")

    def list_cluster(self):
        try:
            clusters = pc.list_clusters()
            for i, cluster in enumerate(clusters['clusters']):
                print(f'CLUSTER no.{i+1}:')
                pp.pprint(cluster)
                print(f'HEAD NODE:')
                self.get_cluster_head_node(cluster['clusterName'])
                print("\n" + "-" * 50 + "\n")

        except Exception as e:
            print(f"ERROR: An error occurred while creating the cluster: {e}")

    def create_cluster(self, cluster_name, cluster_file):
        try:
            with open(cluster_file, 'r') as f:
                cluster_data = yaml.load(f, Loader=yaml.FullLoader)

            print(f"Creating cluster {cluster_name} with the following data:")
            pp.pprint(
                pc.create_cluster(
                    cluster_name=cluster_name,
                    cluster_configuration=cluster_data,
                    rollback_on_failure=False)
            )
        except FileNotFoundError as e:
            print(f"ERROR: Configuration file has not been found!")
        except BadRequestException as e:
            print(f"ERROR: Invalid cluster action request! {e}")
        except Exception as e:
            print(f"ERROR: An error occurred while creating the cluster: {e}")

    def find_cluster(self, cluster_name):
        try:
            print(f"Fetching information for cluster: {cluster_name}")
            pp.pprint(
                pc.describe_cluster(cluster_name=cluster_name)
            )
        except NotFoundException as e:
            print(f"ERROR: Cluster not found!")
        except Exception as e:
            print(
                f"ERROR: An error occurred while fetching information for the cluster: {e}")

    def delete_cluster(self, cluster_name):
        try:
            print(f"Deleting cluster: {cluster_name}")
            pp.pprint(
                pc.delete_cluster(cluster_name=cluster_name)
            )
        except NotFoundException as e:
            print(f"ERROR: Cluster not found!")
        except Exception as e:
            print(f"ERROR: An error occurred while deleting the cluster: {e}")
