from jnpr.junos.exception import RpcError


def provides_facts():
    """
    Returns a dictionary keyed on the facts provided by this module. The value
    of each key is the doc string describing the fact.
    """
    return {'srx_cluster': "A boolean indicating if the device is part of an "
                           "SRX cluster.",
            'srx_cluster_id': "A string containing the configured cluster id",
            'srx_cluster_redundancy_group': "A multi-level dictionary of "
                                            "information about the SRX "
                                            "cluster redundancy groups "
                                            "on the device. The first-level "
                                            "key is the redundancy group id. "
                                            "The second-level keys are: "
                                            "cluster_id, failover_count, "
                                            "node0, and node1. The node0 and "
                                            "node1 keys have third-level keys "
                                            "of priority, preempt, status, "
                                            "and failover_mode. The values "
                                            "for this fact correspond to the "
                                            "values of the 'show chassis "
                                            "cluster status' CLI command."}


def get_facts(device):
    """
    Gathers facts from the <get-chassis-cluster-status/> RPC.
    """
    srx_cluster = None
    srx_cluster_id = None
    redundancy_group = None

    try:
        rsp = device.rpc.get_chassis_cluster_status(normalize=True)
        if rsp is not None:
            if rsp.tag == 'error':
                srx_cluster = False
            else:
                srx_cluster = True
                srx_cluster_id = rsp.findtext('cluster-id')
                groups = rsp.findall('redundancy-group')
                if groups is not None:
                    redundancy_group = {}
                    for group in groups:
                        group_id = group.findtext('redundancy-group-id')
                        redundancy_group[group_id] = {
                            'cluster_id': group.findtext('cluster-id'),
                            'failover_count': group.findtext(
                                'redundancy-group-failover-count'),
                        }
                        # Iterate over the broken XML in <device-stats>
                        for stats in group.findall('./device-stats'):
                            for node in zip(stats.findall('device-name'),
                                            stats.findall('device-priority'),
                                            stats.findall(
                                                'redundancy-group-status'),
                                            stats.findall('preempt'),
                                            stats.findall('failover-mode')):
                                redundancy_group[group_id][node[0].text] = {
                                    'priority': node[1].text,
                                    'status': node[2].text,
                                    'preempt': node[3].text,
                                    'failover_mode': node[4].text,
                                }
    except RpcError:
        # Likely a device that doesn't implement the
        # <get-chassis-cluster-status/> RPC.
        # That's OK. Just ignore it and leave srx_cluster = None.
        pass
    return {'srx_cluster': srx_cluster,
            'srx_cluster_id': srx_cluster_id,
            'srx_cluster_redundancy_group': redundancy_group}
