# Copyright (C) 2019 OpenEye Scientific Software, Inc.
"""
usage:
    add_experiment_to_context.py CONTEXT_ID CODE
    add_experiment_to_context.py CONTEXT_ID --list LIST
    add_experiment_to_context.py CONTEXT_ID --dir DIR

arguments:
    CODE    Single experiment code to add
    LIST    Text file with one experiment code per line
    DIR     Directory containing structure files. List of codes
            will be created for all files found.

"""
import docopt
import os
import sys

import mmdsclient as mmdscli
from oeclient_utils.exceptions import BadResponse, NotFoundResponse

from spruce.utils.file_utils import get_pdbname_from_filename, get_pdb_files

def main():
    args = docopt.docopt(__doc__)

    codes = []
    if args['CODE']:
        codes.append(args['CODE'])
    elif args['DIR']:
        for filename in get_pdb_files(args['DIR']):
            code = get_pdbname_from_filename(filename)
            codes.append(code)
    else:
        if os.path.isfile(args['LIST']):
            with open(args['LIST']) as ifp:
                for line in ifp:
                    codes.append(line.strip())

    if len(codes) == 0:
        print("No codes input, exiting.")
        return 1

    if 'MMDS_PROFILE' not in os.environ:
        print("Please set MMDS_PROFILE environment variable")
        return 1

    session = mmdscli.get_session(os.environ['MMDS_PROFILE'])

    # verify we have a valid context
    context_id = args['CONTEXT_ID']
    try:
        context = session.get_resource(mmdscli.Context, context_id)
    except (NotFoundResponse, BadResponse):
        print('{} is not a valid MMDS Context'.format(context_id))
        return 1

    print("Adding {} experiments to context id: {}".format(len(codes), context_id))
    for code in codes:
        print('adding {}'.format(code))
        try:
            for expt in session.list_resources(mmdscli.Experiment, {"code": code}):
                for frame in mmdscli.add_experiment_to_context(session, context_id, expt.id):
                    print("  code: {} -> frame: {}".format(code, frame.id))
        except (NotFoundResponse, BadResponse) as e:
            msg = str(e)
            if "already added" in e.reason:
                print(msg)
            else:
                print('  Unable to add code: {}'.format(code))


if __name__ == "__main__":
    sys.exit(main())