#!/usr/bin/env python
# -*-python-*-

import os
import sys
import glob
import getopt
import platform
import re
import shutil
import string
import subprocess

def usage(err=0):
    print >> sys.stderr, """
Usage: %s [-c path/to/previous/config/dir] [-a AUTO] [-d FILENAME] [-n] [-s FACTOR]

  -c <path/to/previous/config/dir>
       -- example: -c /etc/opt/membase/1.6.5.3.1

  -a <yes|no>
       -- automatic or non-interactive mode; default is 'no';
          'yes' to force automatic 'yes' answers to all questions

  -d <dbdir_output_file>
       -- retrieve db directory from config file and exit

  -n   -- dry-run; don't actually change anything

  -s <free_disk_space_needed_factor>
       -- free disk space needed, as a factor of current bucket usage
       -- default value is 2.0
       -- example: -s 1.0
""" % (os.path.basename(sys.argv[0]),)
    sys.exit(err)

def parse_args(args):
    prev_cfg_dir = None
    interactive  = True
    dry_run      = False
    space_needed_factor = 2.0
    dbdir_file   = None

    try:
        opts, args = getopt.getopt(args, 'hc:a:nsd:', ['help'])
    except getopt.GetoptError, e:
        usage(e.msg)

    for (o, a) in opts:
        if o == '--help' or o == '-h':
            usage()
        elif o == '-c':
            prev_cfg_dir = a
        elif o == '-a':
            interactive = re.match("y", a) == None
        elif o == '-n':
            dry_run = True
        elif o == '-s':
            space_needed_factor = float(a)
        elif o == '-d':
            dbdir_file = a
        else:
            usage("unknown option - " + o)

    return prev_cfg_dir, interactive, dry_run, space_needed_factor, dbdir_file

def find_cfg(src_dir, root, patterns):
    for pattern in patterns:
        cfgs = glob.glob(os.path.join(src_dir, pattern))
        if len(cfgs) > 1:
            sys.exit("Error: found multiple " + root + " candidates: " + str(cfgs))
        if len(cfgs) == 1:
            return cfgs[0]

    return None

def copy_cfg(cfg, dst, dry_run=False):
    if os.path.exists(dst) and os.path.getsize(dst) > 0:
        sys.exit("Error: " + dst + " already exists while copying " + cfg)

    print(dry_run_prefix(dry_run) + "Copying " + cfg)
    print(dry_run_prefix(dry_run) + "    cp " + cfg + " " + dst)
    if not dry_run:
        if cfg != dst:
            shutil.copyfile(cfg, dst)

    return True

def dry_run_prefix(dry_run):
    if dry_run:
        return "SKIPPED (dry-run): "
    return ""

def prompt(msg, expect='yes'):
    global interactive

    if interactive:
        print(msg)
        if not re.match(expect, sys.stdin.readline()):
            sys.exit("Did not receive a '" + expect + "', leaving.")

def confirm(install_dir, prev_cfg_dir, prev_cfg, prev_ip,
            dbdir, buckets):
    print("\nMembase should not be running.")
    if platform.system() == 'Windows':
      print("  Please use the Control Panel to stop the Membase service.")
    else:
      print("  Please use: /etc/init.d/membase-server stop")
    prompt("Is the Membase server already stopped? [yes|no]")

    print("\nDatabase dir: " + dbdir)
    prompt("Is that the expected database directory to upgrade? [yes|no]")

    print("\nBuckets to upgrade: " + ",".join(buckets or ["[no buckets found]"]))
    prompt("Are those the expected buckets to upgrade? [yes|no]")

def upgrade(bin_dir, install_dir, prev_cfg_dir, prev_cfg, prev_ip,
            dbdir, buckets, dry_run=False):
    """File modification steps are grouped into this upgrade() function.
    """

    if prev_cfg:
        copy_cfg(prev_cfg,
                 os.path.join(install_dir,
                              "var", "lib", "membase", "config", "config.dat"),
                 dry_run=dry_run)
    if prev_ip:
        copy_cfg(prev_ip,
                 os.path.join(install_dir,
                              "var", "lib", "membase", "ip"),
                 dry_run=dry_run)

    print("Ensuring bucket data directories.")
    for bucket in buckets:
        bucket_dir = dbdir + "/" + bucket + "-data"

        print(dry_run_prefix(dry_run) \
              + "Ensuring bucket data directory: " + bucket_dir)
        print(dry_run_prefix(dry_run) \
              + "    mkdir -p " + bucket_dir)
        if (not dry_run) and (not os.path.isdir(bucket_dir)):
            os.makedirs(bucket_dir)

    if platform.system() != 'Windows':
        cmd = ['chown', '-R', 'membase:membase', dbdir]
        print(dry_run_prefix(dry_run) + "Ensuring dbdir owner/group: " + dbdir)
        print(dry_run_prefix(dry_run) + "    " + " ".join(cmd))
        if not dry_run:
            p = subprocess.Popen(cmd)
            err = os.waitpid(p.pid, 0)[1]
            if err != 0:
                sys.exit("ERROR: chown dbdir failed: " + dbdir + " err: " + str(err))

    print("Upgrading buckets.")
    for bucket in buckets:
        # The '/' is required for Windows, too, due to backslash
        # getting treated as an escape character, and python does the
        # right cross-platform thing.
        bucket_dir = dbdir + "/" + bucket + "-data"

        if (os.path.isfile(os.path.join(bucket_dir, bucket))):
            print("Skipping already converted bucket: " + bucket_dir)
            continue

        if platform.system() == 'Windows':
            cmd = ["mbdbupgrade.exe",
                   dbdir + "/" + bucket,
                   bucket_dir]
        else:
            cmd = [os.path.join(bin_dir, "mbdbupgrade"),
                   os.path.join(dbdir, bucket),
                   bucket_dir]

        print(dry_run_prefix(dry_run) + "Upgrading bucket: " + bucket)
        print(dry_run_prefix(dry_run) + "    " + " ".join(cmd))
        if not dry_run:
            p = subprocess.Popen(cmd)
            print("mbdbupgrade pid: " + str(p.pid))

            if platform.system() == 'Windows':
                p.wait()
            else:
                err = os.waitpid(p.pid, 0)[1]
                print("mbdbupgrade err: " + str(err))
                if err != 0:
                    sys.exit("ERROR: upgrade failed for bucket: " + bucket)

        if platform.system() != 'Windows':
            cmd = ['chown', '-R', 'membase:membase', bucket_dir]
            print(dry_run_prefix(dry_run) + "Ensuring bucket owner/group: " + bucket_dir)
            print(dry_run_prefix(dry_run) + "    " + " ".join(cmd))
            if not dry_run:
                p = subprocess.Popen(cmd)
                err = os.waitpid(p.pid, 0)[1]
                if err != 0:
                    sys.exit("ERROR: chown bucket failed: " + bucket_dir + " err: " + str(err))

def bucket_files(dbdir, bucket):
    shards = glob.glob(os.path.join(dbdir, bucket + "-*.mb")) or \
             glob.glob(os.path.join(dbdir, bucket + "-*.sqlite"))
    shards.sort()
    return [os.path.join(dbdir, bucket)] + shards

def get_free_space(path):
    fs_stat = None
    try:
        fs_stat = os.statvfs(path)
        return fs_stat.f_bsize * fs_stat.f_bavail
    except:
        pass

    # http://stackoverflow.com/questions/51658/cross-platform-space-remaining-on-volume-using-python
    #
    if platform.system() == 'Windows':
        import ctypes

        free_bytes = ctypes.c_ulonglong(0)
        ctypes.windll.kernel32.GetDiskFreeSpaceExW(ctypes.c_wchar_p(path),
                                                   None, None,
                                                   ctypes.pointer(free_bytes))
        return free_bytes.value

    return -1 # Don't know.

def main():
    bin_dir     = os.path.dirname(sys.argv[0])
    if platform.system() == 'Windows':
        install_dir = os.path.join(bin_dir, '..')
    else:
        install_dir = os.path.dirname(bin_dir)

    path = [bin_dir,
            os.path.join(bin_dir, 'erlang', 'bin'),
            os.environ['PATH']]

    if os.name == 'posix':
        os.environ['PATH'] = ':'.join(path)
    elif os.name == 'nt':
        os.environ['PATH'] = ';'.join(path)

    global prev_cfg_dir, interactive, dry_run, space_needed_factor, dbdir_file

    prev_cfg_dir, interactive, dry_run, space_needed_factor, dbdir_file = \
       parse_args(sys.argv[1:])

    if dry_run:
        print("Dry-run mode: no actual upgrade changes will be made.")

    if not interactive:
        print("Automatic mode: running without interactive questions or confirmations.")

    if not prev_cfg_dir:
        sys.exit("Error: please provide a previous config directory path (-c path)")
    if not os.path.isdir(prev_cfg_dir):
        sys.exit("Error: previous config directory incorrect: " + prev_cfg_dir)

    # ----------------------------------------------------

    ver = None
    ver_path = os.path.join(install_dir, "VERSION.txt")
    try:
        f = open(ver_path, 'r')
        ver = f.read()
        f.close()
    except:
        sys.exit("Unexpected: missing " + ver_path + " file")
    if not ver:
        sys.exit("Unexpected: empty " + ver_path)

    print("Upgrading your Membase Server to " + string.strip(ver) + ".")
    print("The upgrade process might take awhile.")
    print("Analysing...")

    # ----------------------------------------------------

    if not dbdir_file:
        prev_cfg = find_cfg(prev_cfg_dir, "config.dat",
                        # Look for something like /etc/opt/membase/ns_1/config.dat
                        # Possibly config.dat.rpmsave or config.dat.debsave
                        [os.path.join('ns_1', 'config.dat'),
                         os.path.join('ns_1', 'config.dat.*save'),
                         '/opt/membase/var/lib/membase/config/config.dat.*save'])
    else:
        prev_cfg = find_cfg(prev_cfg_dir, "config.dat",
                        # Look for something like /etc/opt/membase/ns_1/config.dat
                        # Possibly config.dat.rpmsave or config.dat.debsave
                        [os.path.join('ns_1', 'config.dat'),
                         os.path.join('ns_1', 'config.dat.*save'),
                         'config.dat',
                         '/opt/membase/var/lib/membase/config/config.dat.*save'])
    if not prev_cfg:
        print("Done: no previous config.dat found; nothing to upgrade.")
        sys.exit(0)

    print("Previous config.dat file is " + prev_cfg)

    # ----------------------------------------------------

    prev_ip = find_cfg(prev_cfg_dir, "ip", ['/opt/membase/var/lib/membase/ip.*save',
                                            'ip.*save',
                                            'ip',
                                            '/opt/membase/var/lib/membase/ip'])
    ip = None
    if prev_ip:
        try:
            f = open(prev_ip, 'r')
            ip = string.strip(f.read())
            f.close()
        except:
            pass
    if (not ip) or (len(ip) <= 0):
        if platform.system() == 'Windows':
            ip = subprocess.Popen("ip_addr.bat",
                                  stdout=subprocess.PIPE).communicate()[0]
        else:
            ip = '127.0.0.1'

    # ----------------------------------------------------

    node = 'ns_1@' + ip

    print("Target node: " + node)

    escript_cmd = "escript"
    if platform.system() == 'Windows':
        escript_cmd = "escript.exe"

    node_cfg = subprocess.Popen([escript_cmd,
                                 os.path.join(bin_dir, "mbdumpconfig.escript"),
                                 prev_cfg, "node", node],
                                stdout=subprocess.PIPE).communicate()[0]
    node_cfg = string.strip(node_cfg)
    if len(node_cfg) <= 0:
        print("Done: previous node configuration is empty.")
        sys.exit(0)

    node_cfg = node_cfg.split("\n")

    # ----------------------------------------------------

    dbdir = None
    for line in node_cfg:
        m = re.match('\s*{dbdir,\s*"(.+)"}', line)
        if m and m.group(1):
            dbdir = m.group(1)
            break

    if not dbdir:
        sys.exit("ERROR: no previously configured dbdir")
    if not os.path.isdir(dbdir):
        sys.exit("ERROR: dbdir is not a directory: " + dbdir)

    if dbdir_file:
       dbdir_fp = open(dbdir_file, 'w')
       print >>dbdir_fp, dbdir
       dbdir_fp.close()
       print("Db directory: " + dbdir)
       sys.exit(0)

    # ----------------------------------------------------

    buckets = subprocess.Popen([escript_cmd,
                                os.path.join(bin_dir, "mbdumpconfig.escript"),
                                prev_cfg, "buckets", "membase"],
                               stdout=subprocess.PIPE).communicate()[0]
    buckets = string.strip(buckets)
    if len(buckets) <= 0:
        buckets = []
    else:
        buckets = buckets.split("\n")

    buckets_lcase = {}
    buckets_total = 0

    for bucket in buckets:
        bucket_path = os.path.join(dbdir, bucket)
        if not (os.path.isfile(bucket_path) or
                os.path.isdir(bucket_path + "-data")):
            sys.exit("ERROR: bucket " + bucket
                     + " is configured but missing: " + bucket_path)

        if bucket.lower() in buckets_lcase:
            sys.exit("ERROR: bucket " + bucket \
                     + " has a case-insensitive match with another bucket name." \
                     + " Please first rename the bucket before upgrading.")
        buckets_lcase[bucket.lower()] = True

        for path in bucket_files(dbdir, bucket):
            if os.path.isfile(path):
                buckets_total = buckets_total + os.path.getsize(path)

    # ----------------------------------------------------

    # TODO: Ensure Membase isn't running.
    # TODO: Can I write to the right directories?
    # TODO: linux: Am I running as the right user?
    # TODO: linux: Can I sudo to the right user?
    # TODO: linux: chmod of created files and directories.
    # TODO: Rollback files if there was an error.

    # ----------------------------------------------------

    confirm(install_dir, prev_cfg_dir, prev_cfg, prev_ip,
            dbdir, buckets)

    # ----------------------------------------------------

    print("\nChecking disk space available for buckets in directory:\n  " + dbdir)

    want = buckets_total * space_needed_factor
    avail = get_free_space(dbdir)
    if avail < 0:
        sys.exit("ERROR: unable to retrieve amount of free disk space")

    print("  Free disk bucket space wanted: " + str(want))
    print("  Free disk bucket space available: " + str(avail))
    print("  Free disk space factor: " + str(space_needed_factor))

    if avail < want:
        sys.exit("ERROR: not enough free disk space" \
                 + " in " + dbdir + " directory." \
                 + " Consider using the -s flag to supply" \
                 + " a different free disk space needed factor.")

    print("  Ok.")

    # ----------------------------------------------------

    print("\nAnalysis complete.")

    if dry_run:
        prompt("Proceed with config & data upgrade steps? [yes|no]")
    else:
        prompt("Proceed with config & data upgrade steps (writing new files)? [yes|no]")

    print("")

    upgrade(bin_dir, install_dir, prev_cfg_dir, prev_cfg, prev_ip,
            dbdir, buckets, dry_run=dry_run)

    print("\nDone.")

if __name__ == '__main__':
    main()
