#!/usr/bin/env python

import os
import glob
import warnings
import platform

# name generation
import datetime
import random

import encap_lib.encap_settings as settings
from encap_lib.machines import LocalMachine
from encap_lib.encap_lib import get_machine, filename_and_file_extension, get_interpreter_from_file_extension, extract_folder_name, get_job_name
from encap_lib.encap_lib import record_process, remove_process_from_database
from encap_lib import slurm
import encap_lib.pueue
from encap_lib.status import get_status, print_status
import encap_lib.git_tracker as git_tracker

def tail_pull(machine, run_folder_name, pid=None, log_file_name="log"):
    # `pid` can be:
    #   - None                       no tracked process, fall back to EOT loop
    #   - "<digits>"                 OS PID on the remote (non-slurm run)
    #   - "slurm:<jobid>"            slurm job id (set by run_slurm so encap
    #                                tail can poll squeue)
    is_slurm = isinstance(pid, str) and pid.startswith("slurm:")
    is_os_pid = isinstance(pid, str) and pid.isdigit()

    if is_slurm:
        # Stream the log AND poll squeue. Exits on either the EOT marker
        # (normal completion) or when the slurm job leaves the queue (catches
        # scancel, walltime cap, OOM kill, etc.). `-n +1` so we dump the full
        # log on attach, matching the --pid branch below.
        jobid = pid.split(":", 1)[1]
        machine.run_code(f"""#!/bin/bash
touch {run_folder_name}/{log_file_name}
exec 3< <(tail -f -n +1 {run_folder_name}/{log_file_name})
TAIL_PID=$!
trap 'kill $TAIL_PID 2>/dev/null; wait $TAIL_PID 2>/dev/null' EXIT

while true; do
    # Drain currently-available log lines (1s timeout per blocked read).
    while IFS= read -r -t 1 -u 3 LOGLINE; do
        [[ "$LOGLINE" == "{chr(4)}" ]] && exit 0
        printf '%s\\n' "$LOGLINE"
    done
    # No log activity for >=1s -- check whether the slurm job is still queued/running.
    if ! squeue -h -j {jobid} 2>/dev/null | grep -q .; then
        # Job is gone. Final drain to catch any last buffered lines, then exit.
        while IFS= read -r -t 1 -u 3 LOGLINE; do
            [[ "$LOGLINE" == "{chr(4)}" ]] && exit 0
            printf '%s\\n' "$LOGLINE"
        done
        exit 0
    fi
done
""", verbose=True, output=True)
        machine.run_code(f"rm {run_folder_name}/pid -f")
    elif is_os_pid and platform.system() != "Darwin":
        machine.run_code(f"touch {run_folder_name}/{log_file_name}; tail -f -n +1 --pid={pid} -f {run_folder_name}/{log_file_name}", verbose=True, output=True)
        machine.run_code(f"rm {run_folder_name}/pid -f")
    else:
        # Runs the tail command until it sees the special character.
        # `tail -f` is wired through an explicit FD so its PID is captured in
        # $! and killed on EXIT -- otherwise, when the loop exits on EOT,
        # tail keeps sitting in its inotify/kqueue wait and only learns the
        # reader is gone the next time it tries to write. If the script emits
        # nothing after EOT (no atexit / no stderr flush), tail hangs forever,
        # ssh keeps the session up, and encap appears frozen until Ctrl-C.
        # `-n +1` so already-finished logs are dumped fully on attach.
        machine.run_code(f"""#!/bin/bash
touch {run_folder_name}/{log_file_name}
exec 3< <(tail -f -n +1 {run_folder_name}/{log_file_name})
TAIL_PID=$!
trap 'kill $TAIL_PID 2>/dev/null; wait $TAIL_PID 2>/dev/null' EXIT
while IFS= read -r -u 3 LOGLINE || [[ -n "$LOGLINE" ]]; do
    [[ "${{LOGLINE}}" == "{chr(4)}" ]] && exit 0
    printf '%s\\n' "$LOGLINE"
done
""", verbose=True, output=True)

    machine.pull(run_folder_name, run_folder_name, directory=True)


ADJECTIVES = [
    "agile", "alert", "ancient", "ardent", "artful", "astonishing", "bold", "brave", "brilliant",
    "calm", "careful", "cautious", "charming", "clever", "colossal", "courageous", "crafty", 
    "cunning", "daring", "dauntless", "dazzling", "devoted", "eager", "ebullient", "eccentric", 
    "elusive", "energetic", "enlightened", "epic", "fearless", "fiery", "fierce", "flashy", 
    "flawless", "focused", "formidable", "gallant", "gifted", "glorious", "graceful", "grand", 
    "gregarious", "gritty", "heroic", "honest", "intrepid", "invincible", "jolly", "keen", 
    "lively", "luminous", "magnificent", "majestic", "marvelous", "mighty", "mischievous", 
    "mysterious", "nimble", "observant", "omniscient", "outlandish", "outstanding", "passionate", 
    "patient", "perceptive", "persistent", "playful", "plucky", "proud", "quick", "radiant", 
    "reliable", "resilient", "resolute", "resourceful", "rugged", "savvy", "shrewd", "silent", 
    "sleek", "sly", "spirited", "steadfast", "stubborn", "swift", "tenacious", "thoughtful", 
    "undaunted", "unpredictable", "unyielding", "valiant", "vibrant", "vigorous", "vivid", 
    "wary", "wily", "wise", "witty", "zealous"
]
ANIMALS = [
    "aardvark", "albatross", "anaconda", "antelope", "armadillo", "axolotl", "badger", "basilisk", 
    "bat", "beetle", "bison", "blackbird", "bobcat", "buffalo", "butterfly", "caracal", "cassowary", 
    "catfish", "centipede", "cheetah", "cormorant", "cougar", "coyote", "crane", "crocodile", "crow", 
    "dingo", "dolphin", "dragonfly", "eagle", "echidna", "egret", "elephant", "emu", "falcon", 
    "ferret", "firefly", "fox", "frog", "gazelle", "gecko", "giraffe", "glowworm", "gorilla", 
    "grizzly", "hammerhead", "hare", "harrier", "hawk", "hedgehog", "heron", "hornet", "hyena", 
    "ibex", "ibis", "iguana", "jackal", "jaguar", "jellyfish", "kangaroo", "kingfisher", "koala", 
    "kraken", "lemur", "leopard", "lion", "lynx", "mandrill", "meerkat", "mockingbird", "mongoose", 
    "moose", "moth", "narwhal", "ocelot", "octopus", "opossum", "oryx", "osprey", "otter", 
    "owl", "panda", "panther", "parrot", "peacock", "pelican", "penguin", "phoenix", "piranha", 
    "platypus", "porcupine", "puma", "quail", "quokka", "raccoon", "ram", "raven", "red panda", 
    "reindeer", "roadrunner", "salamander", "scorpion", "seahorse", "serval", "shark", "skunk", 
    "sloth", "snow leopard", "squid", "stag", "stingray", "swan", "tarantula", "tiger", "tortoise", 
    "toucan", "viper", "vulture", "wallaby", "warthog", "weasel", "whale", "wildebeest", "wolf", 
    "wolverine", "wombat", "woodpecker", "yak", "zebra"
]

def generate_experiment_name():
    timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
    adjective = random.choice(ADJECTIVES)
    animal = random.choice(ANIMALS)
    return f"{timestamp}_{adjective}_{animal}"


def run(machine, file_extension, target, args, run_folder_name, target_file, target_file_path, interprter_args="", number_of_instances=1, name=None):
    """
    Run the file and save outputs in log, save PID in the file pid.
    """
    # The right file is already on the remote by the time we get here:
    #   - folder-mode fresh-run: source pushed via rsync_push in mode_run_folder
    #   - file-mode fresh-run:   source pushed via mode_run_file's push
    #   - rerun:                 pushed via the `if rerun:` rsync_push in mode_rerun_file
    # An additional self-push here would, for -vm runs where the local
    # 0encap_folder copy is a stale pull-back from a previous run, overwrite
    # the freshly-pushed remote file with that stale copy, silently
    # discarding the user's source edits.

    # Check if the files is executable, if yes then run it directly
    is_executable = os.access(target_file_path, os.X_OK)

    # Get interpreter from file extension
    interpreter = get_interpreter_from_file_extension(file_extension, ignore_file_extensior_if_interpreter_set_in_settings=True, error_if_not_found=not is_executable)

    # Give warning if an interpreter is set in the settings file
    if is_executable and "interpreter" in settings.config:
        warnings.warn(f"WARNING: 'interpreter' is set in the settings files, but your file is executable. The interpreter in the (settings file)/(command line) will be ignored.", UserWarning)
    
    if is_executable and interpreter != "":
        warnings.warn(f"WARNING: Your file is executable, but the interpreter {interpreter} could be infered. The file will be executed, make sure that this is what you want!", UserWarning)

    if isinstance(number_of_instances, int):
        iterator = range(number_of_instances)
    else:
        iterator = number_of_instances

    log_file_names = []
    pids = []
    for i in iterator:

        if i == 0:
            log = "log"
            instance_text = ""
        else:
            log = f"log_{i}"
            instance_text = f"_{i}"
        
        args_replace = args.replace("{i}", f"{i}")
        
        # Define job name
        job_name = get_job_name(target, name, instance_text, args_replace)

        print_instance = ""
        if len(iterator) > 1:
            print_instance = f"echo 'Instance {i}' >> {log} 2>&1"
        

        setsid = "setsid"
        if platform.system() == "Darwin":
            setsid = ''
        if is_executable:
            run_the_experiment = f"""{setsid} nohup bash -c "time (./{target_file} {args_replace}) 2>&1 | tee -a /dev/null; printf '\\004\\n'" >> {log} 2>&1 &"""
        else:
            run_the_experiment = f"""{setsid} nohup bash -c "time ({interpreter} {interprter_args} {target_file} {args_replace}) 2>&1 | tee -a /dev/null; printf '\\004\\n'" >> {log} 2>&1 &"""
        
        code = f'''
        set -e # Exit on error
        
        cd {run_folder_name}
        export ENCAP_PROCID={i}
        export ENCAP_NAME="{run_folder_name}"
        export ENCAP_JOB_NAME="{job_name}"
        export ENCAP_LOG="$PWD/{log}"
        date &> {log}
        echo "host: $(hostname)" >> {log} 2>&1
        {print_instance}
        echo "{target_file_path} {args}  \n" >> {log} 2>&1
        # Tee avoids buffering
        {run_the_experiment}
        PID=$!
        disown
        echo $PID 
        '''

        if i == 0:
            code += f"""
        
        echo "$PID" > pid"""
            pid = machine.run_code(code, output=True)[0]
            print('PID ' + pid)
        else:
            pid = machine.run_code(code, output=True)[0]
        
        pids.append(pid)
        log_file_names.append(log)

    return pids[0], log_file_names[0]

def run_slurm(machine, file_extension, target, args, run_folder_name, target_file, target_file_path, slurm_settings, interpreter_args="", name=None):
    # (Redundant self-push removed -- see comment in run() above.)

    # Delete previous log file and create a new one
    machine.run_code(f"""rm {run_folder_name}/pid -f
    rm {run_folder_name}/.slurm_files/log* -f
    touch {run_folder_name}/log
    mkdir -p {run_folder_name}/.slurm_files
    """)
    if not ("i" in slurm_settings):
        slurm_settings["i"] = 1
    
    if not ("ntasks-per-node" in slurm_settings):
        slurm_settings["ntasks-per-node"] = 1

    if isinstance(slurm_settings["i"], int):
        iterator = range(slurm_settings["i"])
    else:
        iterator = slurm_settings["i"]
    
    log_file_name = ""
    # Capture the first instance's slurm job id and persist it as
    # `slurm:<id>` in the pid file so `encap tail` can poll squeue and detect
    # both normal completion (EOT marker in log) and abnormal terminations
    # (scancel, walltime cap, OOM kill -- where no EOT is ever written).
    first_slurm_jobid = None
    # Create the slurm script for all slurm instances
    for i in iterator:
        if i == 0:
            slurm_instance_text = ""
        else:
            slurm_instance_text = f"_{i}"
        
        if log_file_name == "":
            log_file_name = f"log{slurm_instance_text}"
            # Delete previous log file and create a new one if it is the first instance
            machine.run_code(f"""rm {run_folder_name}/{log_file_name} -f
                                 touch {run_folder_name}/{log_file_name }""")
        
        runslurm_file_name = f"{run_folder_name}/.slurm_files/run{slurm_instance_text}.slurm"
        executable_file_name = f"{run_folder_name}/.slurm_files/run{slurm_instance_text}.sh"
        log_file_name_slurm = f"{run_folder_name}/.slurm_files/log{slurm_instance_text}.slurm"
        # We need args_replace which is just args replacing {i} with slurm_instance
        args_replace = args.replace("{i}", f"{i}")
        
        # Define job name
        job_name = get_job_name(target, name, slurm_instance_text, args_replace)

        # Script to run the file and save outputs in log
        code, _ = slurm.generate_slurm_executable(file_extension, run_folder_name, target_file_path, args, target_file=target_file, slurm_instance=i, ntpn=slurm_settings["ntasks-per-node"], job_name=job_name)
        
        # Save the script in the run_folder
        machine.write_file(executable_file_name, code, verbose=True)
        
        # Generate the slurm settings script
        code = slurm.generate_code_for_slurm_script(run_folder_name, slurm_settings, runslurm_file_name, executable_file_name, log_file_name_slurm, job_name=job_name)

        # Write the slurm file
        machine.write_file(runslurm_file_name, code, verbose=True)

        # Run the slurm file
        out = machine.run_code(f"sbatch {runslurm_file_name}", verbose=True)
        if "sbatch: command not found" in out[0]:
            print("Error: Slurm does not seem to be installed on this machine.")
            exit(1)

        # Extract the slurm job id from sbatch's output ("Submitted batch job
        # 12345") and persist the first instance's id as "slurm:<id>" in the
        # pid file. tail_pull uses this to poll squeue.
        if i == 0:
            for line in out:
                parts = line.strip().split()
                if parts and parts[-1].isdigit():
                    first_slurm_jobid = parts[-1]
                    break
            if first_slurm_jobid is not None:
                machine.run_code(
                    f"echo 'slurm:{first_slurm_jobid}' > {run_folder_name}/pid"
                )

    return log_file_name, first_slurm_jobid

def get_script_name(run_folder_name, script_name):
    
    # Search in the run_folder for the script_name with * as a wildcard
    script_name_ = os.path.join(run_folder_name, script_name)
    script_names = glob.glob(script_name_)
    assert len(script_names) == 1, f"Found {len(script_names)} scripts matching {script_name} in {run_folder_name}."
    script_name_ = script_names[0]

    # Remove the run_folder_name from the script_name
    script_name = os.path.basename(script_name_)
    
    return script_name

def create_folder_and_check_if_experiment_exists(machine, pargs, folder_name, run_folder_name):
    """
    Creates folder if it does not exist and checks if the caspsule already exists.
    """
    code = f'''
    mkdir -p {folder_name}
    if [ -d "{run_folder_name}" ]
    then
        echo "exists"
    else
        mkdir -p {run_folder_name}
        echo "ok"
    fi
    '''
    out = machine.run_code(code)
    assert len(out) == 1, str(out)
    out = out[0]

    if out == "ok":
        pass

    elif out == "exists":
        if not pargs.yes:
            c = input(f"The experiment {run_folder_name} already exists. This action will overwrite the {run_folder_name} folder. Do you whish to continue y/n? ")
            if c == "y" or c == "Y":
                pass
            else:
                quit()
    else:
        raise Exception(f"Unexpected value {out}.")
    
def mode_run_file(folder_name, run_folder_name, source_file_path, local_project_dir, target_file_path, machine, pargs, interpreter_args=""):
    # If no name is provided, generate one
    experiment_name = pargs.name if pargs.name else generate_experiment_name()

    #Syncs folders
    machine.sync_files()

    # Creates folder if it does not exist and checks if the caspsule already exists.
    create_folder_and_check_if_experiment_exists(machine, pargs, folder_name, run_folder_name)

    # Copy the local copy of pargs.target to the run_folder
    machine.push(source_file_path, target_file_path, directory=False, verbose=True)

    # Run the file and save outputs in log, save PID in the file pid.
    mode_rerun_file(run_folder_name=run_folder_name, local_project_dir=local_project_dir, machine=machine, pargs=pargs, rerun=False)

def mode_run_folder(folder_name, run_folder_name, source_file_path, local_project_dir, machine, pargs, interpreter_args=""):
    assert pargs.name is not None, "The experiment_name -n must be specified."

    # Syncs folders
    machine.sync_files()

    # Creates folder if it does not exist and checks if the caspsule already exists.
    create_folder_and_check_if_experiment_exists(machine, pargs, folder_name, run_folder_name)

    # create_folder_and_check_if_experiment_exists only mkdirs on the remote
    # for SSH machines. Mirror it locally so subsequent local-filesystem
    # operations (config writes, isdir asserts) succeed in -vm mode.
    os.makedirs(run_folder_name, exist_ok=True)

    # Copy the local version of pargs.target to the run_folder.
    # Use rsync_push so that (a) rsync_exclude_push is honored, (b) only
    # changed files cross the network on subsequent runs. Trailing "/" on
    # source_file_path mirrors the previous `copy_full_dir=False` semantics
    # (copy contents, not the source folder itself).
    # -fp/--force-push: drop --update and add --ignore-times to transfer
    # regardless of mtime/size comparison.
    machine.rsync_push(source_file_path + "/", run_folder_name + "/",
                       directory=True, verbose=True,
                       last_timestamp_prevails=not pargs.force_push,
                       ignore_times=pargs.force_push,
                       ignore=["sending incremental", "sent ", "total size"])

    # Patch settings with all .encap.conf files. Walk from the *source* folder:
    # in -vm mode the run_folder only exists on the remote, so walking from
    # the local mirror would miss any topic-level .encap.conf entirely. The
    # source folder has the same configs (it's what was just pushed) and is
    # guaranteed to exist locally.
    settings.load_encap_config_files_recursive(os.path.join(os.getcwd(), source_file_path))

    # Infer the proper script_name
    if "script_name" in settings.config:
        script_name = settings.config["script_name"]
    else:
        # Defaul script name is run.*
        script_name = "run.*"

    # Glob against the source folder for the same reason as the config walk
    # above: the run folder may only exist on the remote in -vm mode.
    script_name = get_script_name(source_file_path, script_name)

    settings.config["script_name"] = script_name
    settings.args_config["script_name"] = script_name

    # Run the file and save outputs in log, save PID in the file pid.
    mode_rerun_file(run_folder_name=run_folder_name,
                    local_project_dir=local_project_dir, machine=machine,
                    pargs=pargs, rerun=False)
    

def mode_rerun_file(run_folder_name, local_project_dir, machine, pargs, rerun=True, interpreter_args=""):
    assert pargs.name is not None, "The experiment_name -n must be specified."

    # Check if the experiment in run_folder_name exists
    assert os.path.isdir(run_folder_name), f"The experiment {run_folder_name} does not exist."


    # Patch settings with all .encap.conf files
    settings.load_encap_config_files_recursive(os.path.join(os.getcwd(), run_folder_name))

    # If any git-tracking is enabled, then we need to save in the config
    if "git-track" in settings.config:
        settings.config["git-track"] = git_tracker.get_current_commit_hashes(settings.config["git-track"])
    
    if "git-track-force" in settings.config:
        settings.config["git-track-force"] = git_tracker.get_current_commit_hashes(settings.config["git-track-force"], force=True, verbose=settings.debug,
                                                                                   commit_message=f"{run_folder_name} automatic commit from encap.")
        
    # Save the settings
    settings.write_config_file(os.path.join(run_folder_name, ".encap_history.conf"), settings.config,
                               comment="This file is automatically generated by encap. Effective encap configuration file from last run that has no effect on future runs.")
    
    pid = None

    # Read settings
    if "args" in settings.config:
        args = settings.config["args"]
    else:
        args = ""
    
    if "script_name" in settings.config:
        target_file = settings.config["script_name"]
    else:
        # Infer the proper script_name, the default is run.* (This can only happen in folder mode)
        target_file = get_script_name(run_folder_name, "run.*")
        settings.config["script_name"] = target_file
        settings.args_config["script_name"] = target_file
    
    target_file_path = os.path.join(run_folder_name, target_file)
    _, file_extension = filename_and_file_extension(target_file_path)

    
    

    # If .encap.conf file exists in the run_folder, load it, merge it with the comannd line arguments and save it
    if os.path.exists(os.path.join(run_folder_name, ".encap.conf")):
        run_folder_config = settings.read_config_file(os.path.join(run_folder_name, ".encap.conf"))

        run_folder_config = settings.merge_dicts(run_folder_config, settings.args_config)
    else:
        run_folder_config = settings.args_config
    
    settings.write_config_file(os.path.join(run_folder_name, ".encap.conf"), run_folder_config,
                               comment="This file is automatically generated by encap. This file will patch the global configuration file if the experiment is rerun.")

    if rerun:
        machine.sync_files()
        # rsync the whole experiment dir up, not just the entry script. The
        # previous single-file push (with a `# TODO: The entire folder should
        # be pushed, not just the file.` comment) meant that edits to helper
        # modules (e.g. files imported by the entry script) inside the local
        # experiment dir never reached the remote on `encap rerun`. rsync
        # honors rsync_exclude_push, so outputs (*.hdf5, *.pdf, etc.) aren't
        # shipped back up; stale log files might get pushed but get truncated
        # by encap's wrapper at job start, so no harm done.
        machine.rsync_push(run_folder_name + "/", run_folder_name + "/",
                           directory=True, verbose=True,
                           last_timestamp_prevails=not pargs.force_push,
                           ignore_times=pargs.force_push,
                           ignore=["sending incremental", "sent ", "total size"])

    # Are we using slurm or pueue?
    if settings.using_pueue:
        pueue_settings = encap_lib.pueue.read_pueue_settings_from_encapconfig(pargs.vm, local_project_dir)
        log_file_name = encap_lib.pueue.run_pueue(machine, file_extension, pargs.target, args, run_folder_name=run_folder_name, target_file=target_file, target_file_path=target_file_path, pueue_settings=pueue_settings, name=pargs.name)
        
    elif not settings.using_slurm:
        if "i" in settings.config:
            number_of_instances = settings.config["i"]
        else:
            number_of_instances = 1

        pid, log_file_name = run(machine, file_extension, pargs.target, args, run_folder_name=run_folder_name, target_file=target_file, target_file_path=target_file_path, number_of_instances=number_of_instances, name=pargs.name)

    else:

        # Load the slurm settings from the loaded config file
        slurm_settings = slurm.read_slurm_settings_from_encapconfig(pargs.vm, local_project_dir)

        # Run the SLURM job with the slurm config file. run_slurm now also
        # returns the first instance's slurm job id; encode as "slurm:<id>"
        # so tail_pull polls squeue rather than relying solely on the EOT
        # marker (which never arrives if the job is killed by slurm itself,
        # OOMs, or hits a walltime cap).
        log_file_name, slurm_jobid = run_slurm(machine, file_extension, pargs.target, args, run_folder_name=run_folder_name, target_file=target_file, target_file_path=target_file_path, slurm_settings=slurm_settings, name=pargs.name)
        if slurm_jobid is not None:
            pid = f"slurm:{slurm_jobid}"

    machine.pull(run_folder_name, run_folder_name, directory=True)

    # Record active process
    record_process(pargs.vm, pargs.target, pargs.name)
    tail_pull(machine, run_folder_name, pid, log_file_name = log_file_name)
    remove_process_from_database(pargs.target, pargs.name)


def parse_arguments():
    import argparse
    from argparse import ArgumentParser
    import textwrap

    parser = ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("mode", help=textwrap.dedent("""The mode of the program. Either run, rerun, tail, copy, status, kill, tar or untar.
run: Copies the target file into the experiment and runs it.
rerun: Reruns the target file in the experiment without copying.
tail: Tails the log file of the experiment.
copy: Copies the experiment specifeid with -cn into the experiment specified with -n.
status: Prints the status of all the experiments running on the machine.
kill: Kills the running experiment. Note that if the flag -i is specified, then only the instance/list of instances specified by -i is killed. Else all instances if -i is not specified.
tar: Tars the experiment, if -n is not specified, all experiments are tarred.
untar: Untars the experiment, if -n is not specified, all experiments are untarred.
                                                     
    """))

    parser.add_argument("target", nargs="?", help="Name of the target file to run.", default=None)
    parser.add_argument("-n", "--name", dest="name", help="Name of the experiment.", default=None)
    parser.add_argument("-cn", "--copy-name", dest="copy_name", help="Name of the experiment to copy from.", required=False, default=None)
    parser.add_argument("-a", "--args", dest="args", help="Arguments for the interpreter, for example --args ' -k 5', forming `<interpreter> <file_name> <args>`.", default=None)
    parser.add_argument("-vm", "--vm_name", dest="vm", help="Name of the Virtual Machine (VM).", default=None)
    parser.add_argument("-fp", "--force-push", dest="force_push",
                        action="store_true", default=False,
                        help="Force-push local files to the remote, ignoring rsync's mtime/size short-circuit (adds --ignore-times, drops --update). Use to recover from clock skew or other sync mysteries; rsync still uses delta-transfer so unchanged bytes don't go over the wire. Has no effect without -vm.")
    parser.add_argument("-sn", "--script_name", dest="script_name", help="Specify the script file to run when a folder is targeted. Default behavior is to find a file named `run.*`.", default=None)
    parser.add_argument("-i", "--number_of_instances", dest="i", help="Specify the number of instances to initiate concurrently. Each instance will have an associated environment variable, ENCAP_PROCID=<instance_number>. This option accepts either an integer, list or a Python expression in string format. When a list is provided, instances corresponding to the list items are initiated. Examples: -i '[0, 2, 3]' or -i 'range(5, 20)'. This feature is similar to -sl_i in Slurm mode.", default=None, type=eval)
    parser.add_argument("-int", "--interpreter", dest="interpreter", help="Specify the interpreter. If not provided, it's inferred from the file extension.", default=None)
    parser.add_argument("-tar_cpus", "--tar_cpus", dest="tar_cpus", help="Number of cpus to use when tarring several experiments. Default is 1.", default=1, type=int)

    # Slurm arguments
    parser.add_argument("-sl", "--slurm", action="store_true", dest="slurm", help="Flag to run the job on Slurm. If any Slurm option is used, this flag is automatically set to True.", default=False)
    parser.add_argument("-sl_nodes", "--slurm_nodes", dest="sl_nodes", help="Number of nodes to start in slurm. - nodes", default=None, type=int)
    parser.add_argument("-sl_ntpn", "--slurm_ntasks-per-node", dest="sl_ntpn", help="Number of tasks per node to start in slurm. - ntasks-per-node", default=None, type=int)
    parser.add_argument("-sl_time", "--slurm_time", dest="sl_time", help="Time to run the job in slurm. - time", default=None)
    parser.add_argument("-sl_partition", "--slurm_partition", dest="sl_partition", help="Partition to run the job in slurm.", default=None)
    parser.add_argument("-sl_account", "--slurm_account", dest="sl_account", help="Account to run the job in slurm.", default=None)
    parser.add_argument("-sl_cpus", "--slurm_cpus-per-task", dest="sl_cpus", help="Number of cpus to run the job in slurm. - cpus-per-task", default=None, type=int)
    parser.add_argument("-sl_nice", "--slurm_nice", dest="sl_nice", help="Nice value to run the job in slurm. Sets the priority, higer nice values equals lower priority. - nice", default=None, type=int)
    parser.add_argument("-sl_i", "--slurm_instances", dest="sl_i", help="Number of seperate slurm instances to start. If you use {i} in args it will be replaced by the index number of the slurm instance. For each instance the ENCAP_PROCID enviromental variable and the ENCAP_SLURM_INSTANCE variable will be set. If you pass a list it will run the instances in that list. Example: -sl_i [0, 2, 3] or range(5, 20)", default=None, type=eval)
    
    # Pueue arguments
    parser.add_argument("-pu", "--pueue", action="store_true", dest="pueue", help="Flag to run the job with Pueue. If any Pueue option is used, this flag is automatically set to True.", default=False)
    parser.add_argument("-pu_i", "--pueue_instances", dest="pu_i", help="Number of separate pueue instances to start. If you use {i} in args it will be replaced by the index number of the pueue instance. For each instance the ENCAP_PROCID and ENCAP_PUEUE_INSTANCE variables will be set. Example: -pu_i [0, 2, 3] or range(5, 20)", default=None, type=eval)
    parser.add_argument("-pu_pt", "--pueue_parallel_tasks", dest="pu_pt", help="Number of tasks that can run in parallel in the pueue group.", default=None, type=int)

    # Other arguments
    parser.add_argument("-y", "--yes",
                        action="store_true", dest="yes", default=False,
                        help="All prompts will be answerd with yes.")

    parser.add_argument("-d", "--debug",
                        action="store_true", dest="debug", default=False,
                        help="All commands will be printed.")

    parser.add_argument("-dr", "--dryrun",
                        action="store_true", dest="dryrun", default=False,
                        help="No command will be executed. Also implies debug.")

    pargs = parser.parse_args()

    # Read the data in
    if pargs.debug:
        settings.debug = True
    if pargs.dryrun:
        settings.dryrun = True
    
    if pargs.target is None:
        assert pargs.mode == "status", "No target specified."
    
    if pargs.name is None:
        assert pargs.mode == "status" or pargs.mode == "tar" or pargs.mode == "untar", "encap: error: the following arguments are required: -n/--name"
    
    if pargs.mode == "copy":
        assert pargs.copy_name is not None, "You need to specify a experiment to copy from with -cn."
    else:
        assert pargs.copy_name is None, "You can only copy from a experiment in copy mode."
    
    settings.read_terminal_arguments(pargs) 
    return pargs

def main():

    pargs = parse_arguments()
    local_project_dir = os.getcwd()
    localmachine = LocalMachine(local_project_dir)

    if pargs.vm is not None:
        machine = get_machine(pargs.vm, local_project_dir=local_project_dir)
    else:
        machine = localmachine
    
    # modes not requiring a target
    if pargs.mode == "status":
        encap_process_dict = get_status(machine)
        print_status(encap_process_dict)

        slurm.print_slurm_status_if_using_slurm(machine)
        exit()

    source_file_path = pargs.target
    # Check if it will run in folder mode or file mode
    if os.path.isdir(source_file_path):
        if source_file_path[-1] == "/":
            source_file_path = source_file_path[:-1]
        
        root_path, folder_name = extract_folder_name(source_file_path)

        folder_name = os.path.join(root_path, "0encap_folder", folder_name)
        run_folder_name = os.path.join(folder_name, str(pargs.name))

        is_file = False

    elif os.path.isfile(source_file_path):
        folder_name, file_extension = filename_and_file_extension(source_file_path)

        # Get the file name
        script_name = os.path.basename(source_file_path)

        run_folder_name = os.path.join(folder_name, str(pargs.name))
        settings.config["script_name"] = script_name
        target_file_path = os.path.join(run_folder_name, script_name)
        is_file = True


    else:
        assert False, f"{source_file_path} is neither a directory nor a file."
    

    if pargs.mode == "run":
        if is_file:
            mode_run_file(folder_name=folder_name, run_folder_name=run_folder_name,
                          source_file_path=source_file_path, local_project_dir=local_project_dir, target_file_path=target_file_path,
                          machine=machine, pargs=pargs)
        else:
            mode_run_folder(folder_name=folder_name, run_folder_name=run_folder_name,
                            source_file_path=source_file_path, local_project_dir=local_project_dir,
                            machine=machine, pargs=pargs)
        
    elif pargs.mode == "rerun":
        mode_rerun_file(run_folder_name=run_folder_name,
                        local_project_dir=local_project_dir,
                        machine=machine, pargs=pargs)
        
        
    elif pargs.mode == "tail":
        # ignore_errors=True so a missing pid file (slurm job that already
        # finished, or a non-slurm experiment whose pid was cleaned up by a
        # previous tail) doesn't raise inside run_code_local before we can
        # decide what to do.
        code = f"cat {run_folder_name}/pid"
        out = machine.run_code(code, verbose=False, output=True, ignore_errors=True)

        pid = None
        if len(out) == 1:
            content = out[0].strip()
            if content.isdigit():
                pid = content
                print("PID: ", pid)
            elif content.startswith("slurm:"):
                pid = content
                print("Slurm Job: ", content.split(":", 1)[1])
            # else: cat error message ("cat: ...: No such file or directory"),
            # or some other unexpected content -- treat as no tracked process.

        machine.pull(run_folder_name, run_folder_name, directory=True)
        tail_pull(machine, run_folder_name, pid)
        if pid is not None:
            remove_process_from_database(pargs.target, pargs.name)

    elif pargs.mode == "copy":
        copy_folder_name = os.path.join(folder_name, pargs.copy_name)
        if not machine.exists(copy_folder_name):
            print(f"encap error: {copy_folder_name} does not exist.")
            exit()
        
        create_folder_and_check_if_experiment_exists(machine, pargs, folder_name, run_folder_name)
        
        if is_file:
            file_list = [script_name]
        else:
            # Copy the files named like files from source_file_path from copy_folder_name to run_folder_name
            file_list = os.listdir(source_file_path) 
        
        if os.path.isfile(os.path.join(copy_folder_name, ".encap.conf")):
            file_list += [".encap.conf"]
        
        for file in file_list:
            copy_file_path = os.path.join(copy_folder_name, file)
            target_file_path = os.path.join(run_folder_name, file)
            machine.push(copy_file_path, target_file_path, directory=False, verbose=True)
        
        if settings.debug : print(f"Files copied from {copy_folder_name} to {run_folder_name}")

        print(f"""You can run the copied experiment with:
encap rerun {pargs.target} -n {pargs.name}""")


    elif pargs.mode == "kill":
        encap_process_dict = get_status(machine, run_folder_name=run_folder_name, id=pargs.i)

        if len(encap_process_dict) == 0:
            print(f"No process {run_folder_name} found.")
            exit()
        
        print("Killing:")
        print_status(encap_process_dict)
        code = ""

        for pid in encap_process_dict:
            code += f"kill {pid}; "
        machine.run_code(code, verbose=True)
    
    elif pargs.mode == "tar":
        if pargs.name is not None:
            machine.tar(run_folder_name, verbose=True)
        else:
            machine.tar(folder_name, verbose=True, subfolders=True, threads=pargs.tar_cpus)

    elif pargs.mode == "untar":
        if pargs.name is not None:
            machine.untar(run_folder_name, verbose=True)
        else:
            machine.untar(folder_name, verbose=True, subfiles=True, threads=pargs.tar_cpus)

    elif pargs.mode == "pull":
        assert pargs.name is not None, "The experiment_name -n must be specified."
        machine.pull(run_folder_name, run_folder_name, directory=True, verbose=True)

    elif pargs.mode == "push":
        assert pargs.name is not None, "The experiment_name -n must be specified."
        machine.push(run_folder_name, run_folder_name, directory=True, verbose=True)
    
    else:
        raise ValueError(f"The mode '{pargs.mode}' is not available.")


if __name__ == "__main__":
    main()
