import os
import logging

logger = logging.getLogger('DockOnSurf')


def check_bak(file_name):
    """Checks if a file already exists and backs it up if so.
    @param file_name: file to be checked if exists
    """
    new_name = file_name
    bak_num = 0
    while os.path.isdir(new_name)  or os.path.isfile(new_name):
        bak_num += 1
        new_name = new_name.split(".bak")[0] + f".bak{bak_num}"
    if bak_num > 0:
        os.rename(file_name, new_name)
        logger.warning(f"'{file_name}' already present. Backed it up to "
                       f"{new_name}")


def prep_cp2k(inp_file, run_type, atms_list):
    """Prepares the directories to run isolated calculation with CP2K.

    @param inp_file: CP2K Input file to run the calculations with.
    @param run_type: Type of calculation. 'isolated', 'screening' or
        'refinement'
    @param atms_list: list of ase.Atoms objects to run the calculation of.
    @return: None
    """
    from shutil import copy
    import ase.io
    from pycp2k import CP2K
    cp2k = CP2K()
    cp2k.parse(inp_file)
    force_eval = cp2k.CP2K_INPUT.FORCE_EVAL_list[0]
    if force_eval.SUBSYS.TOPOLOGY.Coord_file_name is None:
        logger.warning("'COORD_FILE_NAME' not specified on CP2K input. Using\n"
                       "default name 'coord.xyz'. A new CP2K input file with "
                       "the 'COORD_FILE_NAME' variable is created. If there\n"
                       "is a name conflict the old file will be backed up")
        force_eval.SUBSYS.TOPOLOGY.Coord_file_name = 'coord.xyz'
        print(inp_file.split('/')[-1])
        check_bak(inp_file.split('/')[-1])
        cp2k.write_input_file(inp_file.split('/')[-1])

    coord_file = force_eval.SUBSYS.TOPOLOGY.Coord_file_name

    # Creating and setting up directories for every configuration.
    for i, conf in enumerate(atms_list):
        os.mkdir(f'{run_type}/conf_{i}')
        copy(inp_file, f'{run_type}/conf_{i}/')
        ase.io.write(f'{run_type}/conf_{i}/{coord_file}', conf)


def get_jobs_status_sge(job_ids):  # TODO more elegant
    """Returns a list of job status for a list of job ids.

    @param job_ids: list of all jobs to be checked their status.
    @return: list of status for every job.
    """
    from gridtk.tools import qstat
    run_chk = 'usage         1'
    status_list = []
    for job in job_ids:
        if run_chk in qstat(job):
            status_list.append('r')
        elif len(qstat(job)) > 0:
            status_list.append('q')
        else:
            status_list.append('f')
    return status_list


def submit(run_type, q_sys, sub_script, max_qw, name):
    """Submits jobs to the relevant queuing system with the provided script

    @param run_type: Type of calculation. 'isolated', 'screening' or
        'refinement'
    @param q_sys: Batch queuing system used
    @param sub_script: script for the job submission.
    @param max_qw: Maximum number of simultaneous jobs waiting to be executed.
    @param name: name of the project
    """
    from shutil import copy
    from time import sleep
    subm_jobs = []
    if q_sys == 'sge': # TODO Independent functions for every batch_q_sys
        from gridtk.tools import qsub  # TODO CHANGE TO DRMAA
        init_dir = os.getcwd()
        for conf in os.listdir(run_type):
            i = conf.split('_')[1]
            while get_jobs_status_sge(subm_jobs).count('q') >= max_qw:
                sleep(30)
            copy(sub_script, f"{run_type}/{conf}")
            os.chdir(f"{run_type}/{conf}")
            job_name = f'{name[:6].capitalize()}{run_type[:3].capitalize()}{i}'
            subm_jobs.append(qsub(sub_script, name=job_name))
            os.chdir(init_dir)

        logger.info('All jobs have been submitted, waiting them to finish')
        while not all([stat == 'f' for stat in get_jobs_status_sge(subm_jobs)]):
            sleep(30)
        logger.info('All jobs have finished')


def run_calc(run_type, inp_vars, atms_list):
    """Directs the calculation run according to the provided arguments.

    @param run_type: Type of calculation. 'isolated', 'screening' or
    'refinement'
    @param inp_vars: Calculation parameters from input file.
    @param atms_list: List of ase.Atoms objects containing the sets of atoms
    aimed to run the calculations of.
    """
    run_types = ['isolated', 'screening', 'refinement']
    run_type_err = f"'run_type' must be one of the following: {run_types}"
    if not isinstance(run_type, str) or run_type.lower() not in run_types:
        logger.error(run_type_err)
        raise ValueError(run_type_err)

    logger.info(f"Running {run_type} calculation with {inp_vars['code']} on "
                f"{inp_vars['batch_q_sys']}")
    check_bak(run_type)
    os.mkdir(run_type)
    if inp_vars['code'] == 'cp2k':
        prep_cp2k(inp_vars['isol_inp_file'], run_type, atms_list)
    # elif: inp_vars['code'] == 'Other codes here'

    submit(run_type, inp_vars['batch_q_sys'], inp_vars['subm_script'],
           inp_vars['max_qw'], inp_vars['project_name'])