import logging

from apache_beam import Pipeline
import bigflow
from apache_beam.io import WriteToText
from bigflow.dataflow import BeamJob

from .pipeline import dataflow_pipeline_options, workflow_config
from .processing import count_words


logger = logging.getLogger(__name__)


def wordcount_entry_point(pipeline: Pipeline, context: bigflow.JobContext, entry_point_arguments: dict):
    logger.info(f'Running wordcount at {context.runtime_str}')
    count_words(pipeline, WriteToText("gs://{}/beam_wordcount".format(entry_point_arguments['temp_location'])))


wordcount_workflow = bigflow.Workflow(
    workflow_id="wordcount",
    log_config={
        'gcp_project_id': workflow_config['gcp_project_id'],
        'log_level': 'INFO',
    },
    definition=[BeamJob(
        id='wordcount_job',
        entry_point=wordcount_entry_point,
        pipeline_options=dataflow_pipeline_options(),
        entry_point_arguments={'temp_location': workflow_config['temp_location']}
    )])
