#!/usr/bin/python

from collections import defaultdict
from rpmUtils.miscutils import splitFilename
import argparse
import json
import os.path
import re
import rpm
import sys
import yum

conf_file = '/etc/cta/cta-release.conf'

vfiles = {
  'cta': '/etc/yum/pluginconf.d/versionlock.cta',
  'yum': '/etc/yum/pluginconf.d/versionlock.list'
}
versionlock_file = '/etc/yum/pluginconf.d/versionlock.cta'
actions = ('check', 'apply', 'forceupdate', 'checkpkg', 'remove')
configs = ('eos4', 'eos5', 'xrootd4', 'xrootd5')
retcode = 0
jsonoutput = False
no_vlock = '/etc/cta.novlock'
eos_version=0
xrootd_version=0

try:
  FileNotFoundError
except NameError:
  FileNotFoundError = IOError

def usage():
  print("\n%s: command line tool to manage cta packages versionlock\n \
\n \
usage: %s help|check|apply|forceupdate|checkpkg|remove [--json]\n \
       %s config eos4|eos5 [--json]\n \
       %s config xrootd4|xrootd5 [--json]\n \
  help: print this message\n \
  check: show consistency of versionlock.list with cta versionlock requirements\n \
  apply: add cta versions to versionlock.list\n \
  forceupdate: add cta versions to versionlock.list and overwrite already defined packages\n \
  checkpkg: check consistency between installed packages and cta versionlock requirements\n \
  remove: remove all cta packages and dependencies from versionlock.list\n \
  config eos4|eos5: select EOS version\n \
  config xrootd4|xrootd5: select XRootD version\n \
  [--json]: format output in json\n \
  \n \
  to prevent cta-versionlock to alter versionlock.list create a %s file.\n \
  " % (sys.argv[0], sys.argv[0], sys.argv[0], no_vlock))
  exit(1)

START_EXPR = "\s*#+\s*%s-START\s*#+"
END_EXPR = "\s*#+\s*%s-END\s*#+"

yb = yum.YumBase()
yb.preconf.debuglevel = 0
yb.preconf.errorlevel = 0
repos_all = [repo.id for repo in yb.repos.findRepos("cta-*")]

# if path doesn't exist, create it
if not os.path.exists(os.path.dirname(conf_file)):
  os.makedirs(os.path.dirname(conf_file))
# If file doesn't exists, create a new file
if not os.path.isfile(conf_file):
  with open(conf_file, 'w') as f:
    print("Creating new config file %s" % conf_file)
    cta_conf_json = {'eos_major_version': 4, 'xrootd_major_version': 4, 'cta_major_version': 4}
    f.write(json.dumps(cta_conf_json))
with open(conf_file) as f:
  cta_conf_json = json.load(f)

def get_repo_set_for_eos_version(eos_major_version):
  if eos_major_version == 4:
    return {repo_id for repo_id in repos_all if "eos" in repo_id and "diopside" not in repo_id}
  elif eos_major_version == 5:
    return {repo_id for repo_id in repos_all if "eos" in repo_id and "citrine" not in repo_id}
  else:
    raise RuntimeError("Unknown EOS major version '%s'." % eos_major_version)

def get_repo_set_for_cta_version(cta_major_version):
  if cta_major_version == 4:
    return {repo_id for repo_id in repos_all if "eos" not in repo_id and "public-5" not in repo_id}
  elif cta_major_version == 5:
    return {repo_id for repo_id in repos_all if "eos" not in repo_id and "public-4" not in repo_id}
  else:
    raise RuntimeError("Unknown CTA major version '%s'." % cta_major_version)

def switch_repos(repo_id_list, enable):
  for repo_id in repo_id_list:
    yb.repos.enableRepo(repo_id) if enable else yb.repos.disableRepo(repo_id)
    repo_obj = yb.repos.findRepos(repo_id)[0]
    with open(repo_obj.repofile, "w") as f:
      repo_obj.write(f, always=("enabled",))

def enable_repos(repo_id_list):
  switch_repos(repo_id_list, True)

def disable_repos(repo_id_list):
  switch_repos(repo_id_list, False)

def switch_packages(file_path, key, enable_packages):

  start_expr_key = START_EXPR % key
  end_expr_key = END_EXPR % key

  new_lines = []

  with open(file_path) as f:
    plist = f.read().splitlines()
    found_key = False
    for p in plist:

      if not found_key and re.match(start_expr_key, p):
        found_key = True
        new_lines.append(p)
        continue

      if found_key and re.match(end_expr_key, p):
        found_key = False
        new_lines.append(p)
        continue

      if found_key and enable_packages and p.startswith('#'):
        new_lines.append(p[1:].strip())
      elif found_key and not enable_packages and not p.startswith('#'):
        new_lines.append("# " + p.strip())
      else:
        new_lines.append(p.strip())

  with open(file_path, 'w') as f:
    for p in new_lines:
      f.write(p + '\n')

def _exit():
  printer(message)
  exit(retcode)

# output formatting
def printer(msg):
  if jsonoutput:
    print(json.dumps(msg))
  else:
    print("\n%s" % msg.pop('title'))
    for status in ['ok', 'wrong_version', 'missing', 'adding', 'updating', 'removing']:
      if msg[status]:
        print("\n==> %s: %s/%s" % (status, msg[status]['count'], len(versions['cta'])))
        if (status == 'wrong_version' or status == 'updating') and msg[status]['count'] > 0:
          print("\n".join(['%s: %s' % (key, value) for (key, value) in msg[status]['content'].items()]))
        else:
          print("\n".join(msg[status]['content']))
    print("\n")

# read versionlock files
def readVer(filelist):
  versions = defaultdict(dict)
  for fname, vfile in filelist.items():
    if not os.path.isfile(vfile):
      raise FileNotFoundError("file %s not found" % vfile)

    with open(vfile) as f:
      plist = f.read().splitlines()
      for p in plist:
        if p == "" or p.startswith('#'):
          continue
        (n, v, r, e, a) = splitFilename(p)
        versions[fname][n] = [e, v, r, a]
  return versions

# Compare versions in both versionlock lists
def checkVfile(pkglist):
  result = defaultdict(dict)
  for p, version in pkglist.items():
    try: 
      versions['yum'][p]
    except:
      result.setdefault('missing', []).append(p)
      continue
    else:
      if versions['yum'][p][:-1] == version[:-1]:
        result.setdefault('ok',[]).append(p)
      else:
        result['wrong_version'][p]= {'current': ("%s:%s-%s" % tuple(versions['yum'][p][:-1])), 'required': ("%s:%s-%s" % tuple(version[:-1]))}
  return result

# add CTA packages to versionlock.list
def addtoVfile(pkglist):
  with open(vfiles['yum'], 'a' ) as f:
    for p in pkglist:
      (e, v, r, a) = versions['cta'][p][:]
      package = ("%s:%s-%s-%s.%s" % (e, p, v, r, a))
      f.write(package + '\n')

# update existing packages in versionlock.list with CTA versions required
def updateVfile(pkglist):
  with open(vfiles['yum'], 'r+') as f:
    content = f.read()
    for p in pkglist:
      ver = {}
      for t in 'yum', 'cta':
        (e, v, r, a) = versions[t][p][:]
        ver[t] = ("%s:%s-%s-%s.%s" % (e, p, v, r, a))
      content = re.sub(ver['yum'], ver['cta'], content)

    f.seek(0)
    f.write(content)
    f.truncate()

# check installed packages version
def checkPkg(pkglist):
  result = defaultdict(dict)
  ts = rpm.TransactionSet()
  for pname, verlist in pkglist.items():
    mi = ts.dbMatch('name', pname)
    for item in mi:
      ie = str("0" if item['epoch'] is None else item['epoch'])
      (e, v, r, a) = verlist
      if (e, v, r) == (ie, item['version'], item['release']):
        result['ok'][pname] = ("%s:%s-%s" % (e, v, r))
      else:
        result['wrong_version'][pname] = {'required': ("%s:%s-%s" % (e, v, r)), 'installed': ("%s:%s-%s" % (ie, item['version'], item['release']))}
  return result

# remove CTA packages from versionlock.list
def clearVfile(pkglist):
  with open(vfiles['yum'], 'r+') as f:
    content = f.read()
    for p in pkglist:
      (e, v, r, a) = versions['cta'][p][:]
      line = ("%s:%s-%s-%s.%s\n" % (e, p, v, r, a))
      content = re.sub(line, '', content)

    f.seek(0)
    f.write(content)
    f.truncate()

# Write previous statements but using argparse
parser = argparse.ArgumentParser(description='Manage CTA dependencies')
subparsers = parser.add_subparsers(dest='command')
config_parser = subparsers.add_parser('config')
config_parser.add_argument('config_name', choices=configs, help='Name of the configuration file')
for action in actions:
  subparsers.add_parser(action)
parser.add_argument('--file', help='use yum versionlock file')
parser.add_argument('--json', action='store_true', help='output in json format')
args = parser.parse_args()

# Parse config
if args.command == 'config':
  if args.config_name == 'eos4':
    eos_version = 4
  elif args.config_name == 'eos5':
    eos_version = 5
  elif args.config_name == 'xrootd4':
    xrootd_version = 4
  elif args.config_name == 'xrootd5':
    xrootd_version = 5

if args.json:
  jsonoutput = True

if args.file:
  versionlock_file = args.file
  print("Using yum versionlock file: %s" % versionlock_file)  

# check arguments
if (args.command == 'check'):
  if (eos_version and xrootd_version):
    print("Error: cannot check both EOS and XRootD versions")
    usage()
  elif (not eos_version and not xrootd_version):
    print("Error: no EOS or XRootD version was provided")
    usage()

if (not eos_version and not xrootd_version)  and (action == 'config'):
  print("Error: no EOS version was provided")
  usage()

message = defaultdict(dict)

if (args.command == 'config'):
  if eos_version is 4:
    message['title'] = "Enabling EOS-4 dependencies"
    cta_conf_json["eos_major_version"] = 4
    switch_packages(versionlock_file, "EOS-4", True)
    switch_packages(versionlock_file, "EOS-5", False)
  elif eos_version is 5:
    message['title'] = "Enabling EOS-5 dependencies"
    cta_conf_json["eos_major_version"] = 5
    switch_packages(versionlock_file, "EOS-4", False)
    switch_packages(versionlock_file, "EOS-5", True)
  if xrootd_version is 4:
    message['title'] = "Enabling XROOTD-4 dependencies"
    cta_conf_json["xrootd_major_version"] = 4
    switch_packages(versionlock_file, "XROOTD-4", True)
    switch_packages(versionlock_file, "XROOTD-5", False)
  elif xrootd_version is 5:
    message['title'] = "Enabling XROOTD-5 dependencies"
    cta_conf_json["xrootd_major_version"] = 5
    switch_packages(versionlock_file, "XROOTD-4", False)
    switch_packages(versionlock_file, "XROOTD-5", True)

  # Update value in config file
  with open(conf_file, "w") as f:
    f.write(json.dumps(cta_conf_json))
  _exit()

repos_selected = set()
repos_selected |= get_repo_set_for_cta_version(cta_conf_json["cta_major_version"])
repos_selected |= get_repo_set_for_eos_version(cta_conf_json["eos_major_version"])

# check if CTA packages exist in versionlock.list (ignore arch)
versions = readVer(vfiles)
versionlock = checkVfile(versions['cta'])

# return versionlock summary
if (action == 'check'):
  message['title'] = "Yum versionlock status for CTA packages and dependencies"
  for status, content in versionlock.items():
    message[status]['count'] = len(content)
    message[status]['content'] = content
  if versionlock['missing'] or versionlock['wrong_version']:
    retcode = 2

# add CTA packages to versionlock.list
elif (action == 'apply'):
  message['title'] = "Adding CTA packages and dependencies to versionlock.list"

  # Enable CTA repos
  enable_repos(repos_selected)

  if os.path.isfile(no_vlock):
    message['adding']['count'] = 0
    message['adding']['content'] = [ "cta_novlock file present, doing nothing" ]
    retcode = 1
    _exit()
  if versionlock['missing']:
    message['adding']['count'] = len(versionlock['missing'])
    message['adding']['content'] = versionlock['missing']
    addtoVfile(versionlock['missing'])
  else:
    message['adding']['count'] = 0
    message['adding']['content'] = [ "Nothing to do" ]
  if versionlock['wrong_version']:
    message['title'] += "\nWARNING: some packages have a different version specified in versionlock.list than required by CTA, it will not be changed unless you use the 'forceupdate' option"
    message['wrong_version']['count'] = len(versionlock['wrong_version'])
    message['wrong_version']['content'] = versionlock['wrong_version']
    retcode = 2

# add CTA packages and overwrite existing versions in versionlock.list
elif (action == 'forceupdate'):
  message['title'] = "Adding and updating CTA packages and dependencies in versionlock.list"

  # Enable CTA repos
  enable_repos(repos_selected)

  if os.path.isfile(no_vlock):
    message['updating']['count'] = 0
    message['updating']['content'] = [ "cta_novlock file present, doing nothing" ]
    retcode = 1
    _exit()
  if versionlock['wrong_version']:
    message['updating']['count'] = len(versionlock['wrong_version'])
    message['updating']['content'] = versionlock['wrong_version']
    updateVfile(versionlock['wrong_version'])
  if versionlock['missing']:
    message['adding']['count'] = len(versionlock['missing'])
    message['adding']['content'] = versionlock['missing']
    addtoVfile(versionlock['missing'])
  if (not versionlock['missing'] and not versionlock['wrong_version']):
    message['updating']['count'] = 0
    message['updating']['content'] = [ "Nothing to do" ]

# check version of currently installed packages
elif (action == 'checkpkg'):
  message['title'] = "CTA packages and dependencies versions currently installed"
  packages = checkPkg(versions['cta'])
  for status, content in packages.items():
    message[status]['count'] = len(content)
    message[status]['content'] = content
  if packages['wrong_version']:
    retcode = 2

# remove CTA related packages from versionlock.list
elif (action == 'remove'):
  message['title'] = "Removing CTA packages and dependencies from versionlock.list"

  # Disable all CTA repos
  disable_repos(repos_all)

  if os.path.isfile(no_vlock):
    message['removing']['count'] = 0
    message['removing']['content'] = [ "cta_novlock file present, doing nothing" ]
    retcode = 1
    _exit()
  if versionlock['ok']:
    message['removing']['count'] = len(versionlock['ok'])
    message['removing']['content'] = versionlock['ok']
    clearVfile(versionlock['ok'])
  else:
    message['removing']['count'] = 0
    message['removing']['content'] = [ "Nothing to do" ]

_exit()