#!/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

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 __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: {sorted([int(cpu.id) for cpu in self.cpusPrimaries])}"
    ret = f"{ret}\n\tSecondaries Cpus: {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 : {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} | {'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} | {self.primary:<8} | {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} | {'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} | {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:
      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} | {'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
    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
    else:
      self.realCpuAffinity = None

  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} | {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} on Numa Node {self.nvme.numaNode.cpuAffinity}")
      return True
    else:
      logger(f"OSD {self.id} does not run to correct CPU ({self.realCpu} in place of CPU Affinity {self.realCpuAffinity}) (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:
    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:
    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/"
  nvmes = [ i for i in os.listdir(dir_to_nvme) if i.startswith('nvme')]
  for nv in nvmes:
    print('.', file=sys.stderr, end='', flush=True)
    with open(os.path.join(dir_to_nvme, nv, 'device', 'numa_node')) as f:
      node = f.read().splitlines()[0]
    nvme = Nvme.create(nv, 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 ({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 ({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)
