#!/usr/bin/env python2
import argparse
import os
import sys
import getpass
import re
import time
import shutil

localied_work_dir = ""
localied_data_dir = ""
localied_conf_path = ""

def param_replace(string, rep_dict):
    pattern = re.compile("|".join([re.escape(k) for k in rep_dict.keys()]), re.M)
    return pattern.sub(lambda x: rep_dict[x.group(0)], string)

def prepare_directorys(args):
  global localied_conf_path
  global localied_data_dir
  global localied_vec_tmpdir
  global localied_work_dir
  localied_install_port_dir = "{ins_prefix}/{pt}".format(ins_prefix=args.install_prefix,pt=args.pg_protocal_port)
  if os.path.exists(localied_install_port_dir):
    sys.stderr.write("{lipd} already exists".format(lipd=localied_install_port_dir))
    sys.exit(1)
  os.makedirs(localied_install_port_dir)

  localied_data_dir = "{dpr}/{pgport}".format(dpr=args.datadir_prefix,pgport=args.pg_protocal_port)
  if os.path.exists(localied_data_dir) and len(os.listdir(localied_data_dir))!=0:
    sys.stderr.write("Invalid data dir: {} already exists and is not empty".format(localied_data_dir))
    sys.exit(1)
  os.makedirs(localied_data_dir)
  
  localied_vec_tmpdir = "{dpr}/{pgport}/vec_tmpdir".format(dpr=args.datadir_prefix,pgport=args.pg_protocal_port)
  #if os.path.exists(localied_vec_tmpdir) and len(os.listdir(localied_vec_tmpdir))!=0:
  #  sys.stderr.write("Invalid vec tmp dir: {} already exists and is not empty".format(localied_vec_tmpdir))
  #  sys.exit(1)
  #os.makedirs(localied_vec_tmpdir)

  # copy package to dest
  package_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  cmd = "cp -r {pkg} {lopr}".format(pkg= package_path,lopr=localied_install_port_dir)
  os.system(cmd)
  localied_work_dir = "{lopr}/{prgname}".format(lopr=localied_install_port_dir,prgname=str(args.prog_name))

def init_db(args):
  global localied_conf_path
  global localied_data_dir
  global localied_work_dir
  cmd0 = "export LD_LIBRARY_PATH=" + localied_work_dir + "/lib:$LD_LIBRARY_PATH;"
  initcmd = cmd0 + localied_work_dir + '/bin/initdb -D ' + localied_data_dir + ' 2>&1'
  os.system(initcmd) 

def modify_spec_conf(localied_conf_path):
  var_params = args.var_params
  
  var_dict={}
  if var_params != 'default':
    var_list = var_params.split(";")
    for item in var_list:
      if len(item):
        tmp = item.split("=")
        var_dict[tmp[0]] = tmp[1]
  
  new_content = ""
  del_list = []
  with open(localied_conf_path, 'r') as f:
    for line in f:
      print(line)
      if args.single_shard_mode == 'on':
        if line.startswith('single_shard_mode'):
          new_content = new_content + 'single_shard_mode = true\n'
          continue
      elif args.enable_global_mvcc == 'on':
        if line.startswith('enable_global_mvcc'):
          new_content = new_content + 'enable_global_mvcc = true\n'
          continue
      else:
        match_flag = False
        for key, val in var_dict.items():
          if line.startswith(key):
            new_content = new_content + key + ' = ' + val + '\n'
            del_list.append(key)
            match_flag = True
            
        
        if match_flag:
          continue

      new_content = new_content + line;
    
    f.close()

  for item in del_list:
    del var_dict[item]
  
  with open(localied_conf_path, 'w') as f:
    f.write(new_content)
    for key, val in var_dict.items():
      f.write(key + ' = ' + val + '\n')
    f.close()

def open_pg_extension(localied_conf_path):
  if os.path.exists(localied_vec_tmpdir) == False:
    os.makedirs(localied_vec_tmpdir)

  with open(localied_conf_path, 'a') as f:
    f.write("tornado_cn.enable_custom_planner = true\n")
    vec_tmpdir = "tornado_cn.data_dir='{tddir}'\n".format(tddir=localied_vec_tmpdir)
    f.write(vec_tmpdir)
    f.close()

def deal_conf(args):
  global localied_conf_path
  global localied_data_dir
  global localied_work_dir
  temp_conf_file = "{lwdr}/resources/postgresql.conf".format(lwdr= localied_work_dir)
  cp_conf = 'cp ' + temp_conf_file + ' ' + localied_data_dir
  os.system(cp_conf)
  localied_conf_path = "{ddr}/postgresql.conf".format(ddr=localied_data_dir)
  config_template = open(localied_conf_path,'r').read()
  replace_items = {}
  replace_items['port_placeholder'] = str(args.pg_protocal_port)
  replace_items['mysql_port_placeholder'] = str(args.mysql_protocal_port)
  replace_items['comp_node_id_placeholder'] = str(args.compute_node_id) 
  replace_items['unix_socket_dir_place_holder'] = localied_data_dir

  conf_content = param_replace(config_template,replace_items)
  conf_file = open(localied_conf_path,'w')
  conf_file.write(conf_content)
  conf_file.close()
    
  if args.single_shard_mode == 'on' or args.enable_global_mvcc == 'on' or args.var_params != 'default':
    modify_spec_conf(localied_conf_path)
  open_pg_extension(localied_conf_path)

  if args.with_hostname:
    os.system('echo "host all all ' + args.bind_address + ' trust" >> ' + localied_data_dir + '/pg_hba.conf')
  else:
    os.system('echo "host all all ' + args.bind_address + '/32 trust" >> ' + localied_data_dir + '/pg_hba.conf')
  os.system('echo "host all all 127.0.0.1/32  trust" >> ' + localied_data_dir + '/pg_hba.conf')
  os.system('echo "host all agent 0.0.0.0/0 reject" >> ' + localied_data_dir + '/pg_hba.conf')
  os.system('echo "host all all 0.0.0.0/0 md5" >> ' + localied_data_dir + '/pg_hba.conf') 

def boot_postgres(args):
  global localied_conf_path
  global localied_data_dir
  global localied_work_dir
  pg_logfp = localied_data_dir + "/logfile-" + str(args.pg_protocal_port)

  prefix = ""
  if args.valgrind:
      prefix = "valgrind --leak-check=full --trace-children=yes --undef-value-errors=no "

  cmd0 = "export LD_LIBRARY_PATH=" + localied_work_dir + "/lib:$LD_LIBRARY_PATH;"
  cmd1 = 'export LD_PRELOAD="' + localied_work_dir + '/resources/libjemalloc.so"; ulimit -c unlimited; '
  startup_cmd = cmd0 + cmd1 + prefix + localied_work_dir + '/bin/postgres -D ' + localied_data_dir + " > " + pg_logfp + " 2>&1 &"
  ret = os.system(startup_cmd)
  if ret != 0:
    sys.stderr.write("Start postgres failed, please check the logfile {pglfp} for more information".format(pglfp=pg_logfp))
    sys.exit(1)
  time.sleep(5); # wait for postgres to startup

  # add initial user for clients to use later.
  # TODO: use more restricted privs than superuser
  sql = "set skip_tidsync = true; CREATE USER agent PASSWORD 'agent_pwd' superuser; CREATE USER " + args.user + " PASSWORD '" + args.password + '\' superuser;'
  psql_cmd = cmd0 + localied_work_dir + "/bin/psql -h {datadir} -p".format(datadir=localied_data_dir) + str(args.pg_protocal_port) + " -U " + getpass.getuser() + " -d postgres -c \"" + sql + "\""
  ret = os.system(psql_cmd)
  if ret != 0:
    sys.stderr.write("Kunlun-server instance initialization failed, please check the logfile {pglfp} for more information".format(pglfp=pg_logfp))
    sys.exit(1)
  # add mysql functions for mysql clients.
  psql_cmd2 = cmd0 + localied_work_dir + "/bin/psql -h {datadir} -p".format(datadir=localied_data_dir) + str(args.pg_protocal_port) + " -U " + getpass.getuser() + " -d template1 -f " + localied_work_dir + "/scripts/mysql_funcs.sql"
  ret = os.system(psql_cmd2)
  if ret != 0:
    sys.stderr.write("Kunlun-server instance initialization failed, please check the logfile {pglfp} for more information".format(pglfp=pg_logfp))
    sys.exit(1)
  psql_cmd2 = cmd0 + localied_work_dir + "/bin/psql -h {datadir} -p".format(datadir=localied_data_dir) + str(args.pg_protocal_port) + " -U " + getpass.getuser() + " -d postgres -f " + localied_work_dir + "/scripts/mysql_funcs.sql"
  ret = os.system(psql_cmd2)
  if ret != 0:
    sys.stderr.write("Kunlun-server instance initialization failed, please check the logfile {pglfp} for more information".format(pglfp=pg_logfp))
    sys.exit(1)

  # append the new instance's port to datadir mapping into instance_list.txt
  etc_path = localied_work_dir + "/etc"
  list_file = etc_path+"/instances_list.txt" 
  if not os.path.exists(etc_path):
      os.mkdir(etc_path)
  os.system("echo \"" + str(args.pg_protocal_port) + "==>" + localied_data_dir + "\" >> " + list_file)

def rollBackInstallaiton(args):
    error_count = 0
    workpath = "{prefix}/{pt}/{prgname}".format(
        prefix=os.path.abspath(args.install_prefix), pt=str(args.pg_protocal_port), prgname=str(args.prog_name)
    )
    bin_path = "{wp}/bin".format(wp=workpath)

    datadir = os.path.abspath(args.datadir_prefix) + "/" + str(args.pg_protocal_port)

    # 1. Stop postgres instance if already running by invoking stop_mysql.sh
    cmd_stop_pg = "cd {dp};./pg_ctl -D {d_dir} stop".format(dp=bin_path,d_dir=str(datadir))
    try:
        os.system(cmd_stop_pg)
    except OSError as err:
        sys.stderr.write("Stop Postgres instance failed: {}".format(str(err)))
        error_count = error_count +1

    # 2. backup pg_log
    file_cwd = os.path.dirname(os.path.realpath(__file__))
    package = os.path.dirname(file_cwd)
    program_path = os.path.dirname(package)
    install_err_log_pg = "{pp}/install_pg_err".format(pp=program_path)

    if os.path.isdir(install_err_log_pg) == False:
        os.makedirs(install_err_log_pg, 0o777)
    data_dir = "{dpr}/{pgport}".format(dpr=args.datadir_prefix,pgport=args.pg_protocal_port)
    pg_logfp = data_dir + "/logfile-" + str(args.pg_protocal_port)
    cmd_backup_log = "cp {pgl} {ielp}/".format(pgl=pg_logfp,ielp=install_err_log_pg)
    #print(cmd_backup_log)
    try:
        os.system(cmd_backup_log)
    except OSError as err:
        sys.stderr.write("Stop Postgres instance failed: {}".format(str(err)))
        error_count = error_count +1
    
    # 3. remove the related directory
    try:
        # 3.1 remove the Postgres binary residential directory
        runing_dir = os.path.abspath(args.install_prefix) + "/" + str(args.pg_protocal_port)
        shutil.rmtree(runing_dir,ignore_errors=True)

        # 3.2 remove the Postgres data residential directory
        data_dir = os.path.abspath(args.datadir_prefix) + "/" + str(args.pg_protocal_port)
        shutil.rmtree(data_dir,ignore_errors=True)

    except OSError as err:
        sys.stderr.write("Remove the Instance data and log failed: {}".format(str(err)))
        error_count = error_count+1
    
    if error_count > 0:
      return False
    
    return True


if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="Install computing node")
  parser.add_argument('--user',type=str,help="The user for connect postgresql process",required=True)
  parser.add_argument('--password',type=str,help="The password for user connect Postgres",required=True)
  parser.add_argument('--bind_address',type=str,help="The binding ip address of Postgres",required=True)
  parser.add_argument('--pg_protocal_port',type=int,help="The port listen in Postgresql protocal",required=True)
  parser.add_argument('--mysql_protocal_port',type=int,help="The port listen in MySQL Protocal",required=True)
  parser.add_argument('--compute_node_id',type=int,help="The compute node id in the cluster",required=True)
  parser.add_argument('--datadir_prefix',type=str,help="Postgres data dir prefix",required=True)
  parser.add_argument('--install_prefix',type=str,help="Postgres install prefix",required=True)
  parser.add_argument('--prog_name',type=str,help="Postgres pkg name",required=True)
  parser.add_argument('--single_shard_mode',type=str,help="Postgres single shard mode",required=True)
  parser.add_argument('--enable_global_mvcc',type=str,help="Postgres enable global mvcc",required=True)
  parser.add_argument('--var_params',type=str,help="Postgres global variable params",required=True)
  parser.add_argument('--valgrind',type=bool,help="Start Postgresql with valgrind or not",required=False,default=False)
  parser.add_argument('--with_hostname',type=bool,help="Start Postgresql with hostname",required=False,default=False)
  args= parser.parse_args()

  try:
    # Prepare the related directorys
    prepare_directorys(args)

    # do initdb
    init_db(args)
    
    # deal conf
    deal_conf(args)
    
    # boot postgres
    boot_postgres(args)
  except:  
    rollBackInstallaiton(args)
    sys.stderr.flush()
    sys.exit(1)

  sys.exit(0)