import argparse
import os
import stat

def print_cmd(ofile, url, cadence_str, k2id, campaign, cmdtype):
    # Print a data download command to an opened file given a K2 ID and campaign.
    this_filename = "ktwo" + k2id + '-c' + '{0:02d}'.format(int(campaign)) + cadence_str
    if cmdtype == "curl":
        cmd_str_to_add = " -f -R -o " + this_filename + " "
    else:
        cmd_str_to_add = " "
    this_cmd = (cmdtype + cmd_str_to_add + url + 'c'+'{0:d}'.format(int(campaign)) + '/' +
                k2id[0:4] + '00000/' + k2id[4:6] + '000/' + this_filename)
    ofile.write(this_cmd + '\n')
    
def get_k2(idict):
    # Extract command-line arguments from the input dict.
    k2ids = idict["k2ids"]
    cadence = idict["cadence"]
    data_type = idict["data_type"]
    ofile = idict["ofile"]
    cmdtype = idict["cmdtype"]
    if "campaigns" in idict:
        campaigns = idict["campaigns"]
    else:
        campaigns = None

    # This is the base URL to use when requesting K2 files.
    lc_base_url = "http://archive.stsci.edu/missions/k2/lightcurves/"
    tp_base_url = "http://archive.stsci.edu/missions/k2/target_pixel_files/"
    if data_type == "lightcurve":
        base_url = lc_base_url
    else:
        base_url = tp_base_url

    # Define the file extensions to use.
    if data_type == "lightcurve":
        if cadence == "long":
            cadence_str = "_llc.fits"
        else:
            cadence_str = "_slc.fits"
    else:
        if cadence == "long":
            cadence_str = "_lpd-targ.fits.gz"
        else:
            cadence_str = "_spd-targ.fits.gz"
    
    # Use all available campaigns if 'campaigns' are not specified.
    if campaigns is None:
        campaigns = [str(x) for x in xrange(26)]
    
    # If 'k2ids' is a scalar string put it in a list.  Same with 'epochs' and 'campaigns'
    if isinstance(k2ids,str):
        k2ids = [k2ids]
    if isinstance(campaigns, str):
        campaigns = [campaigns]

    # Make sure K2 IDs are zero-padded.
    k2ids = ["{0:09d}".format(int(x)) for x in k2ids]

    # Write to script file.
    with open(ofile, 'w') as output_file:
        output_file.write("#!/bin/sh\n")
        for k2id in k2ids:
            for campaign in campaigns:
                print_cmd(output_file, base_url, cadence_str, k2id, campaign, cmdtype)
        # Print a warning message to let users know there will sometimes be Error 404's.
        output_file.write("echo\n")
        output_file.write('echo Script completed.  Note: Some \\"Error 404\\" messages are expected,'
                          ' depending on your parameters.  It is always recommended to confirm the'
                          ' expected number of files are successfully downloaded.\n')

    # Make sure file is executable.
    os.chmod(ofile, 0744)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Script to generate download commands given one or"
                                     " more K2 IDs.", epilog="Ex.: python get_k2.py"
                                     " 202717132 204313397 -t target_pixel_file"
                                     " --campaigns 2")
    parser.add_argument("k2ids", action="store", nargs='+', help="One or more K2 IDs to"
                        " retrieve data from.")
    parser.add_argument("-c", action="store", type=str.lower, dest="cadence", default="long",
                        choices=["short", "long"], help="Specify the type of cadence.  Default="
                        "'%(default)s'.")
    parser.add_argument("-t", action="store", type=str.lower, dest="data_type", default=
                        "lightcurve", choices=["lightcurve","target_pixel_file"], help="What type"
                        " of data product to retrieve.  Default = '%(default)s'.")
    parser.add_argument("-o", action="store", dest="ofile", default="get_k2.sh", help="Full"
                        " path and file name of the output file to create.  Any existing file with"
                        " the same name will be overwritten.  Default = '%(default)s'.")
    parser.add_argument("--campaigns", action="store", type=str, dest="campaigns", nargs="*",
                        help="One or more campaign numbers to retrieve with each K2"
                        " ID.")
    parser.add_argument("--cmdtype", action="store", type=str.lower, dest="cmdtype", default="curl",
                        choices=["curl","wget"], help="What download tool should be used?  Default"
                        " = '%(default)s'.")
    args = vars(parser.parse_args())
    get_k2(args)