from types import SimpleNamespace
from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, Mapping

import fsspec
import logging

# NEW: PyArrow FileSystem imports
import pyarrow.fs as pafs

# ------------------ Lazy Accessors ------------------


class _SubdirAccessor:
    """
    Lazy accessor for a specific depot's subdirectories.
    Paths are computed on-demand using the parent manager.
    """

    __slots__ = ("_mgr", "_depot", "_subdirs")

    def __init__(self, manager, depot_name: str, subdirs: Iterable[str]):
        self._mgr = manager
        self._depot = depot_name
        self._subdirs = set(subdirs)  # Fast lookup

    def __getattr__(self, name: str) -> str:
        if name in self._subdirs:
            return self._mgr.join_paths(self._mgr.storage_path, self._depot, name)
        raise AttributeError(f"Depot '{self._depot}' has no subdirectory '{name}'")

    def __dir__(self):
        return list(self._subdirs)


class _DepotAccessor:
    """
    Lazy accessor for depots.
    """

    __slots__ = ("_mgr", "_config")

    def __init__(self, manager, depots_config: Mapping[str, Iterable[str]]):
        self._mgr = manager
        self._config = depots_config

    def __getattr__(self, name: str) -> _SubdirAccessor:
        if name in self._config:
            return _SubdirAccessor(self._mgr, name, self._config[name])
        raise AttributeError(f"No depot named '{name}'")

    def __dir__(self):
        return list(self._config.keys())


class StorageManager:
    def __init__(
        self,
        storage_path: str,
        fs_type: str = "file",
        fs_options: Optional[Mapping[str, Any]] = None,
        logger: Optional[logging.Logger] = None,
    ):
        """
        Initialises the StorageManager with the base storage path and file system settings.
        Supports generating both fsspec (standard) and PyArrow (high-performance) connections.
        """
        self.logger = logger or logging.getLogger(__name__)

        # 1. Store the raw path cleaned of trailing slashes
        self.raw_storage_path = storage_path.rstrip("/")
        self.fs_type = fs_type
        self.fs_options = fs_options or {}

        # 2. Path Normalization for fsspec
        # fsspec usually requires the protocol (e.g., s3://) to be present for consistency
        # if the user didn't provide it, we add it based on fs_type.
        if self.fs_type == "s3" and not self.raw_storage_path.startswith("s3://"):
            self.storage_path = f"s3://{self.raw_storage_path}"
        else:
            self.storage_path = self.raw_storage_path

        # 3. Instantiate the main fsspec filesystem (used for general file ops)
        try:
            self.fs = fsspec.filesystem(fs_type, **self.fs_options)
        except Exception as e:
            self.logger.error(f"Failed to initialize filesystem '{fs_type}': {e}")
            raise

        self._depot_config: Dict[str, List[str]] = {}
        self.depot_names: Any = None  # Will be _DepotAccessor

    # ------------------ NEW: PyArrow Integration ------------------

    def get_pyarrow_fs(self) -> Tuple[pafs.FileSystem, str]:
        """
        Returns a tuple of (NativePyArrowFilesystem, StrippedPath).

        This generates a high-performance C++ filesystem instance compatible with
        PyArrow Datasets, handling credential mapping and path cleaning automatically.
        """

        # CASE A: Local Filesystem
        if self.fs_type == "file":
            # Native PyArrow LocalFileSystem
            arrow_fs = pafs.LocalFileSystem()
            # Clean protocol if present (PyArrow local fs handles paths generally, but cleaner is better)
            clean_path = self.raw_storage_path.replace("file://", "")
            return arrow_fs, clean_path

        # CASE B: S3 Filesystem (Optimized C++)
        elif self.fs_type == "s3":
            # Map fsspec options (boto3 style) to PyArrow options (C++ SDK style)
            arrow_kwargs = {
                "access_key": self.fs_options.get("key"),
                "secret_key": self.fs_options.get("secret"),
                "session_token": self.fs_options.get("token"),
                "region": self.fs_options.get("client_kwargs", {}).get("region_name")
                or "us-east-1",
            }

            # Handle Custom Endpoints (MinIO, LocalStack, R2, etc.)
            client_kwargs = self.fs_options.get("client_kwargs", {})
            endpoint = client_kwargs.get("endpoint_url")

            if endpoint:
                arrow_kwargs["endpoint_override"] = endpoint
                # PyArrow S3 usually infers scheme, but explicit is safer for custom endpoints
                arrow_kwargs["scheme"] = "https" if "https" in endpoint else "http"

            try:
                arrow_fs = pafs.S3FileSystem(**arrow_kwargs)
            except Exception as e:
                self.logger.warning(
                    f"Could not create Native S3FileSystem: {e}. Falling back to wrapper."
                )
                return self._get_fallback_arrow_fs()

            # CRITICAL: PyArrow S3FS crashes if path starts with "s3://"
            # We strip the protocol to return a clean bucket/key path
            clean_path = self.raw_storage_path.replace("s3://", "")
            return arrow_fs, clean_path

        # CASE C: Fallback (GCS, Azure, or unknown types)
        # We wrap the existing fsspec instance so it "looks like" a PyArrow FS
        else:
            return self._get_fallback_arrow_fs()

    def _get_fallback_arrow_fs(self) -> Tuple[pafs.FileSystem, str]:
        """
        Wraps the existing fsspec instance in a PyArrow interface.
        It is slower (Python overhead) but ensures compatibility for unsupported fs_types.
        """
        try:
            handler = pafs.FSSpecHandler(self.fs)
            arrow_fs = pafs.PyFileSystem(handler)
            return arrow_fs, self.storage_path
        except ImportError:
            self.logger.error(
                "PyArrow is installed but FSSpecHandler is missing. Update pyarrow."
            )
            raise

    # ------------------ Existing Functionality ------------------

    @staticmethod
    def join_paths(*parts: str) -> str:
        """
        Joins paths using '/' as a separator, compatible with remote file systems.
        Prevents directory traversal (e.g. '..') in non-base parts.
        """
        clean_parts = []
        for i, part in enumerate(parts):
            # Check for path traversal in non-root parts
            if i > 0 and (".." in part or part.startswith("/")):
                 # We log and raise to ensure developers don't accidentally leak paths
                 raise ValueError(f"Security Risk: Path traversal or absolute path detected in join part: {part}")

            if i == 0:
                # Don't strip leading slashes from the first part (preserves /absolute/path or s3://)
                clean_parts.append(part.rstrip("/"))
            else:
                clean_parts.append(part.strip("/"))

        return "/".join(clean_parts)

    def setup_directories(
        self,
        base_path: str,
        dirs_to_create: Iterable[str],
        clear_existing: bool = False,
    ) -> None:
        """
        Creates directories using fsspec.
        Safely handles clearing existing directories by targeting specific subdirs.
        """
        self.logger.debug(f"Setting up directories under: {base_path}")

        # Ensure the base path exists
        self.fs.mkdirs(base_path, exist_ok=True)

        for sub_directory in dirs_to_create:
            sub_path = self.join_paths(base_path, sub_directory)

            if clear_existing and self.fs.exists(sub_path):
                self.logger.warning(f"Clearing existing directory: {sub_path}")
                try:
                    self.fs.rm(sub_path, recursive=True)
                except Exception as e:
                    self.logger.warning(f"Failed to remove {sub_path}: {e}")

            if self.logger.isEnabledFor(logging.DEBUG):
                self.logger.debug(f"Creating directory: {sub_path}")

            self.fs.mkdirs(sub_path, exist_ok=True)

    def rebuild_depot_paths(
        self,
        depots: Mapping[str, Iterable[str]],
        clear_existing: bool = False,
        write_mode: Literal["full-access", "read-only"] = "full-access",
    ) -> Any:
        """
        Rebuilds depot structures using Lazy Accessor.
        """
        for depot, sub_directories in depots.items():
            depot_path = self.join_paths(self.storage_path, depot)

            if self.logger.isEnabledFor(logging.DEBUG):
                self.logger.debug(f"Rebuilding depot at: {depot_path}")

            if write_mode == "full-access":
                self.setup_directories(
                    depot_path, sub_directories, clear_existing=clear_existing
                )

        # Store config for lazy access
        self._depot_config = depots

        # Initialize Lazy Accessor
        self.depot_names = _DepotAccessor(self, depots)

        return self.depot_names

    def get_fs_instance(self) -> Any:
        """
        Returns the standard fsspec filesystem instance.
        """
        return self.fs

    def upload_file(self, local_path: str, remote_path: str) -> None:
        """
        Upload a file to the storage via fsspec.
        """
        self.logger.debug(f"Uploading {local_path} -> {remote_path}")
        self.fs.put(local_path, remote_path)

    def download_file(self, remote_path: str, local_path: str) -> None:
        """
        Download a file from the storage via fsspec.
        """
        self.logger.debug(f"Downloading {remote_path} -> {local_path}")
        self.fs.get(remote_path, local_path)
