# 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 pydantic import BaseModel, computed_field

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

if not os.environ.get('AWS_DEFAULT_REGION'):
    os.environ['AWS_DEFAULT_REGION'] = 'us-east-1'


class Config:
    arbitrary_types_allowed = True


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

    @classmethod
    def from_json(cls, json_file):
        """
        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 = {}
                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 cls(
            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"]


class InstanceTypesData(BaseModel):
    data: Dict

    @computed_field
    @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)


class PClusterInstanceTypes(BaseModel):
    instance_type_data: Dict[str, InstanceTypesData]

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

        return cls(instance_type_data=instance_type_defs)
