#!/usr/bin/python3

"""
ct-create-pinning-osd

Create OSD and pinning in the good CPU
"""

import os
import sys
from subprocess import run, STDOUT
import json
from argparse import ArgumentParser
import datetime
import psutil
import random
import inspect
import socket


default_nvme_to_configure = 10
default_osd_by_nvme = 4


def listing(args):
    """
    Callback function for the list subcommand
    :param args: args of subcommand
    """
    init()
    List = classDict[args.type].List
    json = []
    ret = ""
    for id, ob in sorted(List.items()):
        if args.json:
            json.append(ob.toJson())
        else:
            ret += f"{ob}\n"
    if args.json:
        print(json)
    else:
        print(classDict[args.type].header())
        print(ret, end='')
        print(classDict[args.type].footer())


def checking(args):
    """
    Callback function for the check subcommand
    :param args: args of subcommand
    """
    init()
    for nvme in Nvme.get_list_by_args():
        # Get OSDs in NVME
        for osdid, osd in sorted(
          {osdid: osd for (osdid, osd)
              in Osd.List.items() if osd in nvme.osds}.items()):
            osd.check()


def pinning(args):
    """
    Callback function for the pin subcommand
    :param args: args of subcommand
    """
    init()
    for nvme in Nvme.get_list_by_args():
        # Get OSDs in NVME
        for osdid, osd in sorted(
          {osdid: osd for (osdid, osd)
              in Osd.List.items() if osd in nvme.osds}.items()):
            osd.pinning()


def creating(args):
    """
    Callback function for the create subcommand
    :param args: args of subcommand
    """
    init()
    run(["ceph", "osd", "set", 'noin'])
    if args.devices:
        # Get Nvme in list passed in argument (or all Nvme)
        for nvme in Nvme.get_list_by_args():
            nvme.create_osds()
    else:
        for i in range(args.nvmes):
            NumaNode.create_osd_on_nvme()


def parse_Cpu_Affinity(cpuAffinity):
    """
    Fonction to parse cpuAffinity format (0-15,32-47) to to list of Cpus
    :param cpuAffinity: String cpuAffinity
    :return: list of Cpu Number
    """
    cpuRange = []
    for pranges in cpuAffinity.split(','):
        prange = pranges.split('-')
        if len(prange) == 1:
            prange.append(prange[0])
        for i in range(int(prange[0]), int(prange[1])+1):
            cpuRange.append(i)
    return cpuRange


class NumaNode():
    """Class representing a NumaNode"""

    List = {}

    def create(id):
        """
        Class Method for creating a Numa Node and store it
        in 'List' class attribut

        :param id: the id of Numa Node
        :return: the new Numa Node
        """

        if id not in NumaNode.List:
            NumaNode(id)
        return NumaNode.List[id]

    def create_osd_on_nvme():
        """
        Class Method for creating OSDs in least filled Numa Node
        """

        minNvmeBusy = len(Nvme.List)
        for id, node in NumaNode.List.items():
            nodeNvmeBusy = node.get_nb_nvme_busy()
            if nodeNvmeBusy < minNvmeBusy:
                minNvmeBusy = nodeNvmeBusy
                nodeMinNvmeBusy = node
        nvmes = [nvme for nvme in nodeMinNvmeBusy.nvmes if not nvme.osds]
        nvmes[0].create_osds()

    def get_free_cpus(self):
        """
        Class Method for getting free CPUs

        :return: list of free CPUs
        """

        cpus = [cpu for cpu in self.cpusPrimaries if not cpu.osds]
        if not cpus:
            cpus = [cpu for cpu in self.cpusSecondaries if not cpu.osds]
        return cpus

    def header():
        """
        Class Method for printing header of table
        """
        return ''

    def footer():
        """
        Class Method for printing header of table
        """
        return ''

    def __init__(self, id):
        """
        Instance Method for creating an instance of Numa Node

        :param id: the id of Numa Node
        """

        self.id = id
        self.cpusPrimaries = []
        self.cpusSecondaries = []
        self.nvmes = []
        self.cpuAffinity = ""
        self.cpuRange = []

        NumaNode.List[id] = self

    def __str__(self):
        """
        Instance Method for displaying an instance of Numa Node

        :return: string representing instance
        """

        ret = f"Node {self.id}"
        ret = (f"{ret}\n\tPrimaries Cpus: "
               f"{sorted([int(cpu.id)for cpu in self.cpusPrimaries])}")
        ret = (f"{ret}\n\tSecondaries Cpus: "
               f"{sorted([int(cpu.id) for cpu in self.cpusSecondaries])}")
        ret = f"{ret}\n\tNVME: {sorted([nvme.id for nvme in self.nvmes])}"
        ret = f"{ret}\n\tCpu Affinity: {self.cpuAffinity}"
        ret = f"{ret}\n\tCpu Range: {self.cpuRange}"
        ret = (f"{ret}\n\tNumber of Nvme configured : "
               f"{len({nvme for nvme in self.nvmes if nvme.osds})}")
        return ret

    def toJson(self):
        """
        Instance Method for json output an instance of Numa Node

        :return: string representing instance in json format
        """

        ret = {'Id': self.id}
        ret['Primaries Cpus'] = sorted([cpu.id for cpu in self.cpusPrimaries])
        ret['Secondaries Cpus'] = sorted([cpu.id
                                         for cpu in self.cpusSecondaries])
        ret['NVME'] = sorted([nvme.id for nvme in self.nvmes])
        ret['Cpu Affinity'] = self.cpuAffinity
        ret['Cpu Range'] = self.cpuRange
        ret["Number of Nvme configured"] = len({nvme
                                               for nvme in self.nvmes
                                               if nvme.osds})
        return ret

    def set_cpuAffinity(self, cpuAffinity):
        """
        Instance Method for json setting CpuAffinity

        :param cpuAffinity: The cpuAffinity
        """

        self.cpuAffinity = cpuAffinity
        self.cpuRange = parse_Cpu_Affinity(cpuAffinity)

    def checkAffinity(self, realCpu):
        """
        Instance Method for checking if a Cpu in CpuAffinity

        :param realCpu: Cpu to check
        :return: boolean True if realCpu in cpuAffinity
        """

        return realCpu in self.cpuRange

    def get_nb_nvme_busy(self):
        """
        Instance Method for getting number of nvme busy

        :return: Number of nvme busy for this Numa Node
                 or max nvme of the system if all nvme
                 of this Numa Node are busy
        """

        if len([nvme for nvme in self.nvmes if not nvme.osds]) != 0:
            return len([nvme for nvme in self.nvmes if nvme.osds])
        else:
            return len(Nvme.List)


class Cpu():
    """Class representing a CPU"""

    List = {}

    def create(id, core, node):
        """
        Class Method for creating a Cpu and store it in 'List' class attribut

        :param id: the id of Cpu
        :param core: the core of Cpu
        :param node: the Numa Node of Cpu
        :return: the new Cpu
        """

        if id not in Cpu.List:
            Cpu(id, core, node)
        return Cpu.List[id]

    def header():
        """
        Class Method for printing header of table
        """
        ret = "-" * 80
        ret = (f"{ret}\n| {'Cpu':8} | {'Core':8} | {'NumaNode':8} | "
               f"{'Primary':8} | {'OSDs':32} |\n{ret}")
        return ret

    def footer():
        """
        Class Method for printing header of table
        """
        return "-" * 80

    def __init__(self, id, core, node):
        """
        Instance Method for creating an instance of Cpu

        :param id: the id of Cpu
        :param core: the core of Cpu
        :param node: the Numa Node of Cpu
        """

        self.id = id
        self.numaNode = NumaNode.create(node)
        self.osds = []
        self.core = core
        self.primary = (id == core)
        Cpu.List[id] = self
        # Append if cpusPrimaries if it's a physical Cpu ou cpuSecondaries
        # if it's a HyperThreading
        if id == core:
            self.numaNode.cpusPrimaries.append(self)
        else:
            self.numaNode.cpusSecondaries.append(self)

    def __str__(self):
        """
        Instance Method for displaying an instance of Numa Node

        :return: string representing instance
        """
        return (f"| {self.id:8} | {self.core:8} | {self.numaNode.id:8} | "
                f"{self.primary:<8} | "
                f"{str(sorted([int(osd.id) for osd in self.osds ])):32} |")

    def toJson(self):
        """
        Instance Method for json output an instance of Cpu

        :return: string representing instance in json format
        """

        ret = {'Id': self.id}
        ret['Core'] = self.core
        ret['Numa Node'] = self.numaNode.id
        ret['Primary'] = self.primary
        ret['OSDs'] = sorted([osd.id for osd in self.osds])
        return ret


class Nvme():
    """Class representing a NVME"""
    List = {}

    def create(name, node):
        """
        Class Method for creating a NVME and store it in 'List' class attribut

        :param name: the name of NVME
        :param node: the Numa Node of NVME
        :return: the new NVME
        """

        id = int(name[4:].split('n')[0])
        if id not in Nvme.List:
            Nvme(name, node)
        return Nvme.List[id]

    def get_list_by_args():
        """
        Class Method for getting a list off nvmes by args

        :return: List of Nvme
        """
        if args.devices:
            nvmes = []
            for nvme in args.devices:
                nvmes.append(Nvme.List[int(nvme[4:].split('n')[0])])
        else:
            nvmes = Nvme.List.values()
        return nvmes

    def header():
        """
        Class Method for printing header of table
        """
        ret = "-" * 80
        ret = (f"{ret}\n| {'ID':4} | {'Name':12} | {'NumaNode':8} | "
               f"{'OSDs':43} |\n{ret}")
        return ret

    def footer():
        """
        Class Method for printing header of table
        """
        return "-" * 80

    def __init__(self, name, node):
        """
        Instance Method for creating an instance of NVME

        :param name: the name of NVME
        :param core: the core of NVME
        """

        self.id = int(name[4:].split('n')[0])
        self.name = name
        self.numaNode = NumaNode.create(node)
        self.osds = []
        Nvme.List[self.id] = self
        self.numaNode.nvmes.append(self)

    def __str__(self):
        """
        Instance Method for displaying an instance of NVME

        :return: string representing instance
        """
        return (f"| {self.id:4} | {self.name:12} | {self.numaNode.id:8} | "
                f"{str(sorted([ int(osd.id) for osd in self.osds ])):43} |")

    def toJson(self):
        """
        Instance Method for json output an instance of Nvme

        :return: string representing instance in json format
        """

        ret = {'Id': self.id}
        ret['Name'] = self.name
        ret['Numa Node'] = self.numaNode.id
        ret['OSDs'] = sorted([osd.id for osd in self.osds])
        return ret

    def load_osds(self):
        """
        Instance Method for load OSDS in this NVME
        """
        c = run(["ceph-volume",
                 "lvm",
                 "list",
                 f"/dev/{self.name}",
                 "--format",
                 "json"],
                capture_output=True)
        try:
            osds = json.loads(c.stdout)
        except Exception:
            logger(f"No OSD found on NVME {self.id}")
            return
        for osdid, v in osds.items():
            self.osds.append(Osd.create(osdid, self))

    def create_osds(self):
        """
        Instance Method for create and pinning OSDS in this NVME
        """
        print(f"Creating OSD on {self.name}", file=sys.stderr)
        os.makedirs("/var/log/ceph-tools/", mode=0o0755, exist_ok=True)
        with open("/var/log/ceph-tools/create_osd.log", 'w') as f:
            run(["ceph-volume",
                 "lvm",
                 "batch",
                 "--osds-per-device",
                 str(args.osds),
                 f"/dev/{self.name}",
                 "--dmcrypt",
                 "--yes"],
                stdout=f,
                stderr=STDOUT)
        self.load_osds()
        for osd in self.osds:
            osd.pinning()


class Osd():
    """Class representing an OSD"""
    List = {}

    def create(id, nvme):
        """
        Class Method for creating an OSD and store it in 'List' class attribut

        :param id: the id of OSD
        :param nvme: the NVME where this OSD is located
        :return: the new OSD
        """

        if id not in Osd.List:
            Osd(id, nvme)
        return Osd.List[id]

    def header():
        """
        Class Method for printing header of table
        """
        ret = "-" * 50
        ret = (f"{ret}\n| {'OSD':4} | {'NVME':12} | {'CPU':4} | "
               f"{'Cpu Affinity':17} |\n{ret}")
        return ret

    def footer():
        """
        Class Method for printing header of table
        """
        return "-" * 50

    def __init__(self, id, nvme):
        """
        Instance Method for creating an instance of OSD

        :param id: the id of OSD
        :param nvme: the NVME where this OSD is located
        """

        self.id = id
        self.nvme = nvme
        Osd.List[id] = self
        self.realCpu = None
        self.realCpuAffinity = ''
        systemd_dir = "/etc/systemd/system"
        unit_path = os.path.join(systemd_dir,
                                 f"ceph-osd@{self.id}.service.d",
                                 "osd_cpu_pinning.conf")
        # Check if the pinning for the service is set
        if os.path.isfile(unit_path):
            with open(unit_path, 'r') as f:
                for line in f:
                    if 'CPUAffinity' in line:
                        self.realCpuAffinity = line.split('=')[1].strip()
                        break

    def __str__(self):
        """
        Instance Method for displaying an instance of Osd

        :return: string representing instance
        """

        return (f"| {self.id:4} | {self.nvme.name:12} | {self.realCpu:<4} | "
                f"{self.realCpuAffinity:17} |")

    def toJson(self):
        """
        Instance Method for json output an instance of Osd

        :return: string representing instance in json format
        """

        ret = {'Id': self.id}
        ret['NVME'] = self.nvme.id
        ret['CPU'] = self.realCpu
        ret['CPU Affinity'] = self.realCpuAffinity

    def check(self):
        """
        Instance Method for checking if OSD is correctly pinning

        :return: Result of checking
        """

        print(f"Check OSD {self.id}: ", end='', flush=True)
        systemd_dir = "/etc/systemd/system"
        unit_path = os.path.join(systemd_dir,
                                 f"ceph-osd@{self.id}.service.d",
                                 "osd_cpu_pinning.conf")
        # Check if the pinning for the service is set
        if not os.path.isfile(unit_path):
            logger(f"OSD {self.id} has no pinning configuration file")
            return False
        if self.nvme.numaNode.checkAffinity(self.realCpu) \
           and self.realCpu in parse_Cpu_Affinity(self.realCpuAffinity):
            print(f"OK -> {self.realCpu} = CPU Affinity {self.realCpuAffinity}"
                  f" on Numa Node {self.nvme.numaNode.cpuAffinity}")
            return True
        else:
            logger(f"OSD {self.id} does not run to correct CPU ({self.realCpu}"
                   f" in place of CPU Affinity {self.realCpuAffinity}) "
                   f"(Numa Node {self.nvme.numaNode.cpuAffinity})")
        return False

    def pinning(self):
        """
        Instance Method for pinning an OSD on a good CPU
        """

        numaNode = self.nvme.numaNode
        # If demand of pinning on NumaNode (range of CPUs) or not
        if args.numanode:
            self.realCpuAffinity = numaNode.cpuAffinity
        else:
            cpus = numaNode.get_free_cpus()
            self.realCpuAffinity = random.choice(cpus).id
        print(f"Pin OSD {self.id}: ", end='', flush=True)
        systemd_dir = "/etc/systemd/system"
        unit_dir_path = os.path.join(systemd_dir,
                                     f"ceph-osd@{self.id}.service.d")
        os.makedirs(unit_dir_path, mode=0o0755, exist_ok=True)
        unit_path = os.path.join(unit_dir_path, "osd_cpu_pinning.conf")
        with open(unit_path, 'w') as f:
            f.write(f"[Service]\n\tCPUAffinity={self.realCpuAffinity}\n")
        run(["systemctl", "daemon-reload"])
        run(["systemctl", "restart", f"ceph-osd@{self.id}.service"])
        map_osds_cpus()
        self.check()


def logger(message, priority="Warning"):
    """
    Function for logging

    :param message: The message to log
    :param priority: The priority of log
    """

    date = datetime.datetime.now()
    print(f"[{date}] {priority}: {message}", file=sys.stderr)


def init():
    """
    Function for initialisation instance of numaNode, Cpu, Nvme and Osd
    """

    print("Analyse Cpus", file=sys.stderr, end='', flush=True)
    c = run(["lscpu", "-e=CPU,CORE,NODE", "-J"], capture_output=True)
    try:
        cpus = json.loads(c.stdout)
    except Exception:
        logger("lscpu Error")

    for c in cpus['cpus']:
        print('.', file=sys.stderr, end='', flush=True)
        Cpu.create(c['cpu'], c['core'], c['node'])
    print('', file=sys.stderr)

    c = run(["lscpu", "-J"], capture_output=True)
    try:
        out = json.loads(c.stdout)
    except Exception:
        logger("lscpu Error")
    for id, node in NumaNode.List.items():
        node.set_cpuAffinity([cpuA['data'] for cpuA in out['lscpu']
                              if cpuA['field'] == f"NUMA node{id} CPU(s):"][0])

    print("Analyse NVMEs", file=sys.stderr, end='', flush=True)
    dir_to_nvme = "/sys/block/"
    output = run(['ceph',
                  'device',
                  'ls-by-host',
                  socket.gethostname(),
                  '-f',
                  'json'], capture_output=True)
    jso = json.loads(output.stdout)
    for devid in jso:
        for loc in devid['location']:
            print('.', file=sys.stderr, end='', flush=True)
            try:
                with open(os.path.join(dir_to_nvme, loc['dev'], 'device', 'numa_node')) as f:
                    node = f.read().splitlines()[0]
            except Exception as e:
                logger(f"Error for device {dir_to_nvme}": {e})
            nvme = Nvme.create(loc['dev'], node)
            nvme.load_osds()
    print('', file=sys.stderr)
    map_osds_cpus()


def map_osds_cpus():
    """
    Function for gwt the Realcpu of each OSD
    """
    for p in psutil.process_iter():
        if 'ceph-osd' in p.name():
            osdid = p.cmdline()[5]
            Osd.List[osdid].realCpu = int(p.cpu_num())
            Cpu.List[str(p.cpu_num())].osds.append(Osd.List[osdid])


if __name__ == "__main__":
    classDict = dict(inspect.getmembers(
          sys.modules[__name__],
          lambda member: inspect.isclass(member)
          and member.__module__ == __name__
      ))
    classtype = list(classDict.keys())

    parser = ArgumentParser(prog="ct-create-osd",
                            description="Create or check OSDs on NVME \
                                        and pinning to correct Cpu")
    subParser = parser.add_subparsers(title='Commands',
                                      help='Command Help',
                                      dest='command',
                                      required=True)

    createParser = subParser.add_parser('create',
                                        help='Create and pinning OSD')
    createParser.set_defaults(func=creating)
    createParser.add_argument("-n",
                              "--nvmes",
                              action="store",
                              default=default_nvme_to_configure,
                              type=int,
                              help=(f"Number of NVMEs to configure "
                                    f"({default_nvme_to_configure})"))
    createParser.add_argument("-o",
                              "--osds",
                              action="store",
                              default=default_osd_by_nvme,
                              type=int,
                              help=(f"Number of OSDs by Nvme "
                                    f"({default_osd_by_nvme})"))
    createParser.add_argument("-N",
                              "--numanode",
                              action="store_true",
                              help="Pinning on Numa Node instead CPU")
    createParser.add_argument("devices",
                              action="store",
                              type=str,
                              nargs='*',
                              help="NVME devices (all device if absent) , \
                                    eg: nvme0n1 nvme1n1")
    listParser = subParser.add_parser('list',
                                      help='List OSD pinning')
    listParser.set_defaults(func=listing)
    listParser.add_argument('-t',
                            '--type',
                            choices=classtype,
                            default='Nvme',
                            type=str,
                            help='wich type to list')
    listParser.add_argument('-j',
                            '--json',
                            action="store_true",
                            help="Display output in Json format")
    checkParser = subParser.add_parser('check',
                                       help='Check OSD pinning')
    checkParser.add_argument("devices",
                             action="store",
                             type=str,
                             nargs='*',
                             help="NVME devices (all device if absent) , \
                                  eg: nvme0n1 nvme1n1")
    checkParser.set_defaults(func=checking)
    pinningParser = subParser.add_parser('pin',
                                         help='Pinning OSD by NVME')
    pinningParser.set_defaults(func=pinning)
    pinningParser.add_argument("-N",
                               "--numanode",
                               action="store_true",
                               help="Pinning on Numa Node instead CPU")
    pinningParser.add_argument("devices",
                               action="store",
                               type=str,
                               nargs='*',
                               help="NVME devices (all device if absent) , \
                                     eg: nvme0n1 nvme1n1")

    args = parser.parse_args()
    args.func(args)
