#!/usr/bin/env python
# PYTHON_ARGCOMPLETE_OK

"""Script that helps create and validate pipelines from command line
"""
from argparse import ArgumentParser
from datetime import timedelta
from pytimeparse import parse

from dataduct.utils.cli import *  # noqa

import logging
logger = logging.getLogger(__name__)

PIPELINE = 'pipeline'
CREATE = 'create'
VALIDATE = 'validate'
ACTIVATE = 'activate'
VISUALIZE = 'visualize'
SQL_SHELL = 'sql_shell'

CONFIG = 'config'
CONFIG_TO_S3 = 'sync_to_s3'
CONFIG_FROM_S3 = 'sync_from_s3'

DATABASE = 'database'
DROP = 'drop'
GRANT = 'grant'
RECREATE = 'recreate'

REDSHIFT = 'redshift'
MYSQL = 'mysql'


def initialize_etl_objects(pipeline_definitions, time_delta=None,
                           frequency_override=None, backfill=False):
    """Generate etl objects from yaml files
    """
    from dataduct.etl import create_pipeline
    from dataduct.etl import read_pipeline_definition

    # Convert the time_delta if it exists
    if time_delta is not None:
        time_delta = timedelta(seconds=parse(time_delta))
        if backfill:
            time_delta *= -1

    etls = []
    for pipeline_definition in pipeline_definitions:
        definition = read_pipeline_definition(pipeline_definition)
        if time_delta is not None:
            definition.update({'time_delta': time_delta})
        if frequency_override is not None:
            definition.update({'frequency': frequency_override})
        etls.append(create_pipeline(definition))
    return etls


def config_actions(action, filename=None, **kwargs):
    """Config related actions are executed in this block
    """
    from dataduct.config.config_actions import sync_to_s3
    from dataduct.config.config_actions import sync_from_s3

    if action == CONFIG_TO_S3:
        return sync_to_s3()
    return sync_from_s3(filename)


def pipeline_actions(action, pipeline_definitions, force=None, time_delta=None,
                     frequency_override=None, activities_only=None,
                     filename=None, backfill=False, **kwargs):
    """Pipeline related actions are executed in this block
    """
    from dataduct.etl import activate_pipeline
    from dataduct.etl import validate_pipeline
    from dataduct.etl import visualize_pipeline

    for etl in initialize_etl_objects(pipeline_definitions, time_delta,
                                      frequency_override, backfill):
        if action in [VALIDATE, ACTIVATE]:
            validate_pipeline(etl, force)
        if action == ACTIVATE:
            activate_pipeline(etl)
        if action == VISUALIZE:
            visualize_pipeline(etl, activities_only, filename)


def database_actions(action, table_definitions, filename=None, execute=False,
                     **kwargs):
    """Database related actions are executed in this block
    """
    from dataduct.database import Database
    from dataduct.data_access import redshift_connection
    import pandas.io.sql as pdsql

    script = None
    database = Database(files=table_definitions)
    if action == CREATE:
        script = database.create_relations_script()
    elif action == DROP:
        script = database.drop_relations_script()
    elif action == GRANT:
        script = database.grant_relations_script()
    elif action == RECREATE:
        script = database.recreate_relations_script()
    elif action == VISUALIZE:
        database.visualize(filename)

    # TODO: Build execution options
    if script:
        print script

    if execute:
        logger.info('Getting redshift connection...')
        connection = redshift_connection()
        logger.info('Executing query...')
        pdsql.execute(script.wrap_transaction().sql(), connection)
        logger.info('Query executed successfully!')


def main():
    """Main function that parses the command line arguments
    """
    parser = ArgumentParser(
        description='Run Dataduct commands',
        add_help=False,
        parents=[help_parser],
        formatter_class=formatter_class,
    )
    subparsers = parser.add_subparsers(
        dest='command',
        help='Actions for various features',
    )

    # Pipeline parser
    pipeline_parser = subparsers.add_parser(
        PIPELINE,
        formatter_class=formatter_class,
        add_help=False,
        parents=[help_parser]
    )
    pipeline_subparsers = pipeline_parser.add_subparsers(
        dest='action',
        help='Pipeline actions',
    )

    # Pipeline subparsers_action
    pipeline_subparsers.add_parser(
        CREATE,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            pipeline_run_options,
            pipeline_definition_parser,
        ],
        help='Create a pipeline locally',
    )
    pipeline_subparsers.add_parser(
        VALIDATE,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            pipeline_run_options,
            pipeline_definition_parser,
        ],
        help='Validate a pipeline with AWS without activating',
    )
    pipeline_subparsers.add_parser(
        ACTIVATE,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            pipeline_run_options,
            pipeline_definition_parser,
        ],
        help='Activate the pipeline on AWS',
    )
    visualize_pipeline_parser = pipeline_subparsers.add_parser(
        VISUALIZE,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            file_parser,
            pipeline_definition_parser,
        ],
        help='Visualize the pipeline',
    )
    visualize_pipeline_parser.add_argument(
        '--activities_only',
        action='store_true',
        default=False,
        help='Visualize only activities',
    )

    # Config parser
    config_parser = subparsers.add_parser(
        CONFIG,
        formatter_class=formatter_class,
        add_help=False,
        parents=[help_parser]
    )
    config_subparsers = config_parser.add_subparsers(
        dest='action',
        help='config actions',
    )

    # config subparsers_action
    config_subparsers.add_parser(
        CONFIG_TO_S3,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
        ],
        help='sync config file from local to s3',
    )
    config_subparsers.add_parser(
        CONFIG_FROM_S3,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            file_parser,
        ],
        help='sync config file from s3 to local file',
    )

    # Database parser
    database_parser = subparsers.add_parser(
        DATABASE,
        formatter_class=formatter_class,
        add_help=False,
        parents=[help_parser]
    )
    database_subparsers = database_parser.add_subparsers(
        dest='action',
        help='database actions',
    )

    # database subparsers_action
    database_subparsers.add_parser(
        CREATE,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            table_definition_parser,
            execute_sql_parser,
        ],
        help='Create tables',
    )
    database_subparsers.add_parser(
        DROP,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            table_definition_parser,
            execute_sql_parser,
        ],
        help='Drop views and tables',
    )
    database_subparsers.add_parser(
        GRANT,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            table_definition_parser,
            execute_sql_parser,
        ],
        help='Grant permissions to neccessary groups',
    )
    database_subparsers.add_parser(
        RECREATE,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            table_definition_parser,
            execute_sql_parser,
        ],
        help='Recreate tables, load new data, drop old tables',
    )
    database_subparsers.add_parser(
        VISUALIZE,
        formatter_class=formatter_class,
        parents=[
            mode_parser,
            file_parser,
            table_definition_parser,
        ],
        help='Visualize the database er-diagram',
    )

    # SQL Shell parser
    sql_shell_parser = subparsers.add_parser(
        SQL_SHELL,
        formatter_class=formatter_class,
        add_help=False,
        parents=[help_parser]
    )
    sql_shell_subparser = sql_shell_parser.add_subparsers(
        dest='database_type',
        help='Database Type (MySQL or Redshift)',
    )
    sql_shell_subparser.add_parser(
        REDSHIFT,
        formatter_class=formatter_class,
        parents=[mode_parser, help_parser],
        add_help=False
    )
    sql_shell_subparser.add_parser(
        MYSQL,
        formatter_class=formatter_class,
        parents=[mode_parser, help_parser, host_alias_parser],
        add_help=False
    )

    # Check if autocomplete is possible
    try:
        import argcomplete
        argcomplete.autocomplete(parser)
    except ImportError:
        pass
    args = parser.parse_args()
    config = config_singleton_setup(args)

    # Frequency override
    if hasattr(args, 'frequency') and args.frequency is not None:
        frequency_override = args.frequency
    else:
        # Certain modes in the config can override frequency of a pipeline
        frequency_override = config.etl.get('FREQUENCY_OVERRIDE', None)

    arg_vars = vars(args)

    # Action parse
    if args.command == CONFIG:
        config_actions(**arg_vars)
    elif args.command == PIPELINE:
        pipeline_actions(frequency_override=frequency_override, **arg_vars)
    elif args.command == DATABASE:
        database_actions(**arg_vars)
    elif args.command == SQL_SHELL:
        open_sql_shell(**arg_vars)
    else:
        raise ValueError('Unknown argument provided, use dataduct -h')


if __name__ == '__main__':
    main()
