#!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

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()
