# generated by datamodel-codegen:
#   filename:  instance-types-data.json
#   timestamp: 2022-05-24T09:17:47+00:00

from __future__ import annotations
from datasize import DataSize
from copy import deepcopy
from typing import Dict

from mypy_boto3_ec2.type_defs import (InstanceTypeDef)
from pydantic.dataclasses import dataclass

from devtools import PrettyFormat, pprint, pformat, debug
import json
import os
import yaml


class Config:
    arbitrary_types_allowed = True


@dataclass(config=Config)
class InstanceTypesMappings:
    ec2_instance_types: Dict
    sinfo_instance_types: Dict

    @classmethod
    def from_json(cls, json_file) -> InstanceTypesMappings:
        """
        Data looks like this:
        "basic": {
        "t32xlarge": "t3.2xlarge",
        "t3alarge": "t3a.large",
        "t3axlarge": "t3a.xlarge"
        },
        A dictionary of queue name and sinfo instance type mapped back to the ec2 instance name
        """
        json_data = json.load(open(json_file))
        inverse_data = {}
        data = {}
        for queue in json_data.keys():
            instance_types = json_data[queue]
            for sinfo_instance_type in instance_types:
                t_inverse_data = {}
                t_data = {}
                instance_type = instance_types[sinfo_instance_type]
                t_inverse_data['ec2_instance_type'] = instance_type
                t_inverse_data['sinfo_instance_type'] = sinfo_instance_type
                if instance_type not in inverse_data.keys():
                    inverse_data[instance_type] = {}
                if sinfo_instance_type not in data.keys():
                    data[sinfo_instance_type] = {}
                inverse_data[instance_type].update(t_inverse_data)
                data[sinfo_instance_type].update(t_inverse_data)
                if 'queues' not in inverse_data[instance_type].keys():
                    inverse_data[instance_type]['queues'] = [queue]
                    data[sinfo_instance_type]['queues'] = [queue]
                else:
                    inverse_data[instance_type]['queues'].append(queue)
                    data[sinfo_instance_type]['queues'].append(queue)
        return InstanceTypesMappings(ec2_instance_types=inverse_data, sinfo_instance_types=data)


def size_in_gib(mib: int) -> int:
    mib_bytes = DataSize(f'{mib}Mi')
    return mib_bytes / mib_bytes.IEC_prefixes['Gi']


@dataclass(config=Config)
class InstanceTypesData:
    data: Dict

    @property
    def instance_type_data(self) -> InstanceTypeDef:
        if 'Hypervisor' in self.data.keys():
            if 'nitro' in self.data['Hypervisor']:
                self.data['Hypervisor'] = 'xen'
        if 'MemoryInfo' in self.data.keys():
            if 'SizeInMib' in self.data['MemoryInfo'].keys():
                size_in_mib = self.data['MemoryInfo']['SizeInMib']
                bytes = DataSize(f'{size_in_mib}Mi')
                gibs = bytes / bytes.IEC_prefixes['Gi']
                self.data['MemoryInfo']['SizeInGib'] = gibs
        else:
            debug(self.data)
        return deepcopy(self.data)

    @classmethod
    def from_yaml(cls, yaml_file: str) -> InstanceTypesData:
        d = yaml.safe_load(open(yaml_file))
        return InstanceTypesData(data=d)


@dataclass
class PClusterInstanceTypes:
    instance_type_data: Dict[str, InstanceTypesData]

    @classmethod
    def from_json(cls, json_file: str) -> PClusterInstanceTypes:
        data = json.load(open(json_file))
        instance_type_defs = {}
        for k in data.keys():
            instance_type_defs[k] = InstanceTypesData(data={'data': data[k]})

        return PClusterInstanceTypes(instance_type_data=instance_type_defs)
