"""
Base Feed Ingestion Work Flow.

Base Work Flow class that new Feed Ingestion flows should extend. Standardizes
how flow domain, name, and workflow_id are generated and adds default
on_exception method.

Flows should override the decider method.
"""

import datetime
from functools import wraps
import json
import logging
import os

import boto3
from botocore.config import Config
from garcon import activity
from garcon import param
from garcon import runner
from garcon import task
from garcon.param import StaticParam
from garcon_contrib.dynamo_feed_status import garcon_feed_status

from feed_ingestion import conf
from feed_ingestion.common import status_runner
from feed_ingestion.conf.config import AWS_REGION
from feed_ingestion.tasks import date_tasks
from feed_ingestion.tasks import load_fact_tables_tasks_sf
from feed_ingestion.tasks import load_marketshare_table_tasks
from feed_ingestion.tasks import load_raw_table_tasks_sf
from feed_ingestion.tasks import notification_tasks
from feed_ingestion.tasks import overall_status_tasks
from feed_ingestion.tasks import youtube_tasks
from feed_ingestion.util import sentry_util
from feed_ingestion.util.query_tag import clear_query_tag
from feed_ingestion.util.query_tag import set_query_tag
from feed_ingestion.util.sentry_util import SentryRunnerMixin

PROD_DOMAIN = 'prod_swf_feed_ingestion'
QA_DOMAIN = 'qa_swf_feed_ingestion'
TEST_DOMAIN = 'test'
DEV_DOMAIN = 'dev'
SHADOW_DOMAIN = 'shadow_swf_feed_ingestion'

logger = logging.getLogger('garcon')


class FlowBase(object):
    """Base Garcon Workflow that other Flows can extend."""

    def __init__(self, feed_name, version):
        """Initialize Flow.

        Initializes the Flow's:
            domain - 'prod_swf_feed_ingestion' or 'dev
            feed_name - name of feed used in the status table
            name - '<feed_name>_feed_ingestion' ex. qq_feed_ingestion
                (this is the workflow name as opposed to the feed name)

        Args:
            feed_name (str): Name of the feed the workflow is processing,
                with underscores replacing spaces. ex. 'qq'
            version (str): Version of the workflow type to filter on. should
                be a 'x.x' float (ex. 1.0, 1.1, etc.)
        """
        self.domain = get_domain()
        self.name = generate_feed_name(feed_name)  # name = workflow_name
        self.feed_name = feed_name  # feed name verified in generate_feed_name
        self.version = version
        self.client = boto3.client('swf', config=Config(
            connect_timeout=60,
            region_name=AWS_REGION,
            read_timeout=180,
            retries={'max_attempts': 2}))
        self.create = activity.create(
            self.client, self.domain, self.name, version=self.version,
            on_exception=self.on_exception)
        self.create_status_runner = status_runner.create(feed_name)

        self.task_timeout = '1200'
        # 4 hours by default for start-to-close timeout for SWF run
        self.timeout = 3600 * 4

    @staticmethod
    def get_kwargs_namespace_for(activity_name):
        """Get kwargs namespace for specified activity.

        If activity defines kwargs it needs to define if within namespace
        returned by this function.

        Args:
            activity_name (str): name of the activity.

        Returns:
            str: namespace.
        """
        return activity_name + '_kwargs'

    def on_exception(self, actor, exception):
        """Capture an exception that has occurred in the application.

        Args:
            actor (ActivityWorker, DeciderWorker): The actor that has received
                the exception.
            exception (Exception): The exception to capture.
        """
        # client grabs Sentry DSN from SENTRY_DSN environment variable
        if os.environ.get('SENTRY_DSN'):
            # we can add a non-actionable exception to
            # the sentry_util.SEND_AS_WARNINGS list
            sentry_util.send_error_or_warning(exception)

        if isinstance(actor, activity.Activity):
            actor.logger.error(exception, exc_info=True)
        else:
            logger.error(exception, exc_info=True)

    def workflow_id(self, initial_context):
        """Generate workflow id.

        Assumes a unique workflow is defined by it's context date and
        defaults to today's date. If that is not the case this method
        should be overridden.

        Args:
            initial_context (dict): The initial context for the flow.
        Returns:
            str: A unique identifier for a workflow being executed
                In the forms of '<self.flow_name>-YYYY-MM-DD', where
                YYYY-MM-DD is the context date or, if none passed the
                current date.
        """
        if 'context_date' not in initial_context:
            date = datetime.datetime.today().strftime('%Y-%m-%d')
        else:
            date = initial_context['context_date']

        return '{flow_name}-{date}'.format(
            flow_name=self.name,
            date=date)

    def contextified_feed_name(self, context):
        """Get feed_name in context.

        Get feed name if it depends on context.
        Defaults to flow's feed_name otherwise this method should
        be overwritten.
        Typically it's used to set status in DynamoDB
        (e.g. if status name in DynamoDB is constructed from
        feed_name and values from context).

        Args:
            context (dict): The context of the flow.
        Returns:
            str: Contextified feed name.
        """
        return self.feed_name

    def create_kwargs_collect_activity_for(self, activity_name, mapping):
        """Create activity that collects values specified by mapping.

        Usage:
            @property
            def collect_kwargs_activity(self):
                return self.create_kwargs_collect_activity_for(
                    activity_name='load_staging_raw_table',
                    mapping=dict(
                        source_file_name='bootstrap.source_file_name',
                        source_file_size='move_raw_file.source_file_size'
                    )
                )
        Args:
            activity_name (str): name of the activity accepting kwargs.
            mapping (dict): namespace key mapping.

        Returns:
            Activity: activity object.
        """
        class NamespaceMappingParam(param.BaseParam):
            def __init__(self, **mapping):
                super().__init__()
                self.mapping = mapping

            @property
            def requirements(self):
                """Return all reqs defined in mapping."""
                yield from self.mapping.values()

            def get_data(self, context):
                """Resolve and return values for keys from mapping."""
                return {
                    k: context.get(v, None) for k, v in self.mapping.items()
                }

        @task.decorate(timeout=300)
        def collect_arguments(activity, kwargs):
            """Collect kwargs."""
            return {
                'kwargs': dict(kwargs)
            }

        return self.create(
            name=self.get_kwargs_namespace_for(activity_name),
            tasks=runner.Sync(
                collect_arguments.fill(
                    namespace=self.get_kwargs_namespace_for(activity_name),
                    kwargs=NamespaceMappingParam(**mapping)
                )
            )
        )

    # TODO: Maybe we should add 'context' parameter to match the signature of
    # almost all the overriding methods in derived Flow classes
    def decider(self, schedule):
        """Activity decider.

        Args:
            schedule (callable): The scheduler method.
        """
        raise NotImplementedError()

    @property
    def set_status_to_ingested(self):
        """Set overall feed status to INGESTED."""
        return self.create(
            name='set_status_to_ingested',
            tasks=runner.Sync(
                overall_status_tasks.set_overall_status.fill(
                    namespace='set_status_to_ingested',
                    feed_name=param.StaticParam(self.feed_name),
                    date='bootstrap.date',
                    status=param.StaticParam(
                        garcon_feed_status.STATUS_INGESTED))))

    @property
    def set_overall_status(self):
        """Set overall status."""
        return self.create(
            name='set_overall_status',
            tasks=runner.Sync(
                overall_status_tasks.set_overall_status.fill(
                    feed_name='bootstrap.feed_name',
                    date='bootstrap.date',
                    status='status')))


class FlowConfigMixin:
    """Mixin that Flows can inherit from and use to get config properties."""

    @property
    def conf_aws(self):
        """Get AWS credentials.

        Returns:
            dict: AWS credentials - {access_key: ABC, access_secret: XYZ}.
        """
        return get_aws_config()

    @property
    def conf_env(self):
        """Get env app is running in.

        Returns:
            str: Name of the environment (either 'dev' or 'prod').
        """
        return conf.getconf('env')['env']


class FlowLoadFactMixinSF:
    """Mixin that Flows can use to move data from staging_raw -> fact tables.

    Assumes the bootstrap has injected the following into the context:
        - bootstrap.date: download_data that flow is processing
        - bootstrap.feed_name: name of feed (used to get flow specific SF
            executor class.)
    """

    @property
    def bootstrap_namespace(self):
        """Return namespace of bootstrap task for activities."""
        return 'bootstrap'

    @property
    def load_staging_fact_table(self):
        """Activity to load staging_fact_analytics_{feed_name}_{datestamp}."""
        return self.create(
            name='load_staging_fact_table',
            tasks=runner.Sync(
                load_fact_tables_tasks_sf.create_staging_fact.fill(
                    namespace='create_staging_fact',
                    feed_name='{}.feed_name'.format(self.bootstrap_namespace),
                    secrets_path='{}.secrets_path'.format(
                        self.bootstrap_namespace),
                    date='{}.date'.format(self.bootstrap_namespace),
                    sfdb_params='sfdb_params'),
                load_fact_tables_tasks_sf.load_staging_fact.fill(
                    namespace='load_staging_fact',
                    feed_name='{}.feed_name'.format(self.bootstrap_namespace),
                    secrets_path='{}.secrets_path'.format(
                        self.bootstrap_namespace),
                    date='{}.date'.format(self.bootstrap_namespace),
                    sfdb_params='sfdb_params')))

    @property
    def load_fact_tables(self):
        """Activity to load fact_analytics & fact_analytics_error.

        Deletes any existing data for this feed and date in fact_analytics &
        fact_analytics_error.

        After successfully loading the data, sets the feed's status for that
        download_date to INGESTED.
        """
        activity_name = 'load_fact_tables'
        return self.create(
            name=activity_name,
            tasks=runner.Sync(
                load_fact_tables_tasks_sf.load_fact_data.fill(
                    namespace='load_fact_data',
                    feed_name='{}.feed_name'.format(self.bootstrap_namespace),
                    secrets_path='{}.secrets_path'.format(
                        self.bootstrap_namespace),
                    date='{}.date'.format(self.bootstrap_namespace),
                    date_as_in_uuid='{}.date_as_in_uuid'.format(
                        self.bootstrap_namespace),
                    sfdb_params='sfdb_params',
                    kwargs='{}.kwargs'.format(
                        FlowBase.get_kwargs_namespace_for(activity_name))),
                overall_status_tasks.set_overall_status.fill(
                    feed_name='{}.feed_name'.format(self.bootstrap_namespace),
                    date='{}.date'.format(self.bootstrap_namespace),
                    set_status_once=StaticParam(True),
                    status=StaticParam(garcon_feed_status.STATUS_INGESTED))))


class FlowLoadRawMixinSF:
    """Mixin that Flows can use to load data to staging_raw tables."""

    @property
    def bootstrap_namespace(self):
        """Return namespace of bootstrap task for activities."""
        return 'bootstrap'

    @property
    def create_temp_staging_raw_table(self):
        """Create temp staging raw table."""
        return self.create(
            name='create_temp_staging_raw_table',
            tasks=runner.Sync(
                load_raw_table_tasks_sf.create_temp_staging_raw_table.fill(
                    namespace='create_temp_staging_raw_table',
                    date='{}.date'.format(self.bootstrap_namespace),
                    sfdb_params='{}.sfdb_params'.format(
                        self.bootstrap_namespace),
                    feed_name='{}.feed_name'.format(self.bootstrap_namespace),
                    secrets_path='{}.secrets_path'.format(
                        self.bootstrap_namespace),
                    temp_staging_raw_table=('{}.temp_staging_raw_table'.format(
                        self.bootstrap_namespace)))))

    @property
    def load_temp_staging_raw_table(self):
        """Load temp staging raw table."""
        return self.create(
            name='load_temp_staging_raw_table',
            tasks=runner.Sync(
                load_raw_table_tasks_sf.load_temp_staging_raw_table.fill(
                    namespace='load_temp_staging_raw_table',
                    aws=StaticParam(get_aws_config()),
                    feed_name='{}.feed_name'.format(self.bootstrap_namespace),
                    secrets_path='{}.secrets_path'.format(
                        self.bootstrap_namespace),
                    date='{}.date'.format(self.bootstrap_namespace),
                    sfdb_params='sfdb_params',
                    # TODO: rename to s3_temp_staging_raw_path
                    key_dir='{}.s3_temp_staging_raw_bucket'.format(
                        self.bootstrap_namespace),
                    temp_staging_raw_table=('{}.temp_staging_raw_table'.format(
                        self.bootstrap_namespace)))))

    @property
    def load_staging_raw_table(self):
        """Load permanent staging raw table."""
        activity_name = 'load_staging_raw_table'
        return self.create(
            name=activity_name,
            tasks=runner.Sync(
                load_raw_table_tasks_sf.load_staging_raw_table.fill(
                    namespace=activity_name,
                    date='{}.date'.format(self.bootstrap_namespace),
                    feed_name='{}.feed_name'.format(self.bootstrap_namespace),
                    secrets_path='{}.secrets_path'.format(
                        self.bootstrap_namespace),
                    sfdb_params='sfdb_params',
                    temp_staging_raw_table='{}.temp_staging_raw_table'.format(
                        self.bootstrap_namespace),
                    staging_raw_table='{}.staging_raw_table'.format(
                        self.bootstrap_namespace),
                    clean='clean',
                    kwargs='{}.kwargs'.format(
                        FlowBase.get_kwargs_namespace_for(activity_name))),
                overall_status_tasks.set_overall_status.fill(
                    date='{}.date'.format(self.bootstrap_namespace),
                    feed_name='{}.feed_name'.format(self.bootstrap_namespace),
                    set_status_once=StaticParam(True),
                    status=StaticParam(
                        garcon_feed_status.STATUS_POPULATED_RAW_TABLE))))


class FlowLoadMarketshareMixinSF:
    """Mixin that Flows can use to move data from staging_raw -> marketshare.

    Assumes the bootstrap has injected the following into the context:
        - bootstrap.date: download_data that flow is processing
        - bootstrap.feed_name: name of feed (used to get flow specific SF
            executor class.)
    """

    @property
    def bootstrap_namespace(self):
        """Return namespace of bootstrap task for activities."""
        return 'bootstrap'

    @property
    def get_first_day_of_month(self):
        """Get the first day of the month of date."""
        return self.create(
            name='get_first_day_of_month',
            tasks=runner.Sync(
                date_tasks.get_first_day_of_month.fill(
                    namespace='get_first_day_of_month',
                    date='context_date')))

    @property
    def load_marketshare_table(self):
        """Activity to load main_market_share.

        Deletes any existing data for this feed and date in main_market_share.
        After that populate main market share table.
        """
        activity_name = 'load_market_share_table'
        return self.create(
            name=activity_name,
            tasks=runner.Sync(
                load_marketshare_table_tasks.load_marketshare_data.fill(
                    namespace='load_marketshare_data',
                    feed_name='{}.feed_name'.format(self.bootstrap_namespace),
                    secrets_path='{}.secrets_path'.format(
                        self.bootstrap_namespace),
                    date='{}.date'.format(self.bootstrap_namespace),
                    sfdb_params='sfdb_params',
                    kwargs='{}.kwargs'.format(
                        FlowBase.get_kwargs_namespace_for(activity_name)))))

    @property
    def notify_new_ms_files(self):
        """Activity to notify about new market share files."""
        return self.create(
            name='notify_new_ms_files',
            tasks=runner.Sync(
                notification_tasks.notify_new_ms_files.fill(
                    namespace='notify_new_ms_files',
                    date='{}.date'.format(self.bootstrap_namespace),
                    feed_name='{}.feed_name'.format(
                        self.bootstrap_namespace))))

    @property
    def reset_dynamo_db_status(self):
        """Activity to reset DynamoDB status."""
        return self.create(
            name='reset_dynamo_db_status',
            tasks=runner.Sync(
                overall_status_tasks.delete_overall_status.fill(
                    namespace='reset_dynamo_db_status',
                    date='{}.date'.format(self.bootstrap_namespace),
                    feed_name='{}.feed_name'.format(
                        self.bootstrap_namespace))))


class FlowYouTubeMixin:
    """Mixin that flows can use to perform YouTube specific tasks.

    Assumes the bootstrap has injected the following into the context:
        - bootstrap.date: download_data that flow is processing
        - bootstrap.feed_name: name of feed (used to get flow specific SF
            executor class.)
        - bootstrap.secrets_path: flow's secrets manager path.
    """

    @property
    def bootstrap_namespace(self):
        """Return namespace of bootstrap task for activities."""
        return 'bootstrap'

    @property
    def update_channel_names_table(self):
        """Activity to update channel names mapping table."""
        return self.create(
            name='update_channel_names_table',
            tasks=runner.Sync(
                youtube_tasks.update_channel_names_table.fill(
                    namespace='update_channel_names_table',
                    date='{}.date'.format(self.bootstrap_namespace),
                    feed_name='{}.feed_name'.format(
                        self.bootstrap_namespace),
                    sfdb_params='sfdb_params',
                    secrets_path='{}.secrets_path'.format(
                        self.bootstrap_namespace))))

    @property
    def grab_reports_files(self):
        """Archive report files to archive location."""
        return self.create(
            name='grab_reports_files',
            schedule_to_start=48000,
            tasks=runner.Sync(
                youtube_tasks.grab_reports_files.fill(
                    namespace='grab_reports_files',
                    report_name=f'{self.bootstrap_namespace}.report_name',
                    report_status_name=f'{self.bootstrap_namespace}.feed_name',
                    date=f'{self.bootstrap_namespace}.report_date',
                    archive_path=f'{self.bootstrap_namespace}.s3_archive_path',
                    credentials_path=(
                        f'{self.bootstrap_namespace}.credentials_path'),
                    api_service_name=(
                        f'{self.bootstrap_namespace}.api_service_name'),
                    api_version=(
                        f'{self.bootstrap_namespace}.api_version'),
                    jobs_meta_path=(
                        f'{self.bootstrap_namespace}.jobs_meta_path'),
                    cms_dict=(
                        f'{self.bootstrap_namespace}.cms_dict'),
                    gz=param.StaticParam(False))))


class QueryTagRunnerMixin:
    """Runner mixin that propagates query_tag to Snowflake.

    Save query_tag value at task level, so later it will be
    used to add tag to Snowflake queries.
    (see connector.connect patch in feed_ingestion/__init__.py).
    """

    def requirements(self, context):
        """Require 'backfill', when present, so it reaches every activity.

        Activities normally only receive the context keys their tasks
        declare via `.fill()`. If the workflow was started with a
        'backfill' input param, include it for every activity so it's
        available in execute() below, without editing each flow's tasks.
        """
        requirements = super().requirements(context)
        if 'backfill' in context:
            requirements = requirements | {'backfill'}
        return requirements

    def execute(self, activity, context):
        """Execute the activity with the query tag set for each task."""
        # Capture the pristine, unwrapped tasks once. self.tasks is always
        # restored to this after execute() returns (see finally blocks
        # below), so it never accumulates wrapping from a previous call.
        original_tasks = getattr(self, '_original_tasks', None)
        if original_tasks is None:
            original_tasks = self.tasks
            self._original_tasks = original_tasks

        tag_value = {
            'workflow_id': context.get('execution.workflow_id'),
            'run_id': context.get('execution.run_id'),
            'env': os.environ.get('Environment')
        }
        is_backfill = context.get('backfill', 'False')

        # add tag only if backfill is True, to avoid adding it to all queries
        if is_backfill != 'True':
            self.tasks = original_tasks
            return super().execute(activity, context)

        tag_value['backfill'] = True
        query_tag = json.dumps(tag_value)

        def tagged(task_fn):
            @wraps(task_fn)
            def wrapper(task_context, **kwargs):
                activity.logger.debug(
                    'Wrapping runner to set query tag in thread to: '
                    f'{query_tag}')
                set_query_tag(query_tag)
                try:
                    return task_fn(task_context, **kwargs)
                finally:
                    clear_query_tag()
            return wrapper

        self.tasks = tuple(tagged(task_fn) for task_fn in original_tasks)
        try:
            return super().execute(activity, context)
        finally:
            # Always leave self.tasks pointing at the unwrapped originals so
            # a subsequent call never wraps already-wrapped tasks and never
            # observes a stale tag, regardless of call order.
            self.tasks = original_tasks


class SyncRunner(QueryTagRunnerMixin, SentryRunnerMixin, runner.Sync):
    """Sync Runner with Sentry support and query tag propagation."""

    pass


class AsyncRunner(QueryTagRunnerMixin, SentryRunnerMixin, runner.Async):
    """Async Runner with Sentry support and query tag propagation."""

    pass


class FlowLicensor(FlowBase):
    """Flow class can be used to generate workflow_id using licensor.

    Assumes that there is licensor in context.
    """

    def workflow_id(self, initial_context):
        """Generate workflow id.

        Args:
            initial_context (dict): The initial context for the flow.
        Returns:
            str: A unique identifier for a workflow being executed.
        """
        licensor = initial_context['licensor']
        flow_name = '_'.join([self.name, licensor])

        if 'context_date' not in initial_context:
            date = datetime.datetime.today().strftime('%Y-%m-%d')
        else:
            date = initial_context['context_date']

        return '{flow_name}-{date}'.format(
            flow_name=flow_name,
            date=date)

    def contextified_feed_name(self, context):
        """Get feed_name in context.

        Args:
            context (dict): The context of the flow.
        Returns:
            str: Contextified feed name.
        """
        assert 'licensor' in context, 'There is no licensor in context'
        return '_'.join([self.feed_name, context['licensor']])


def generate_feed_name(feed_name):
    """Generate workflow name.

    To play nice with non-class based Flows.

    Args:
        feed_name (str): Name of the feed the work flow is processing,
            must be lower case and have spaces replaced with underscores.
            ex. 'qq'
    Returns:
        str: SWF name in the format of '<feed_name>_feed_ingestion'.
    """
    assert feed_name.islower() and ' ' not in feed_name, (
        'feed name must be lower case with no spaces')

    return '{feed_name}_feed_ingestion'.format(feed_name=feed_name)


def get_aws_config():
    """Get AWS credentials.

    Returns:
        dict: AWS credentials - {access_key: ABC, access_secret: XYZ}.
    """
    env = os.environ.get('Environment')
    if env == 'dev':
        dev_key_id = os.environ.get('AWS_WORKFLOW_ACCESS_KEY_ID')
        dev_secret = os.environ.get('AWS_WORKFLOW_SECRET_ACCESS_KEY')
        dev_token = os.environ.get('AWS_WORKFLOW_SESSION_TOKEN')
        if dev_key_id and dev_secret:
            return dict(
                access_key=dev_key_id,
                access_secret=dev_secret,
                access_token=dev_token or '',
            )

    credentials = boto3.Session().get_credentials()
    aws_config = dict(
        access_key=credentials.access_key if credentials else '',
        access_secret=credentials.secret_key if credentials else '',
        access_token=credentials.token or '' if credentials else '',
    )
    return aws_config


def get_domain():
    """Get workflow domain.

    To play nice with non-class based Flows.

    Returns:
        str: SWF domain the workflow should use.
    """
    environment = conf.getconf('env')['env']
    if environment == 'prod':
        return PROD_DOMAIN
    if environment == 'qa':
        return QA_DOMAIN
    if environment == 'test':
        return TEST_DOMAIN
    if environment == 'shadow':
        return SHADOW_DOMAIN
    return os.environ.get('SWF_DOMAIN', DEV_DOMAIN)