import os
import sys
from oslo_config import cfg
from swift.common.ring import Ring as swiftRing
from random import randint
import json
from socket import gethostbyaddr, herror
import re

# Initialize Arguments
CLI_OPTS = [
    cfg.IntOpt('nb_tests', short='n', default=500,
               help='Number of partition to test'),
    cfg.BoolOpt('json', short='j',
                help='Show full differences in json formatted string'),
    cfg.StrOpt('pattern', short='p', default=r".*-(?P<idx>\w+)$",
               help='Pattern to extract idx of swiftstore'),
    cfg.StrOpt('Ring1', positional=True, required=True,
               help='New Ring File to compare'),
    cfg.StrOpt('Ring2', positional=True, required=True,
               help='Old Ring File to compare'),
]


def _prepare_config():
    """
    Prepare the oslo_config of scripts by analyse arguments
    return: the oslo_config object
    """
    CONF = cfg.ConfigOpts()
    CONF.register_cli_opts(CLI_OPTS)
    return CONF


def get_short_hostname(ip_address):
    """
    Return short hostname for an IP address
    param ip_address: IP address to resolve
    return: short hostname
    """
    try:
        hostname = gethostbyaddr(ip_address)[0].split('.')[0]
        return hostname
    except herror:
        return "Unknown host"


def extract_host_idx(hostname, pattern):
    """
    Extract an index from a swiftstore name
    param hostname: hostname for exctract index
    param pattern: a regexp pattern to extract index
    return: index
    """
    m = re.match(pattern, hostname)
    try:
        return int(m.group('idx'))
    except Exception:
        return -1


class Ring():
    """Class representing an Ring"""

    def __init__(self, file):
        """
        Instance method to initialize a Ring
        param file: file for the ring (with ring.gz extension)
        """
        self.ring_name = os.path.basename(file)[:-len('.ring.gz')]
        self.ring = swiftRing(file)
        if '-' in self.ring_name and self.ring_name.startswith('object'):
            self.loc = 'objects-' + self.ring_name.split('-', 1)[1]
        else:
            self.loc = self.ring_name + 's'

    def get_nodes(self, part):
        """
        Instance method that return list of node for a partition
        param part: partition to test
        return: list of node and device that have the partition
        """
        part = int(part)
        primary_nodes = self.ring.get_part_nodes(part)
        return primary_nodes

    def get_md5(self):
        """
        Instance method that return de md5 for the ring
        return: md5
        """
        return self.ring.md5

    def get_devs(self):
        """
        Get device in ring
        return: dict with device id and (hostname, device_name)
        """
        devs = [d for d in self.ring.devs if d]
        devs_ip = {dev['id']: (get_short_hostname(dev['ip']), dev['device'])
                   for dev in devs}
        return devs_ip

    def get_parts(self):
        """
        Get parts in ring
        return: dict of list of partitions by device id
        """
        partitions = {}
        for replica in self.ring._replica2part2dev_id:
            for partition, device in enumerate(replica):
                partitions.setdefault(device, set())
                partitions[device].add(partition)
        return partitions

    def partitions_count(self):
        """
        Instance method that return number of partition for the ring
        return: number of partitions
        """
        return len(self.ring._replica2part2dev_id[0])

    def compare_with(self, other_r, part):
        """
        Instance method that compare this ring with another
        If more than 1 difference, display an Error Message and exit
        param other_r: the other ring to compare
        param part: the partition to test
        return: the count of difference and nodes/device affected
        """
        my_nodes = self.get_nodes(part)
        other_n = other_r.get_nodes(part)
        diff_count = 0
        diff = {}
        for i in range(3):  # 3 because replica 3 in both rings
            diff_items = {k: my_nodes[i][k] for k in my_nodes[i]
                          if k in other_n[i] and
                          my_nodes[i][k] != other_n[i][k]}
            if diff_items and not (len(diff_items) == 1
               and "weight" in diff_items):
                diff[self.ring_name] = {'ip': my_nodes[i]['ip'],
                                        'port': my_nodes[i]['port'],
                                        'device': my_nodes[i]['device']}
                diff[other_r.ring_name] = {'ip': other_n[i]['ip'],
                                           'port': other_n[i]['port'],
                                           'device': other_n[i]['device']}
                diff_count += 1
                if diff_count > 1:
                    print("[ERROR] more than 1 differences for one partition, \
                           new ring is not valid.", file=sys.stderr)
                    sys.exit(1)
        return diff_count, diff

    def get_in_out(self, other_r):
        """
        Instance method to extract partitions in/out devices between 2 rings
        param other_r: the other ring to compare
        return: 2 dicts for partitions In and Out of each devices
        """
        parts_in = {}
        parts_out = {}
        self_parts = self.get_parts()
        other_r_parts = other_r.get_parts()
        list_host = list(self_parts.keys()) + list(other_r_parts.keys())
        for k in list_host:
            parts_in.setdefault(k, 0)
            parts_out.setdefault(k, 0)
            if k not in other_r_parts:
                other_r_parts[k] = set()
            if k not in self_parts:
                self_parts[k] = set()
            parts_in[k] += len(self_parts[k] - other_r_parts[k])
            parts_out[k] += len(other_r_parts[k] - self_parts[k])
        return parts_in, parts_out


def main():
    """
    Entry point for the script
    """
    args = _prepare_config()
    try:
        args(sys.argv[1:])
    except cfg.RequiredOptError as E:
        print(E)
    # Get informations for 2 rings
    ring1 = Ring(args.Ring1)
    ring2 = Ring(args.Ring2)
    md5_1, md5_2 = None, None
    try:
        md5_1 = ring1.get_md5()
        md5_2 = ring2.get_md5()
    except Exception:
        pass
    partitions_count = ring1.partitions_count()

    diffs = {}
    nb_diffs = 0
    # Compare the rings for args.nb_test random partitions
    for i in range(args.nb_tests):
        part = randint(1, partitions_count)
        nb_diff, diff = ring1.compare_with(ring2, part)
        nb_diffs += nb_diff
        if diff:
            diffs[f"Partition {part}"] = diff

    # Extract In and Out partitions by devices
    parts_in, parts_out = ring1.get_in_out(ring2)

    # Show result
    print()
    if args.json:
        diffs[f"Ring {ring1.ring_name} md5"] = md5_1
        diffs[f"Ring {ring2.ring_name} md5"] = md5_2
        diff['Partitions IN'] = parts_in
        diff['Partitions OUT'] = parts_out
        print(json.dumps(diffs))
    else:
        if md5_1 == md5_2 and md5_1 is not None and md5_2 is not None:
            print("md5 for two are identical !")
            sys.exit(0)
        else:
            print(f"Ring {ring1.ring_name} md5 : {md5_1}")
            print(f"Ring {ring2.ring_name} md5 : {md5_2}")
        if nb_diffs:
            print(f"Number of differences between {ring1.ring_name} \
                    and {ring2.ring_name}: {nb_diffs}/{args.nb_tests} \
                    partitions tested")
        else:
            print(f"No differences find between rings {ring1.ring_name} \
and {ring2.ring_name} , you can re-run the script for confirmation")
            return 0

    print()
    devs = ring1.get_devs()
    devs.update(ring2.get_devs())
    sum_in = 0
    hosts_move = {}
    # Group partitions IN and OUT by host
    for k, v in parts_in.items():
        idx = extract_host_idx(devs[k][0], args.pattern)
        hosts_move.setdefault(idx, {'name': devs[k][0], 'in': 0, 'out': 0})
        sum_in += v
        hosts_move[idx]['in'] += v
    sum_out = 0
    for k, v in parts_out.items():
        idx = extract_host_idx(devs[k][0], args.pattern)
        hosts_move.setdefault(idx, {'name': devs[k][0], 'in': 0, 'out': 0})
        sum_out += v
        hosts_move[idx]['out'] += v

    print(f"Total partitions moved: {sum_in} In, {sum_out} Out.")
    maxi = 0
    for k, v in hosts_move.items():
        maxi = v['in'] if v['in'] > maxi else maxi
        maxi = v['out'] if v['out'] > maxi else maxi

    div = (maxi / 50)
    len_idx = int(max(list(hosts_move.keys())) / 4) + 1
    format_in = [0] * (50 * len_idx)
    format_out = [0] * (50 * len_idx)
    for k, v in sorted(hosts_move.items()):
        idx = int((k - 1) / 4) if k >= 0 else -1
        reste = 2 ** ((k - 1) % 4)
        nb_in = int(v['in'] / div)
        nb_out = int(v['out'] / div)
        for i in range(0, 50):
            format_in[idx * 50 + i] += (reste if nb_in > i else 0)
            format_out[idx * 50 + i] += (reste if nb_out > i else 0)
    # For grouping graph line by 4 hosts
    char_graph = [' ', '⠉', '⠒', '⠛', '⠤', '⠭', '⠶', '⠿',
                  '⣀', '⣉', '⣒', '⣛', '⣤', '⣭', '⣶', '⣿']
    print()
    pr = f"{'Hosts':28}: \x1b[1;31;40m{'IN max: ' + str(int(div * 50)):<50}"
    pr = f"{pr}\x1b[0;37;40m|\x1b[1;32;40m"
    pr = f"{pr}{'OUT max: ' + str(int(div * 50)):>50}\x1b[0;37;40m"
    print(pr)
    print('-' * 131)
    for idx in range(0, len_idx):
        try:
            hosts = f"{hosts_move[1 + idx * 4]['name']}…{idx * 4 + 4}"
        except Exception:
            if -1 in hosts_move.keys():
                hosts = f"{hosts_move[-1]['name']}"
            else:
                hosts = "Hostname unknown"
        print(f"{hosts:28}: ", end='')
        print('\x1b[1;31;40m', end='')
        for i in range(49, -1, -1):
            try:
                print(char_graph[format_in[idx * 50 + i]], end='')
            except Exception:
                print(' ' * 50, end='')
        print('\x1b[0;37;40m|\x1b[1;32;40m', end='')
        for i in range(0, 50):
            try:
                print(char_graph[format_out[idx * 50 + i]], end='')
            except Exception:
                print(' ' * 50, end='')
        print('\x1b[0;37;40m')
    sys.exit(nb_diffs)


if __name__ == "__main__":
    sys.exit(main())
