import collections
import logging
from functools import partial
from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, cast

import fsspec

from snowflake import snowpark
from snowflake.connector import connection
from snowflake.ml._internal import telemetry
from snowflake.ml._internal.exceptions import (
    error_codes,
    exceptions as snowml_exceptions,
)
from snowflake.ml._internal.utils import identifier
from snowflake.ml.fileset import stage_fs
from snowflake.ml.utils import connection_params

PROTOCOL_NAME = "sfc"

_SFFilePath = collections.namedtuple("_SFFilePath", ["database", "schema", "stage", "filepath"])
_PROJECT = "FileSet"

_FILESYSTEM_KWARGS_KEY = "kwargs"
_RECREATE_FROM_SERIALIZED = "recreate_from_serialized"


class SFFileSystem(fsspec.AbstractFileSystem):
    """A filesystem that allows user to access Snowflake stages and stage files with valid Snowflake locations.

    The file system is is based on fsspec (https://filesystem-spec.readthedocs.io/). It is a file system wrapper
    built on top of SFStageFileSystem. It takes Snowflake stage file path as the input and supports read operation.
    A valid Snowflake location will have the form "@{database_name}.{schema_name}.{stage_name}/{path_to_file}".

    Example 1: Create a file system object and do file operation
    --------
    >>> conn = snowflake.connector.connect(**connection_parameters)
    >>> sffs = SFFileSystem(sf_connection=conn)
    >>> sffs.ls("@MYDB.public.FOO/nytrain")
    ['@MYDB.public.FOO/nytrain/data_0_0_0.csv', '@MYDB.public.FOO/nytrain/data_0_0_1.csv']
    >>> with sffs.open('@MYDB.public.FOO/nytrain/nytrain/data_0_0_1.csv', mode='rb') as f:
    >>>     print(f.readline())
    b'2014-02-05 14:35:00.00000054,13,2014-02-05 14:35:00 UTC,-74.00688,40.73049,-74.00563,40.70676,2\n'

    Example 2: Use SFC protocol to create a file system object via fsspec
    --------
    >>> conn = snowflake.connector.connect(**connection_parameters)
    >>> sffs = fsspec.filesystem("sfc", sf_connection=conn)
    >>> sffs.ls("@MYDB.public.FOO/nytrain")
    ['@MYDB.public.FOO/nytrain/data_0_0_0.csv', '@MYDB.public.FOO/nytrain/data_0_0_1.csv']
    >>> with sffs.open('@MYDB.public.FOO/nytrain/nytrain/data_0_0_1.csv', mode='rb') as f:
    >>>     print(f.readline())
    b'2014-02-05 14:35:00.00000054,13,2014-02-05 14:35:00 UTC,-74.00688,40.73049,-74.00563,40.70676,2\n'

    Example 3: Use SFC protocol and open files via fsspec without creating a file system object explicitly
    --------
    >>> conn = snowflake.connector.connect(**connection_parameters)
    >>> with fsspec.open("sfc://@MYDB.public.FOO/nytrain/data_0_0_1.csv", mode='rb', sf_connection=conn) as f:
    >>>     print(f.readline())
    b'2014-02-05 14:35:00.00000054,13,2014-02-05 14:35:00 UTC,-74.00688,40.73049,-74.00563,40.70676,2\n'
    """

    def __init__(
        self,
        sf_connection: Optional[connection.SnowflakeConnection] = None,
        snowpark_session: Optional[snowpark.Session] = None,
        **kwargs: Any,
    ) -> None:
        """Initialize file system with a Snowflake Python connection.

        Args:
            sf_connection: A Snowflake python connection object. Either it or snowpark_session must be non-empty.
            snowpark_session: A Snowpark session. Either it or sf_connection must be non-empty.
            kwargs : Optional. Other parameters that can be passed on to fsspec. Currently supports:
                - skip_instance_cache: Int. Controls reuse of instances.
                - cache_type, cache_options, block_size: Configure file buffering.
                See more information of these options in https://filesystem-spec.readthedocs.io/en/latest/features.html

        Raises:
            ValueError: An error occurred when not exactly one of sf_connection and snowpark_session is given.
            SnowflakeMLException: A failure was encountered while recreating the SFFileSystem from a serialized state.
        """
        if kwargs.get(_RECREATE_FROM_SERIALIZED):
            try:
                snowpark_session = self._create_default_session()
            except Exception as e:
                raise snowml_exceptions.SnowflakeMLException(
                    error_code=error_codes.SNOWML_DESERIALIZATION_FAILED,
                    original_exception=ValueError("Unable to deserialize SFFileSystem."),
                ) from e

        if sf_connection:
            self._conn = sf_connection
        elif snowpark_session:
            self._conn = snowpark_session._conn._conn
        else:
            raise ValueError("Either sf_connection or snowpark_session has to be non-empty!")
        self._kwargs = kwargs
        self._stage_fs_set: Dict[Tuple[str, str, str], stage_fs.SFStageFileSystem] = {}

        super().__init__(**kwargs)

    def _create_default_session(self) -> snowpark.Session:
        """Create a Snowpark Session from default login options.

        Returns:
            An active Snowpark Session.

        Raises:
            ValueError: Snowflake Login Options could not be retrieved from default locations.
            ValueError: Snowflake Connection could not be created.

        """
        try:
            snowflake_config = connection_params.SnowflakeLoginOptions()
        except Exception as e:
            raise ValueError("Unable to retrieve Snowflake Login Options.") from e

        try:
            session = snowpark.Session.builder.configs(snowflake_config).create()
        except Exception as e:
            raise ValueError("Unable to create Snowflake connection.") from e

        assert isinstance(session, snowpark.Session)
        return session

    def __reduce__(self) -> Tuple[Callable[[], Type["SFFileSystem"]], Tuple[()], Dict[str, Any]]:
        """Returns a state dictionary for use in serialization.

        Returns:
            A tuple that is used for recreating the SFFileSystem. For more information, refer to
            https://docs.python.org/3/library/pickle.html#object.__reduce__
            A `partial` is used to generate a callable that accepts kwargs.
        """
        state_dictionary = {_FILESYSTEM_KWARGS_KEY: self._kwargs}

        return partial(self.__class__, **{_RECREATE_FROM_SERIALIZED: True}), (), state_dictionary

    def __setstate__(self, state_dict: Dict[str, Any]) -> None:
        """Sets the dictionary state at deserialization time, and rebuilds a snowflake connection.

        Args:
            state_dict: State dictionary saved at serialization time.

        Raises:
            KeyError: The Kwargs key is not present in the state dictionary.
            ValueError: The value corresponding to the kwargs key is not a dictionary.

        """
        if _FILESYSTEM_KWARGS_KEY not in state_dict:
            raise KeyError(f"Serialized state dictionary missing key {_FILESYSTEM_KWARGS_KEY}.")

        kwargs_dict = state_dict.get(_FILESYSTEM_KWARGS_KEY)
        if not isinstance(kwargs_dict, dict):
            raise ValueError(f"The value corresponding to {_FILESYSTEM_KWARGS_KEY} is not a dictionary.")

        self._kwargs = kwargs_dict

    def _get_stage_fs(self, sf_file_path: _SFFilePath) -> stage_fs.SFStageFileSystem:
        """Get the stage file system for the given snowflake location.

        Args:
            sf_file_path: A SFFILEPATH namedtuple

        Returns:
            A SFStageFileSystem object which supports readonly file operations on Snowflake stages.
        """
        stage_fs_key = (sf_file_path.database, sf_file_path.schema, sf_file_path.stage)
        if stage_fs_key not in self._stage_fs_set:
            cnt_stage_fs = stage_fs.SFStageFileSystem(
                sf_connection=self._conn,
                db=sf_file_path.database,
                schema=sf_file_path.schema,
                stage=sf_file_path.stage,
                **self._kwargs,
            )
            self._stage_fs_set[stage_fs_key] = cnt_stage_fs
        return self._stage_fs_set[stage_fs_key]

    @telemetry.send_api_usage_telemetry(
        project=_PROJECT,
        func_params_to_log=["detail"],
        conn_attr_name="_conn",
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def ls(self, path: str, detail: bool = False, **kwargs: Any) -> Union[List[str], List[Dict[str, Any]]]:
        """Override fsspec `ls` method. List single "directory" with or without details.

        Args:
            path : location at which to list files.
                It should be in the format of "@{database}.{schema}.{stage}/{path}"
            detail : if True, each list item is a dict of file properties; otherwise, returns list of filenames.
            kwargs : additional arguments passed on.

        Returns:
            A list of filename if `detail` is false, or a list of dict if `detail` is true.

        Example:
        >>> sffs.ls("@MYDB.public.FOO/")
        ['@MYDB.public.FOO/nytrain']
        >>> sffs.ls("@MYDB.public.FOO/nytrain")
        ['@MYDB.public.FOO/nytrain/data_0_0_0.csv', '@MYDB.public.FOO/nytrain/data_0_0_1.csv']
        >>> sffs.ls("@MYDB.public.FOO/nytrain/")
        ['@MYDB.public.FOO/nytrain/data_0_0_0.csv', '@MYDB.public.FOO/nytrain/data_0_0_1.csv']
        """
        file_path = _parse_sfc_file_path(path)
        stage_fs = self._get_stage_fs(file_path)
        stage_path_list = stage_fs.ls(file_path.filepath, detail=True, **kwargs)
        stage_path_list = cast(List[Dict[str, Any]], stage_path_list)
        return self._decorate_ls_res(stage_fs, stage_path_list, detail)

    @telemetry.send_api_usage_telemetry(
        project=_PROJECT,
        conn_attr_name="_conn",
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def optimize_read(self, files: Optional[List[str]] = None) -> None:
        """Prefetch and cache the presigned urls for all the given files to speed up the file opening.

        All the files introduced here will have their urls cached. Further open() on any of cached urls will lead to a
        batch refreshment of the cached urls in the same stage if that url is inactive.

        Args:
            files: A list of file paths that needs their presigned url cached.
        """
        if not files:
            return
        stage_file_paths: Dict[Tuple[str, str, str], List[str]] = collections.defaultdict(list)
        for file in files:
            file_path = _parse_sfc_file_path(file)
            stage_file_paths[(file_path.database, file_path.schema, file_path.stage)].append(file_path.filepath)
        for k, v in stage_file_paths.items():
            stage_fs = self._get_stage_fs(_SFFilePath(k[0], k[1], k[2], "*"))
            stage_fs.optimize_read(v)

    @telemetry.send_api_usage_telemetry(
        project=_PROJECT,
        conn_attr_name="_conn",
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def _open(self, path: str, **kwargs: Any) -> fsspec.spec.AbstractBufferedFile:
        """Override fsspec `_open` method. Open a file for reading in 'rb' mode.

        The opened file will be readable for 4 hours. After that you will need to reopen it.

        Fsspec `open` API is built on `_open` method and supports non-binary mode.
        See more details in https://github.com/fsspec/filesystem_spec/blob/2022.10.0/fsspec/spec.py#L1019

        Args:
            path: Path of file in Snowflake stage.
                It should be in the format of "@{database}.{schema}.{stage}/{path}"
            **kwargs : Additional arguments passed on. See more in
                https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.open

        Returns:
            A fsspec AbstractBufferedFile which supports python file operations.
        """
        file_path = _parse_sfc_file_path(path)
        stage_fs = self._get_stage_fs(file_path)
        return stage_fs._open(file_path.filepath, **kwargs)

    @telemetry.send_api_usage_telemetry(
        project=_PROJECT,
        conn_attr_name="_conn",
    )
    @snowpark._internal.utils.private_preview(version="0.2.0")
    def info(self, path: str, **kwargs: Any) -> Dict[str, Any]:
        """Override fsspec `info` method. Give details of entry at path."""
        file_path = _parse_sfc_file_path(path)
        stage_fs = self._get_stage_fs(file_path)
        res: Dict[str, Any] = stage_fs.info(file_path.filepath, **kwargs)
        if res:
            res["name"] = self._stage_path_to_absolute_path(stage_fs, res["name"])
        return res

    def _decorate_ls_res(
        self,
        stage_fs: stage_fs.SFStageFileSystem,
        stage_path_list: List[Dict[str, Any]],
        detail: bool,
    ) -> Union[List[str], List[Dict[str, Any]]]:
        """Add the stage location as the prefix of file names returned by ls() of stagefs"""
        for path in stage_path_list:
            path["name"] = self._stage_path_to_absolute_path(stage_fs, path["name"])
        if detail:
            return stage_path_list
        else:
            return [path["name"] for path in stage_path_list]

    def _stage_path_to_absolute_path(self, stage_fs: stage_fs.SFStageFileSystem, path: str) -> str:
        """Convert the relative path in a stage to an absolute path starts with the location of the stage."""
        return stage_fs.stage_name + "/" + path


fsspec.register_implementation(PROTOCOL_NAME, SFFileSystem)


def _parse_sfc_file_path(path: str) -> _SFFilePath:
    """Parse a snowflake location path.

    The following propertis will be extracted from the path input:
    - database
    - schema
    - stage
    - path (optional)

    Args:
        path: A string in the format of "@{database}.{schema}.{stage}/{path}".

            Example:
                "@my_db.my_schema.my_stage/"
                "@my_db.my_schema.my_stage/file1"
                "@my_db.my_schema.my_stage/dir1/"
                "@my_db.my_schema.my_stage/dir1/file2"

    Returns:
        A namedtuple consists of database name, schema name, stage name and path.

    Raises:
        SnowflakeMLException: An error occurred when invalid path is given.
    """
    sfc_prefix = f"{PROTOCOL_NAME}://"
    if path.startswith(sfc_prefix):
        path = path[len(sfc_prefix) :]
    if not path.startswith("@"):
        raise snowml_exceptions.SnowflakeMLException(
            error_code=error_codes.SNOWML_INVALID_STAGE,
            original_exception=ValueError(
                'Invalid path. Expected path to start with "@". Example: @database.schema.stage/optional_path.'
            ),
        )
    try:
        res = identifier.parse_schema_level_object_identifier(path[1:])
        logging.debug(f"Parsed path: {res}")
        return _SFFilePath(res[0], res[1], res[2], res[3][1:])
    except ValueError:
        raise snowml_exceptions.SnowflakeMLException(
            error_code=error_codes.SNOWML_INVALID_STAGE,
            original_exception=ValueError(
                f"Invalid path. Expected format: @database.schema.stage/optional_path. Getting {path}"
            ),
        )
