#!/usr/bin/env python

__author__ = 'Shyue Ping Ong'
__copyright__ = 'Copyright 2013, The Materials Project'
__version__ = '0.1'
__maintainer__ = 'Shyue Ping Ong'
__email__ = 'ongsp@ucsd.edu'
__date__ = '1/6/14'

import os
import shutil
import sys
from argparse import ArgumentParser
from fireworks.core.firework import FireWork
from fireworks.user_objects.firetasks.fileio_tasks import FileWriteTask
from fireworks.user_objects.firetasks.script_task import ScriptTask
from fireworks.core.launchpad import LaunchPad
import yaml


def create_fw_single(args, fnames, yaml_fname):
    tasks = []
    if fnames:
        files = []
        for fname in fnames:
            with open(fname) as f:
                files.append(
                    {
                        "filename": os.path.basename(fname),
                        "contents": f.read()
                    }
                )
        tasks.append(FileWriteTask({"files_to_write": files}))
    if args.command is not None:
        tasks.append(ScriptTask({"script": [args.command]}))
    fw = FireWork(tasks, name=args.name)
    with open(yaml_fname, "w") as f:
        yaml.dump(fw.to_dict(), f, default_flow_style=False)


def create_fw(args):
    if not args.directory_mode:
        create_fw_single(args, args.files_or_dirs, args.output)
    else:
        output_fname = args.output.replace(".yaml", "{}.yaml")
        for i, d in enumerate(args.files_or_dirs):
            create_fw_single(
                args,
                [os.path.join(d, f) for f in os.listdir(d)],
                output_fname.format(i)
            )

def do_cleanup(args):
    lp = LaunchPad.from_file(
        args.launchpad_file) if args.launchpad_file else LaunchPad(
        strm_lvl=args.loglvl)
    to_delete = []
    for i in lp.get_fw_ids({}):
        fw = lp.get_fw_by_id(i)
        if fw.state == "COMPLETED":
            for l in fw.launches:
                if os.path.isdir(l.launch_dir):
                    to_delete.append((fw.fw_name, fw.fw_id,
                                      l.launch_dir, os.listdir(l.launch_dir)))
    if len(to_delete) == 0:
        print("Your directories are clean!")
        sys.exit(0)
    print("The following will be deleted:")
    for name, fwid, d, files in to_delete:
        print("{}-{} - {} - {}".format(name, fwid, d, ", ".join(files)))
    ans = raw_input("Confirm? (Y/N)")
    if ans.startswith("Y"):
        for name, fwid, d, files in to_delete:
            try:
                shutil.rmtree(d)
                print("Deleted {}".format(d))
            except Exception as ex:
                print("Unable to delete {} because of {}".format(d, ex))
    else:
        print("Canceled!")


if __name__ == '__main__':
    p = ArgumentParser(description="""This script is used to create simple
    fireworks workflows from input files and script commands.
    """)
    sp = p.add_subparsers(help='command', dest='command')

    create_parser = sp.add_parser("create",
        help='Create a new simple workflow. A simple workflow involves '
             'writing a set of input files and/or running a series of '
             'commands.')

    create_parser.add_argument(
        '-c', '--command', default=[], nargs="*",
        type=str, help="Add a ScriptTask at the end to run a series of "
                       "commands.")
    create_parser.add_argument(
        '-o', '--output', default="my_fw.yaml",
        type=str,
        help="Output file name to write to. Defaults to \"my_fw.yaml\".")
    create_parser.add_argument(
        '-d', '--directory_mode', action="store_true",
        help="If specified, files_or_dirs is treated as a series of "
             "input directories and multiple fireworks will be generated with "
             "each directory treated as a single job. A _# will be appended "
             "to the yaml file name.")
    create_parser.add_argument(
        '-n', '--name', default="MyFW",
        type=str, help="Add a name for your firework. Defaults to \"MyFW\".")
    create_parser.add_argument(
        '-f', '--files_or_dirs', default=[], nargs="*",
        type=str, help="Files that need to be added as input files. Supports "
                       "wild cards.")

    create_parser.set_defaults(func=create_fw)

    cleanup_parser = sp.add_parser('cleanup', help='Cleanup the queue.')
    cleanup_parser.set_defaults(func=do_cleanup)

    args = p.parse_args()
    args.func(args)
