import csv
from typing import Iterator, Dict, Any, Union, List, TextIO, Iterable, Optional, Set
from pathlib import Path

from atomingest.normalization.io import LineSource
from atomingest.normalization.steps import BaseNormalizerStep
from atomingest.normalization.sources.base import DataSource
from atomingest.normalization.core import NormalizationReport


class CSVDataSource(DataSource):
    """
    Concrete Strategy for CSV and Text-based files.
    Supports column validation and filtering via expected_columns contract.
    """

    def __init__(
        self,
        delimiter: str = ",",
        quotechar: str = '"',
        expected_columns: Optional[List[str]] = None,
        drop_extra_columns: bool = False,
        require_columns: bool = False,
        **kwargs,
    ):
        """
        Initialize the CSV Data Source.

        Args:
            delimiter: CSV delimiter character.
            quotechar: CSV quote character.
            expected_columns: List of column names expected in the file.
            drop_extra_columns: If True, columns not in expected_columns are removed from rows.
            require_columns: If True, raises ValueError if expected columns are missing.
            **kwargs: Additional arguments passed to csv.DictReader.
        """
        self.delimiter = delimiter
        self.quotechar = quotechar
        self.expected_columns = expected_columns
        self.drop_extra_columns = drop_extra_columns
        self.require_columns = require_columns
        self.csv_kwargs = kwargs

    def read(
        self,
        source: Union[str, Path, TextIO, Iterable[str]],
        line_steps: List[BaseNormalizerStep] = None,
        report: Optional[NormalizationReport] = None,
    ) -> Iterator[Dict[str, Any]]:
        raw_lines = LineSource(source)
        line_iterator = iter(raw_lines)

        # Count lines if report is provided
        if report:
            line_iterator = self._count_lines(line_iterator, report)

        if line_steps:
            for step in line_steps:
                line_iterator = step.process_lines(line_iterator)

        try:
            reader = csv.DictReader(
                line_iterator,
                delimiter=self.delimiter,
                quotechar=self.quotechar,
                **self.csv_kwargs,
            )

            # Column Validation logic
            expected_set: Optional[Set[str]] = None
            if self.expected_columns:
                expected_set = set(self.expected_columns)

                # Accessing reader.fieldnames triggers header reading
                # It will be None if the file is empty or has no valid header
                if reader.fieldnames is None:
                    msg = "CSV file has no header or is empty"
                    if self.require_columns:
                        raise ValueError(msg)
                    elif report:
                        report.add_warning(msg)
                elif reader.fieldnames:
                    file_headers = set(reader.fieldnames)

                    # Check for missing columns
                    missing = expected_set - file_headers
                    if missing:
                        msg = f"Missing columns in CSV: {sorted(list(missing))}"
                        if self.require_columns:
                            raise ValueError(msg)
                        elif report:
                            report.add_warning(msg)

                    # Check for extra columns
                    extra = file_headers - expected_set
                    if extra and report:
                        msg = f"Extra columns found: {sorted(list(extra))}"
                        if self.drop_extra_columns:
                            report.add_warning(f"{msg}. Dropping them.")
                        else:
                            report.add_warning(msg)

            for row in reader:
                # Filter extra columns if configured
                if expected_set and self.drop_extra_columns:
                    # Create a new dict with only expected columns
                    row = {k: v for k, v in row.items() if k in expected_set}

                yield row

        except csv.Error as e:
            raise RuntimeError(f"Error reading CSV/File data: {e}")

    def _count_lines(
        self, lines: Iterator[str], report: NormalizationReport
    ) -> Iterator[str]:
        """Count all lines as they pass through."""
        for line in lines:
            report.increment_total_lines()
            yield line
