import yp.common as ypc
from yp.client import YpClient
from load.projects.tank_finder.lib.config import logger, yp_token


YP_HOST = "{}.yp.yandex.net:8443"
MAIN_YP = "xdc"
PC_RS = "replica_set"
MC_RS = "multi_cluster_replica_set"


def get_client(cluster, transport="http"):
    try:
        return YpClient(YP_HOST.format(cluster), transport=transport, config={'token': yp_token})
    except ypc.YpClientError as error:
        logger.error("YP client error: %s", str(error))
        raise YpSearchError


def search_replica_set(deploy_dict):
    try:
        if PC_RS in deploy_dict:
            return {cluster: value['replica_set_id'] for cluster, value in dict(deploy_dict[PC_RS]['cluster_statuses']).items()}
        elif MC_RS in deploy_dict:
            return deploy_dict[MC_RS]['replica_set_id']
        else:
            logger.error("No specified replica_set in deploy dict. %s", str(deploy_dict))
            return None
    except (KeyError, ValueError) as error:
        logger.error("Some key was no found in deploy dict! %s", str(error))
        return None


def search_per_cluster_pods_set_id(replica_status):
    try:
        return replica_status[0]['deploy_status']['pod_set_id']
    except (IndexError, KeyError, ValueError):
        logger.error("Wrong replica status: %s", str(replica_status))
        return None


def search_multi_cluster_pods_set_id(replica_status, dc):
    try:
        if dc and dc in [deploy_status['key'] for deploy_status in replica_status[0]['cluster_deploy_statuses']]:
            return [{deploy_status['key']: deploy_status['value']['pod_set_id']} for deploy_status in replica_status[0]['cluster_deploy_statuses'] if deploy_status['key'] == dc]
        else:
            return [{deploy_status['key']: deploy_status['value']['pod_set_id']} for deploy_status in replica_status[0]['cluster_deploy_statuses']]
    except (KeyError, ValueError):
        logger.error("Wrong replica status: %s", str(replica_status))
        return []


class YpSearchError(Exception):
    pass


class YpSearch(object):
    def __init__(self, stage, unit, dc):
        self.stage_id = stage
        self.deploy_unit = unit
        self.data_center = dc
        self._yp_client = {MAIN_YP: get_client(MAIN_YP)}

    def get_client(self, cluster):
        if cluster in self._yp_client:
            return self._yp_client[cluster]
        else:
            self._yp_client[cluster] = get_client(cluster)
            return self._yp_client[cluster]

    def __get_request(self, cluster, o_type, o_id, selectors):
        client = self.get_client(cluster)
        try:
            return client.get_object(object_type=o_type, object_identity=o_id, selectors=selectors)
        except (ypc.YpIncorrectResponseError, ypc.YpInvalidObjectIdError, ypc.YpNoSuchObjectError) as error:
            logger.error("YP get object error! %s", str(error))
            raise YpSearchError

    def __select_request(self, cluster, o_type, selectors, pod_filter):
        client = self.get_client(cluster)
        try:
            return client.select_objects(object_type=o_type, selectors=selectors, filter=pod_filter)
        except (ypc.YpIncorrectResponseError, ypc.YpNoSuchObjectError) as error:
            logger.error("YP select object error! %s", str(error))
            raise YpSearchError

    def get_deploy_units(self):
        return self.__get_request(MAIN_YP, "stage", self.stage_id, ['/annotations/notifications_last_state/status/deploy_units'])

    def get_replica_status(self, cluster, replica_type, replica_set_id):
        return self.__get_request(cluster, replica_type, replica_set_id, ['/status'])

    def select_pods(self, cluster, pod_set_id):
        return self.__select_request(cluster, "pod", ['/status'], "[/meta/pod_set_id]=\"{}\"".format(pod_set_id))

    def get_replica_sets(self):
        deploy_units = self.get_deploy_units()
        if isinstance(deploy_units, list) and isinstance(deploy_units[0], dict):
            if self.deploy_unit and self.deploy_unit in deploy_units[0]:
                return [search_replica_set(deploy_units[0][self.deploy_unit])]
            else:
                return [search_replica_set(value) for value in deploy_units[0].values()]
        else:
            logger.error("A list was expected, but a {} was received".format(type(deploy_units)))
            raise YpSearchError

    def get_pods_set_id(self, replica_list):
        pods_list = []
        for replica in replica_list:
            if isinstance(replica, dict):
                if self.data_center and self.data_center in replica:
                    pods_list.extend([{self.data_center: search_per_cluster_pods_set_id(self.get_replica_status(self.data_center, PC_RS, replica[self.data_center]))}])
                else:
                    pods_list.extend([{cluster: search_per_cluster_pods_set_id(self.get_replica_status(cluster, PC_RS, replica_set_id))} for cluster, replica_set_id in replica.items()])
            elif replica:
                pods_list.extend(search_multi_cluster_pods_set_id(self.get_replica_status(MAIN_YP, MC_RS, replica), self.data_center))
            else:
                logger.error("Unknown replica %s", str(replica))
                raise YpSearchError
        return pods_list

    def get_fqdn(self):
        fqdn = []
        for pods_set in self.get_pods_set_id(self.get_replica_sets()):
            for cluster, pods_set_id in pods_set.items():
                try:
                    fqdn.extend([pod[0]['dns']['transient_fqdn'] for pod in self.select_pods(cluster, pods_set_id)])
                except (KeyError, IndexError, ValueError):
                    logger.info("Wrong pod was returned for cluster %s and pods_set_id %s", str(cluster), str(pods_set_id))
        return fqdn

    def get_subnet(self):
        subnet = []
        for pods_set in self.get_pods_set_id(self.get_replica_sets()):
            for cluster, pods_set_id in pods_set.items():
                try:
                    subnet.extend([pod[0]['ip6_subnet_allocations'][0]['subnet'] for pod in self.select_pods(cluster, pods_set_id)])
                except (KeyError, IndexError, ValueError):
                    logger.info("Wrong pod was returned for cluster %s and pods_set_id %s", str(cluster), str(pods_set_id))
        return subnet
