from typing import Any, List, Optional, Type

from snowflake import snowpark
from snowflake.ml._internal import telemetry
from snowflake.ml._internal.lineage import lineage_utils
from snowflake.ml.data import data_connector, data_ingestor, data_source, ingestor_utils
from snowflake.ml.fileset import snowfs

_PROJECT = "Dataset"
_SUBPROJECT = "DatasetReader"


class DatasetReader(data_connector.DataConnector):
    """Snowflake Dataset abstraction which provides application integration connectors"""

    @telemetry.send_api_usage_telemetry(project=_PROJECT, subproject=_SUBPROJECT)
    def __init__(
        self,
        ingestor: data_ingestor.DataIngestor,
        *,
        snowpark_session: snowpark.Session,
    ) -> None:
        super().__init__(ingestor)

        self._session: snowpark.Session = snowpark_session
        self._fs: snowfs.SnowFileSystem = ingestor_utils.get_dataset_filesystem(self._session)
        self._files: Optional[List[str]] = None

    @classmethod
    def from_dataframe(
        cls, df: snowpark.DataFrame, ingestor_class: Optional[Type[data_ingestor.DataIngestor]] = None, **kwargs: Any
    ) -> "DatasetReader":
        # Block superclass constructor from Snowpark DataFrames
        raise RuntimeError("Creating DatasetReader from DataFrames not supported")

    def _list_files(self) -> List[str]:
        """Private helper function that lists all files in this DatasetVersion and caches the results."""
        if self._files:
            return self._files

        files: List[str] = []
        for source in self.data_sources:
            assert isinstance(source, data_source.DatasetInfo)
            files.extend(ingestor_utils.get_dataset_files(self._session, source, filesystem=self._fs))
        files.sort()

        self._files = files
        return self._files

    @telemetry.send_api_usage_telemetry(project=_PROJECT, subproject=_SUBPROJECT)
    def files(self) -> List[str]:
        """Get the list of remote file paths for the current DatasetVersion.

        The file paths follows the snow protocol.

        Returns:
            A list of remote file paths

        Example:
        >>> dsv.files()
        ----
        ["snow://dataset/mydb.myschema.mydataset/versions/test/data_0_0_0.snappy.parquet",
         "snow://dataset/mydb.myschema.mydataset/versions/test/data_0_0_1.snappy.parquet"]
        """
        files = self._list_files()
        return [self._fs.unstrip_protocol(f) for f in files]

    @telemetry.send_api_usage_telemetry(project=_PROJECT, subproject=_SUBPROJECT)
    def filesystem(self) -> snowfs.SnowFileSystem:
        """Return an fsspec FileSystem which can be used to load the DatasetVersion's `files()`"""
        return self._fs

    @telemetry.send_api_usage_telemetry(
        project=_PROJECT,
        subproject=_SUBPROJECT,
        func_params_to_log=["only_feature_cols"],
    )
    def to_snowpark_dataframe(self, only_feature_cols: bool = False) -> snowpark.DataFrame:
        """Convert the DatasetVersion to a Snowpark DataFrame.

        Args:
            only_feature_cols: If True, drops exclude_cols and label_cols from returned DataFrame.
                The original DatasetVersion is unaffected.

        Returns:
            A Snowpark dataframe that contains the data of this DatasetVersion.

        Note: The dataframe generated by this method might not have the same schema as the original one. Specifically,
            - NUMBER type with scale != 0 will become float.
            - Unsupported types (see comments of :func:`Dataset.create_version`) will not have any guarantee.
                For example, an OBJECT column may be scanned back as a STRING column.
        """
        file_path_pattern = ".*data_.*[.]parquet"
        dfs: List[snowpark.DataFrame] = []
        for source in self.data_sources:
            assert isinstance(source, data_source.DatasetInfo) and source.url is not None
            df = self._session.read.option("pattern", file_path_pattern).parquet(source.url)
            if only_feature_cols and source.exclude_cols:
                df = df.drop(source.exclude_cols)
            dfs.append(df)

        combined_df = dfs[0]
        for df in dfs[1:]:
            combined_df = combined_df.union_all_by_name(df)
        return lineage_utils.patch_dataframe(combined_df, data_sources=self.data_sources, inplace=True)
