from typing import Union, Iterable, Iterator, Dict, Any, Tuple, TextIO, List
from pathlib import Path

from atomingest.normalization.core import NormalizationReport, NormalizationConfig
from atomingest.normalization.steps import BaseNormalizerStep, StepFactory
from atomingest.normalization.sources import DataSourceFactory


class NormalizationPipeline:
    """
    Orchestrates the normalization process by applying a series of normalization steps
    to an input data source.
    """

    def __init__(
        self,
        config: NormalizationConfig,
    ):
        self.config = config

    def run(
        self, source: Union[str, Path, TextIO, Iterable[str]]
    ) -> Tuple[Iterator[Dict[str, Any]], NormalizationReport]:
        """
        Executes the pipeline
        1. LineSource -> 2. Line Steps -> 3. CSV Parse -> 4. Row Steps
        """
        report = NormalizationReport()

        active_steps: List[BaseNormalizerStep] = []
        for step_cfg in self.config.steps:
            if step_cfg.enabled:
                step = StepFactory.create_step(step_cfg, report, self.config.strict)
                active_steps.append(step)

        try:
            # Default to CSV if source_type not specified
            source_type = getattr(self.config, "source_type", "csv")
            source_config = getattr(self.config, "source_config", {})
            data_source = DataSourceFactory.create(source_type, source_config)
        except Exception as e:
            if self.config.strict:
                raise e
            report.add_warning(f"Datasource initialization failed: {e}")
            return iter([]), report
        try:
            row_iterator = data_source.read(
                source, line_steps=active_steps, report=report
            )
        except Exception as e:
            if self.config.strict:
                raise e
            report.add_warning(f"Datasource read failed: {e}")
            return iter([]), report

        processed_rows = self._apply_row_steps(row_iterator, active_steps, report)

        # Wrap the iterator to validate row count after all rows are consumed
        validated_rows = self._validate_after_consumption(
            processed_rows, report, self.config.strict
        )

        return validated_rows, report

    def _apply_row_steps(
        self,
        rows: Iterator[Dict[str, Any]],
        steps: List[BaseNormalizerStep],
        report: NormalizationReport,
    ) -> Iterator[Dict[str, Any]]:
        for row in rows:
            current_row = row
            should_drop = False

            for step in steps:
                try:
                    current_row = step.process_row(current_row)
                    if current_row is None:
                        should_drop = True
                        break
                except Exception as e:
                    if self.config.strict:
                        raise e
                    report.add_warning(f"Error in step {step.__class__.__name__}: {e}")
                    should_drop = True
                    break
            if not should_drop:
                report.increment_rows_emitted()
                yield current_row

    def _validate_after_consumption(
        self,
        rows: Iterator[Dict[str, Any]],
        report: NormalizationReport,
        strict: bool,
    ) -> Iterator[Dict[str, Any]]:
        """Yield all rows, then validate row count after consumption."""
        yield from rows
        # After all rows are consumed, validate the row count
        report.validate_row_count(strict=strict)
