import argparse
import logging
from datetime import datetime

import config
from oa_contract_ingest_to_abacus import oa_contract_lambda_execute, oa_contract_ingest_to_abacus

parser = argparse.ArgumentParser(
    description='Ingest OA contract data to Abacus. Load them from Snowflake, then upload to s3 bucket for further processing.')
parser.add_argument('--account_ids',
                    help='List of accounts ids to process separated by comma.',
                    required=False)
parser.add_argument('--batch_keys',
                    help='List of batch keys to process separated by comma. No export from OA is made in this case.',
                    required=False)
parser.add_argument('--skip_abacus_import',
                    help='Use this param to skip the lambda running step. Default is False.',
                    action='store_true',
                    required=False,
                    default=False
                    )
parser.add_argument('--import_user',
                    help='Username that is used to run the lambda. Will be stored to identify who initiated the import.',
                    required=False
                    )
parser.add_argument('--jira_ticket',
                    help='The ticket number which scope includes this import.',
                    required=False
                    )
parser.add_argument('--target_env',
                    help='The environment where the import will be run. Default is QA.',
                    required=False,
                    )

args = parser.parse_args()

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    handlers=[
        logging.FileHandler(f'logs/ingest_logs_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log'),
        logging.StreamHandler()
    ]
)


def main():
    account_ids = []
    if args.account_ids:
        account_ids = args.account_ids.split(',')

    batch_keys = []
    if args.batch_keys:
        batch_keys = args.batch_keys.split(',')

    skip_abacus_import = False
    if args.skip_abacus_import:
        skip_abacus_import = True

    import_user = config.IMPORT_USER
    if args.import_user:
        import_user = args.import_user

    jira_ticket = config.JIRA_TICKET
    if args.jira_ticket:
        jira_ticket = args.jira_ticket

    lambda_function_name = _get_lambda_name()
    s3_bucket = _get_s3_bucket()

    logging.info(
        f'Running with arguments: '
        f'account_id={account_ids}, '
        f'batch_keys={batch_keys}, '
        f'skip_abacus_import={skip_abacus_import}, '
        f'import_user={import_user}, '
        f'jira_ticket={jira_ticket}, '
        f'lambda_function_name={lambda_function_name}, '
        f's3_bucket={s3_bucket}')

    if batch_keys and account_ids:
        raise Exception('Cannot specify both batch_keys and account_id')

    if not import_user or not jira_ticket:
        raise Exception('import_user and jira_ticket are required')

    if batch_keys:
        oa_contract_lambda_execute.run_lambda_for_batches(
            batch_keys,
            import_user,
            jira_ticket,
            lambda_function_name
        )
    else:
        oa_contract_ingest_to_abacus.ingest_contract_to_abacus(
            account_ids,
            skip_abacus_import,
            import_user,
            jira_ticket,
            lambda_function_name,
            s3_bucket,
        )


def _get_s3_bucket():
    if args.target_env:
        return f'{args.target_env}-abacus-json-contract-file-import'
    return config.S3_BUCKET


def _get_lambda_name():
    if args.target_env:
        return f'{args.target_env}-lambda-abacus-json-contract-file-import'
    return config.LAMBDA_FUNCTION_NAME


main()