from typing import Iterable, Any, List
from .base import Source, Transformer, Sink
from ..common.logger import get_logger
from ..common.metrics import Metrics
from ..common.utils import batched


class Pipeline:
    def __init__(self, source: Source, transformers: List[Transformer], sink: Sink, batch_size: int = 500) -> None:
        self.source = source
        self.transformers = transformers
        self.sink = sink
        self.batch_size = batch_size
        self.logger = get_logger("drip_etl.pipeline")
        self.metrics = Metrics()

    def run(self) -> None:
        self.logger.info("pipeline_start")
        self.source.start()
        try:
            for batch in batched(self.source.read(), self.batch_size):
                out = batch
                for t in self.transformers:
                    out = t.transform_batch(out)
                if out:
                    self.sink.write_batch(out)
                    self.metrics.inc("batches_written", 1)
                    self.metrics.inc("records_written", len(out))
            self.sink.flush()
        finally:
            self.source.stop()
            self.sink.close()
        self.logger.info("pipeline_done")
