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

try:
    import pandas as pd
except ImportError:
    pd = None

from atomingest.normalization.steps import BaseNormalizerStep
from atomingest.normalization.sources.base import DataSource

logger = logging.getLogger(__name__)


class ExcelDataSource(DataSource):
    """
    Concrete Strategy for Excel files (XLSX, XLS)
    """

    def __init__(self, sheet_name: Union[str, int] = 0, **kwargs):
        self.sheet_name = sheet_name
        self.read_excel_kwargs = kwargs

    def read(
        self,
        source: Union[str, Path, TextIO, Iterable[str]],
        line_steps: List[BaseNormalizerStep] = None,
        report=None,
    ) -> Iterator[Dict[str, Any]]:
        if pd is None:
            raise ImportError("pandas is required to read Excel files")

        if line_steps:
            logger.warning(
                "Line steps are not supported for ExcelDataSource and will be ignored."
            )

        df = pd.read_excel(source, sheet_name=self.sheet_name, **self.read_excel_kwargs)
        df = df.where(pd.notnull(df), None)

        # Count header row
        if report:
            report.increment_total_lines()

        for row in df.to_dict(orient="records"):
            if report:
                report.increment_total_lines()
            yield row
