#!/usr/bin/env python3
#
# Repatriate directories after submitting several simulations on sator at once

from pathlib import Path
import argparse

import mola.naming_conventions as names
from mola.logging import MolaException
from mola.server import get_one_case, get_all_cases_from_workflow_manager

__DOC__ = (
    "repatriate data from a remote machine after submitting one simulation with a Workflow "
    "or several simulations with a WorkflowManager. "
    f"There has to be a file {names.FILE_LOG_REMOTE_DIRECTORY} "
    f"or a file {names.FILE_WORKFLOW_MANAGER} in the current directory. "
    "This command must be used with at least one option."
)

def parse_args():
    # Predefined file patterns
    BASE_INCLUDES = set(names.STATUS_FILES)
    IMAGE_PATTERNS = ['**.png', '**.pdf', '**.svg']
    LOG_PATTERNS = ['*.log', f'{names.DIRECTORY_LOG}**']
    SIGNALS_PATTERNS = [names.DIRECTORY_OUTPUT, f'*/{names.FILE_OUTPUT_1D}']
    EXTRACTIONS_PATTERNS = [names.DIRECTORY_OUTPUT, f'*/{names.FILE_OUTPUT_2D}']
    OUTPUT_PATTERN = [f'{names.DIRECTORY_OUTPUT}**']

    LIGHT_PATTERNS = IMAGE_PATTERNS + SIGNALS_PATTERNS + EXTRACTIONS_PATTERNS

    parser = argparse.ArgumentParser(usage=__DOC__)
    args_group = parser.add_mutually_exclusive_group()
    args_group.add_argument('-f', '--full', action='store_true', help='Copy all files')
    args_group.add_argument('-l', '--light', help=(
        f'Copy main files for data analysis ({names.FILE_OUTPUT_1D}, '
        f'{names.FILE_OUTPUT_2D}, plots, figures, etc.)'),
        dest='patterns', action='append_const', const=LIGHT_PATTERNS
        )
    
    parser.add_argument('--log', help='Copy log files', 
                      dest='patterns', action='append_const', const=LOG_PATTERNS
                      )
    parser.add_argument('-s', '--signals', help=f'Copy {names.FILE_OUTPUT_1D}', 
                        dest='patterns', action='append_const', const=SIGNALS_PATTERNS
                        )
    parser.add_argument('-e', '--extractions', help=f'Copy {names.FILE_OUTPUT_2D}', 
                        dest='patterns', action='append_const', const=EXTRACTIONS_PATTERNS
                        )
    parser.add_argument('-o', '--output', help=f'Copy directory {names.DIRECTORY_OUTPUT}', 
                        dest='patterns', action='append_const', const=OUTPUT_PATTERN
                        )

    args = parser.parse_args()
    if not args.patterns and not args.full:
        raise MolaException('mola_repatriate must be used with at least one option. See help with mola_repatriate --help')

    included_files = set(BASE_INCLUDES)
    excluded_files = set()

    if not args.full:
        # args.patterns is a list of lists, so we need to flatten it
        patterns = [p for patterns_list in args.patterns for p in patterns_list]
        included_files.update(patterns)
        # exclude all other files
        excluded_files = set('*')

    return included_files, excluded_files

def main():

    included_files, excluded_files = parse_args()

    is_remote_case_file = Path(names.FILE_LOG_REMOTE_DIRECTORY).is_file()
    is_workflow_manager_file = Path(names.FILE_WORKFLOW_MANAGER).is_file()

    if is_remote_case_file and is_workflow_manager_file:
        raise MolaException((
            f'Found both files {names.FILE_LOG_REMOTE_DIRECTORY} and {names.FILE_WORKFLOW_MANAGER}. '
            'Remove the outdated one and try again.'
        ))
    
    elif is_remote_case_file and not is_workflow_manager_file:
        get_one_case(included_files, excluded_files)

    elif not is_remote_case_file and is_workflow_manager_file:
        get_all_cases_from_workflow_manager(included_files, excluded_files)

    else:
        raise MolaException(f'Neither {names.FILE_LOG_REMOTE_DIRECTORY} nor {names.FILE_WORKFLOW_MANAGER} exists.')

if __name__ == '__main__':
    main()
