from __future__ import absolute_import, print_function, unicode_literals

import google.cloud.dataflow as beam

from dojo.dataset import Dataset

from .storage import BeamLocalFileStore, BeamGoogleStorageStore, BeamS3Store
from .transform import Conform, Validate, ConvertFrom, ConvertTo
from .runners.direct import DirectBeamRunner
from .runners.dataflow import DataflowBeamRunner


class BeamDataset(Dataset):

    RUNNERS = {
        'direct': DirectBeamRunner,
        'cloud': DataflowBeamRunner
    }

    STORES = {
        'file': BeamLocalFileStore,
        'gs': BeamGoogleStorageStore,
        's3': BeamS3Store
    }

    def run(self, pipeline):
        self.pipeline = pipeline
        inputs = self.inputs()
        rows = self.process(inputs)
        self.output(rows)
        return rows

    def read_input(self, rows, dataset_name):
        schema = self.input_schema(dataset_name)
        return (rows | 'from_json ' + dataset_name >> beam.ParDo(ConvertFrom())
                     | 'validate ' + dataset_name >> beam.ParDo(Validate(schema)))

    def write_output(self, rows, format):
        # The storage class is expected to do the ConvertTo transform.
        return (rows | 'conform ' >> beam.ParDo(Conform(self.OUTPUT))
                     | 'validate ' >> beam.ParDo(Validate(self.OUTPUT)))
