#!/usr/bin/env python
import sys
import os
import tarfile
import shutil
from ftplib import FTP
import optparse
import string

def detect_domain():
    choices = ["atnf", "astron", "nrao"]
    domain = socket.gethostname().split(".")[-2]
    if domain in choices:
        return domain
    else:
        return choices[0]

usage = """usage: %prog [options] <path_to_measures_data (optional)>a

If no directory is given it will attempt to look in the default locations

   a) measures.directory in rc file given in CASARCFILES environment variable
   b) measures.directory in ~/.casarc
   c) /usr/local/share/casacore/data
"""
parser = optparse.OptionParser(usage, description="""This script updates the casacore measures data directory if it is outdated.
""")

parser.add_option('-l', '--location', dest='location', type="choice",
                  choices=["atnf", "astron", "nrao"],
                  default=detect_domain(),
                  help="Select TARGET from: %s [default=%%default]" % \
                  string.join(["atnf", "astron", "nrao"], ", "))
parser.add_option('-v', '--verbose', dest='verbose',
                  action="store_true",
                  help='show helpful status')
parser.add_option('-c', '--create', dest='create',
                  action="store_true",
                  help='Create the measures data directory')

# parse command line options
(OPTS, args) = parser.parse_args()

def dir_from_casarc(casarcfile):
    if not os.path.exists(casarcfile):
        return None
    f = file(casarcfile, "r")
    for l in f:
        l = l.strip()
        if l.startswith("#"): continue
        k,v = l.split()
        k = k.strip()
        if k.startswith("measures.directory"):
            return v.strip()
        return None

def is_measures_dir(dir):
    return os.path.exists(dir) and \
        os.path.exists(os.path.join(dir, "geodetic")) and \
        os.path.exists(os.path.join(dir, "ephemerides"))
                                        
def get_measures_dir():
    rcfiles = [ i for i in os.environ.get("CASARCFILES", ":").split(":") if i]
    rcfiles.append(os.path.join(os.path.expanduser("~"), ".casarc"))
    rcfiles.append(".")
    rcfiles.append("/usr/local/share/casacore/data/")
    for rc in rcfiles:
        mdir = dir_from_casarc(rc)
        if mdir:
            return mdir
    return None

def vprint(msg):
    if OPTS.verbose:
        print msg

# try to find out where the measures data directories are located
if len(args) == 1:
    if os.path.exists(args[0]):
        measdir = args[0]
    else:
        print "Measures data directory doesn't exist. Please create  it first"
        sys.exit(1)
else:
    measdir = get_measures_dir()

vprint("Using measures data directory '%s'" % measdir)


# globals
md5suff = '.md5sum'
name = 'measures_data.tar.bz2'
dataservers = { 'ATNF': "ftp.atnf.csiro.au",
                'ASTRON': 'ftp.astron.nl',
                'NRAO': 'ftp.nrao.edu' }
datadirs = {"ATNF": 'pub/software/casacore/measures_data',
            "ASTRON": 'pub/software/casacore/measures_data',
            "NRAO": 'pub/software/casacore/measures_data' }
tmpdata = '/tmp/'+ name
tmpmd5 = tmpdata+md5suff

vprint("Connecting to server %s" % dataservers[OPTS.location])
ftp = FTP(dataservers[OPTS.location])   # connect to host, default port
ftp.login()               # user anonymous, passwd anonymous@
ftp.cwd(datadir[OPTS.location])

vprint("Checking if an update is required...")

ftp.retrbinary('RETR %s' % name+md5suff, open(tmpmd5, 'wb').write)
md5file = file(tmpmd5)
md5new = md5file.readlines()[0].split()[0]

try:
    fl = os.path.join(measdir, name+md5suff)
    data_md5 = file(fl)
    ls = data_md5.readlines()[0]
    data_md5.close()
    md5old = ls.split()[0]
except IOError:
    md5old = None

#print md5new, md5old
if md5new != md5old:
    vprint("Update required. Downloading measures data archive...")
    ftp.retrbinary('RETR %s' % name, open(tmpdata, 'wb').write)
    os.chdir(measdir)
    tf = tarfile.TarFile.bz2open(tmpdata)
    vprint("Extracting data archive in %s..." % measdir)
    for member in tf.getmembers():
        tf.extract(member)
    shutil.copy(tmpmd5, measdir)
else:
    vprint("""Data already at latest available version.
If you still get errors running measures, please report this.""")
ftp.quit()
