#!python

import scm.plams
from scm.plams import *

plams_namespace = globals().copy()

VERSION = scm.plams.__version__

import argparse
import os
import sys
import traceback
from os.path import join as opj
import warnings

warnings.warn(
    "PLAMS launch script is deprecated. Please use 'amspython' or another python interpreter and run your PLAMS script directly, making sure to import required components from 'scm.plams'.",
    DeprecationWarning,
)

parser = argparse.ArgumentParser(description="PLAMS launch script")
parser.add_argument("--version", action="version", version="PLAMS " + str(VERSION), help="show version number and exit")
parser.add_argument(
    "-p",
    "--path",
    type=str,
    default=None,
    help="place where the main working folder is created",
    metavar="path",
    dest="path",
)
parser.add_argument(
    "-f", "--folder", type=str, default=None, help="name of the main working folder", metavar="name", dest="folder"
)
parser.add_argument(
    "-v",
    "--var",
    action="append",
    type=str,
    default=[],
    help="declare a variable 'var' with a value 'value' in the global namespace. Multiple variables can be set this way, but each one requires a separate '-v'",
    metavar="var=value",
    dest="vars",
)
parser.add_argument(
    "-l",
    "--load",
    action="append",
    type=str,
    default=[],
    help="load all jobs from the given location before executing the script. Multiple paths can be given, but each one requires a separate '-l'",
    metavar="path",
    dest="load",
)
parser.add_argument(
    "-r",
    "--restart",
    action="store_true",
    help="perform a restart run (import all jobs from the folder given by -f argument and use the same folder for the current run)",
    dest="restart",
)
parser.add_argument("file", nargs="+", type=str, help="file with PLAMS script")
args = parser.parse_args()


# add -v variables to the plams_namespace
for pair in args.vars:
    if "=" in pair:
        var, val = pair.split("=")
        plams_namespace[var] = val


# read and concatenate input file(s)
inputscript = ""
for input_file in args.file:
    if os.path.isfile(input_file):
        sys.path.insert(0, os.path.dirname(os.path.abspath(input_file)))
        with open(input_file, "r", encoding="utf-8") as f:
            inputscript += f.read()
    else:
        print("Error: File {} not found".format(input_file))
        sys.exit(1)


# handle restart
if args.restart:
    if args.folder:
        restartdir = opj(args.path, args.folder) if args.path else args.folder
    else:
        basename = opj(args.path, "plams_workdir") if args.path else "plams_workdir"
        restartdir = basename
        n = 2
        nextname = basename + "." + str(n).zfill(3)
        while os.path.exists(nextname):
            n += 1
            restartdir = nextname
            nextname = basename + "." + str(n).zfill(3)

    if os.path.isdir(restartdir):
        restartdir = restartdir.rstrip("/")
        if os.listdir(restartdir):
            restart_backup = restartdir + ".res"
            n = 1
            while os.path.exists(restart_backup):
                n += 1
                restart_backup = restartdir + ".res" + str(n)
            os.rename(restartdir, restart_backup)
            print("RESTART: Moving {} to {} and restarting from it".format(restartdir, restart_backup))
            args.load.append(restart_backup)
    else:
        print("RESTART: The folder specified for restart does not exist. Ignoring -r flag.")


# initialize PLAMS
init(path=args.path, folder=args.folder)

# write down input script
with open(config.default_jobmanager.input, "w", encoding="utf-8") as f:
    f.write(inputscript)

# load jobs from -l folders
for path in args.load:
    print("LOAD: Loading jobs from " + path)
    load_all(path)

# execute input script
try:
    exec(
        compile(
            open(config.default_jobmanager.input, encoding="utf-8").read(), config.default_jobmanager.input, "exec"
        ),
        plams_namespace,
    )
except KeyboardInterrupt:
    sys.exit(0)
except Exception as e:
    exc_type, exc_obj, exc_tb = sys.exc_info()
    tb = traceback.extract_tb(exc_tb)
    fname, lineno, fn, text = tb[-1]
    err_msg = "Execution interrupted by the following exception:\n"
    err_msg += "{}: {}\n".format(exc_type.__name__, str(e))
    err_msg += "File: {}\n".format(os.path.basename(fname))
    err_msg += "Line {}: {}\n\n".format(lineno, text)
    err_msg += "==============Full traceback========================"
    for fname, lineno, fn, text in tb:
        err_msg += "\nFile: {}".format(os.path.basename(fname))
        err_msg += "\nLine {}: {}".format(lineno, text)
        err_msg += "\n----------------------------------------------------"
    log(err_msg)

# clean the environment
finish()
