#!/usr/bin/env python

import os
import shutil
import sys
from argparse import ArgumentParser

import yaml

from fireworks.core.firework import Firework
from fireworks.core.launchpad import LaunchPad
from fireworks.user_objects.firetasks.fileio_tasks import FileWriteTask
from fireworks.user_objects.firetasks.script_task import ScriptTask

__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"


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)))
    answer = input("Confirm? (Y/N)")
    if answer.startswith("Y"):
        for name, fwid, d, files in to_delete:
            try:
                shutil.rmtree(d)
                print(f"Deleted {d}")
            except Exception as ex:
                print(f"Unable to delete {d} because of {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)
