import re
from abc import ABC, abstractmethod

# Add imports for source customization
from typing import Any, Callable, Dict, Optional, Type, TypeVar
from urllib.parse import unquote

# Don't import TableData at the module level to avoid circular imports
# from datahub.ingestion.source.s3.source import TableData

T = TypeVar("T")


class ObjectStoreInterface(ABC):
    """
    Abstract interface for object store operations.

    This interface defines the operations that any object store connector
    (S3, GCS, ABS, etc.) should implement to provide a consistent API.
    """

    @classmethod
    @abstractmethod
    def is_uri(cls, uri: str) -> bool:
        """
        Check if the given URI is for this object store.

        Args:
            uri: The URI to check

        Returns:
            True if the URI is for this object store, False otherwise
        """
        pass

    @classmethod
    @abstractmethod
    def get_prefix(cls, uri: str) -> Optional[str]:
        """
        Get the prefix for this object store URI (e.g., 's3://', 'gs://').

        Args:
            uri: The URI to get the prefix from

        Returns:
            The prefix if the URI starts with it, None otherwise
        """
        pass

    @classmethod
    @abstractmethod
    def strip_prefix(cls, uri: str) -> str:
        """
        Remove the object store prefix from the URI.

        Args:
            uri: The URI to strip the prefix from

        Returns:
            The URI without the prefix

        Raises:
            ValueError: If the URI does not start with the expected prefix
        """
        pass

    @classmethod
    @abstractmethod
    def get_bucket_name(cls, uri: str) -> str:
        """
        Get the bucket name from the URI.

        Args:
            uri: The URI to get the bucket name from

        Returns:
            The bucket name

        Raises:
            ValueError: If the URI is not valid for this object store
        """
        pass

    @classmethod
    @abstractmethod
    def get_object_key(cls, uri: str) -> str:
        """
        Get the object key/path (excluding the bucket) from the URI.

        Args:
            uri: The URI to get the object key from

        Returns:
            The object key

        Raises:
            ValueError: If the URI is not valid for this object store
        """
        pass

    @classmethod
    def get_object_store_bucket_name(cls, uri: str) -> str:
        """
        Get the bucket name from the URI, handling foreign URIs if supported.

        The default implementation just calls get_bucket_name, but subclasses
        can override this to handle URIs from other object stores.

        Args:
            uri: The URI to get the bucket name from

        Returns:
            The bucket name

        Raises:
            ValueError: If the URI is not supported
        """
        return cls.get_bucket_name(uri)


class S3ObjectStore(ObjectStoreInterface):
    """Implementation of ObjectStoreInterface for Amazon S3."""

    PREFIXES = ["s3://", "s3n://", "s3a://"]

    @classmethod
    def is_uri(cls, uri: str) -> bool:
        return any(uri.startswith(prefix) for prefix in cls.PREFIXES)

    @classmethod
    def get_prefix(cls, uri: str) -> Optional[str]:
        for prefix in cls.PREFIXES:
            if uri.startswith(prefix):
                return prefix
        return None

    @classmethod
    def strip_prefix(cls, uri: str) -> str:
        prefix = cls.get_prefix(uri)
        if not prefix:
            raise ValueError(
                f"Not an S3 URI. Must start with one of the following prefixes: {str(cls.PREFIXES)}"
            )
        return uri[len(prefix) :]

    @classmethod
    def get_bucket_name(cls, uri: str) -> str:
        if not cls.is_uri(uri):
            raise ValueError(
                f"Not an S3 URI. Must start with one of the following prefixes: {str(cls.PREFIXES)}"
            )
        return cls.strip_prefix(uri).split("/")[0]

    @classmethod
    def get_object_key(cls, uri: str) -> str:
        if not cls.is_uri(uri):
            raise ValueError(
                f"Not an S3 URI. Must start with one of the following prefixes: {str(cls.PREFIXES)}"
            )
        parts = cls.strip_prefix(uri).split("/", 1)
        if len(parts) < 2:
            return ""
        return parts[1]

    @classmethod
    def get_object_store_bucket_name(cls, uri: str) -> str:
        """
        Get the bucket name from an S3 URI.

        Args:
            uri: The URI to get the bucket name from

        Returns:
            The bucket name

        Raises:
            ValueError: If the URI is not an S3 URI
        """
        return cls.get_bucket_name(uri)


class GCSObjectStore(ObjectStoreInterface):
    """Implementation of ObjectStoreInterface for Google Cloud Storage."""

    PREFIX = "gs://"

    @classmethod
    def is_uri(cls, uri: str) -> bool:
        return uri.startswith(cls.PREFIX)

    @classmethod
    def get_prefix(cls, uri: str) -> Optional[str]:
        if uri.startswith(cls.PREFIX):
            return cls.PREFIX
        return None

    @classmethod
    def strip_prefix(cls, uri: str) -> str:
        prefix = cls.get_prefix(uri)
        if not prefix:
            raise ValueError(f"Not a GCS URI. Must start with prefix: {cls.PREFIX}")
        return uri[len(prefix) :]

    @classmethod
    def get_bucket_name(cls, uri: str) -> str:
        if not cls.is_uri(uri):
            raise ValueError(f"Not a GCS URI. Must start with prefix: {cls.PREFIX}")
        return cls.strip_prefix(uri).split("/")[0]

    @classmethod
    def get_object_key(cls, uri: str) -> str:
        if not cls.is_uri(uri):
            raise ValueError(f"Not a GCS URI. Must start with prefix: {cls.PREFIX}")
        parts = cls.strip_prefix(uri).split("/", 1)
        if len(parts) < 2:
            return ""
        return parts[1]

    @classmethod
    def get_object_store_bucket_name(cls, uri: str) -> str:
        """
        Get the bucket name from a GCS URI.

        Args:
            uri: The URI to get the bucket name from

        Returns:
            The bucket name

        Raises:
            ValueError: If the URI is not a GCS URI
        """
        return cls.get_bucket_name(uri)


class ABSObjectStore(ObjectStoreInterface):
    """Implementation of ObjectStoreInterface for Azure Blob Storage."""

    PREFIX = "abfss://"
    HTTPS_REGEX = re.compile(r"(https?://[a-z0-9]{3,24}\.blob\.core\.windows\.net/)")

    @classmethod
    def is_uri(cls, uri: str) -> bool:
        return uri.startswith(cls.PREFIX) or bool(cls.HTTPS_REGEX.match(uri))

    @classmethod
    def get_prefix(cls, uri: str) -> Optional[str]:
        if uri.startswith(cls.PREFIX):
            return cls.PREFIX

        # Check for HTTPS format
        match = cls.HTTPS_REGEX.match(uri)
        if match:
            return match.group(1)

        return None

    @classmethod
    def strip_prefix(cls, uri: str) -> str:
        if uri.startswith(cls.PREFIX):
            return uri[len(cls.PREFIX) :]

        # Handle HTTPS format
        match = cls.HTTPS_REGEX.match(uri)
        if match:
            return uri[len(match.group(1)) :]

        raise ValueError(
            f"Not an ABS URI. Must start with prefix: {cls.PREFIX} or match Azure Blob Storage HTTPS pattern"
        )

    @classmethod
    def get_bucket_name(cls, uri: str) -> str:
        if not cls.is_uri(uri):
            raise ValueError(
                f"Not an ABS URI. Must start with prefix: {cls.PREFIX} or match Azure Blob Storage HTTPS pattern"
            )

        if uri.startswith(cls.PREFIX):
            # abfss://container@account.dfs.core.windows.net/path
            return cls.strip_prefix(uri).split("@")[0]
        else:
            # https://account.blob.core.windows.net/container/path
            return cls.strip_prefix(uri).split("/")[0]

    @classmethod
    def get_object_key(cls, uri: str) -> str:
        if not cls.is_uri(uri):
            raise ValueError(
                f"Not an ABS URI. Must start with prefix: {cls.PREFIX} or match Azure Blob Storage HTTPS pattern"
            )

        if uri.startswith(cls.PREFIX):
            # abfss://container@account.dfs.core.windows.net/path
            parts = cls.strip_prefix(uri).split("@", 1)
            if len(parts) < 2:
                return ""
            account_path = parts[1]
            path_parts = account_path.split("/", 1)
            if len(path_parts) < 2:
                return ""
            return path_parts[1]
        else:
            # https://account.blob.core.windows.net/container/path
            stripped = cls.strip_prefix(uri)
            parts = stripped.split("/", 1)
            if len(parts) < 2:
                return ""
            return parts[1]


# Registry of all object store implementations
OBJECT_STORE_REGISTRY: Dict[str, Type[ObjectStoreInterface]] = {
    "s3": S3ObjectStore,
    "gcs": GCSObjectStore,
    "abs": ABSObjectStore,
}


def get_object_store_for_uri(uri: str) -> Optional[Type[ObjectStoreInterface]]:
    """
    Get the appropriate object store implementation for the given URI.

    Args:
        uri: The URI to get the object store for

    Returns:
        The object store implementation, or None if no matching implementation is found
    """
    for object_store in OBJECT_STORE_REGISTRY.values():
        if object_store.is_uri(uri):
            return object_store
    return None


def get_object_store_bucket_name(uri: str) -> str:
    """
    Get the bucket name from any supported object store URI.

    This function acts as a central dispatcher that:
    1. Identifies the appropriate object store implementation for the URI
    2. Uses that implementation to extract the bucket name
    3. Falls back to specific URI format parsing if needed

    Args:
        uri: The URI to get the bucket name from

    Returns:
        The bucket name

    Raises:
        ValueError: If the URI is not supported by any registered object store
    """
    # First try to find the native implementation for this URI
    object_store = get_object_store_for_uri(uri)
    if object_store:
        return object_store.get_bucket_name(uri)

    # If no native implementation, handle specific URI formats directly
    if uri.startswith("gs://"):
        return uri[5:].split("/")[0]
    elif any(uri.startswith(prefix) for prefix in S3ObjectStore.PREFIXES):
        prefix_length = next(
            len(prefix) for prefix in S3ObjectStore.PREFIXES if uri.startswith(prefix)
        )
        return uri[prefix_length:].split("/")[0]
    elif uri.startswith(ABSObjectStore.PREFIX):
        return uri[len(ABSObjectStore.PREFIX) :].split("@")[0]
    elif ABSObjectStore.HTTPS_REGEX.match(uri):
        # Handle HTTPS Azure Blob Storage URLs
        match = ABSObjectStore.HTTPS_REGEX.match(uri)
        if match:
            stripped = uri[len(match.group(1)) :]
            return stripped.split("/")[0]

    raise ValueError(f"Unsupported URI format: {uri}")


def get_object_key(uri: str) -> str:
    """
    Get the object key from any supported object store URI.

    Args:
        uri: The URI to get the object key from

    Returns:
        The object key

    Raises:
        ValueError: If the URI is not supported by any registered object store
    """
    object_store = get_object_store_for_uri(uri)
    if object_store:
        return object_store.get_object_key(uri)

    raise ValueError(f"Unsupported URI format: {uri}")


class ObjectStoreSourceAdapter:
    """
    Adapter for customizing object store source implementations.

    This class provides a way to customize source implementations for different
    object stores (S3, GCS, etc.) without having to directly modify those classes.
    Instead, adapters register customizations that are applied to the source instance.
    """

    @staticmethod
    def create_s3_path(bucket_name: str, key: str) -> str:
        """
        Create a default S3 path.

        Args:
            bucket_name: The bucket name
            key: The object key

        Returns:
            A properly formatted S3 URI
        """
        return unquote(f"s3://{bucket_name}/{key}")

    @staticmethod
    def create_gcs_path(bucket_name: str, key: str) -> str:
        """
        Create a default GCS path.

        Args:
            bucket_name: The bucket name
            key: The object key

        Returns:
            A properly formatted GCS URI
        """
        return unquote(f"gs://{bucket_name}/{key}")

    @staticmethod
    def create_abs_path(container_name: str, key: str, account_name: str) -> str:
        """
        Create a default Azure Blob Storage path.

        Args:
            container_name: The container name
            key: The object key
            account_name: The storage account name

        Returns:
            A properly formatted ABS URI
        """
        return unquote(
            f"abfss://{container_name}@{account_name}.dfs.core.windows.net/{key}"
        )

    @staticmethod
    def get_s3_external_url(
        table_data: Any, region: Optional[str] = None
    ) -> Optional[str]:
        """
        Get the AWS S3 console URL for the given table.

        Args:
            table_data: Table data containing path information
            region: AWS region for the S3 console URL, defaults to us-east-1 if not specified

        Returns:
            The AWS console URL, or None if not applicable
        """
        if not S3ObjectStore.is_uri(table_data.table_path):
            return None

        # Get the bucket name and key from the S3 URI
        bucket_name = get_object_store_bucket_name(table_data.table_path)
        key = get_object_key(table_data.table_path)

        # Use the provided region or default to us-east-1
        aws_region = region or "us-east-1"

        return f"https://{aws_region}.console.aws.amazon.com/s3/buckets/{bucket_name}?prefix={key}"

    @staticmethod
    def get_gcs_external_url(table_data: Any) -> Optional[str]:
        """
        Get the GCS console URL for the given table.

        Args:
            table_data: Table data containing path information

        Returns:
            The GCS console URL, or None if not applicable
        """
        if not GCSObjectStore.is_uri(table_data.table_path):
            return None

        # Get the bucket name and key from the GCS URI
        bucket_name = get_object_store_bucket_name(table_data.table_path)
        key = get_object_key(table_data.table_path)

        # Return the basic GCS console URL
        return f"https://console.cloud.google.com/storage/browser/{bucket_name}/{key}"

    @staticmethod
    def get_abs_external_url(table_data: Any) -> Optional[str]:
        """
        Get the Azure Storage browser URL for the given table.

        Args:
            table_data: Table data containing path information

        Returns:
            The Azure Storage URL, or None if not applicable
        """
        if not ABSObjectStore.is_uri(table_data.table_path):
            return None

        try:
            if table_data.table_path.startswith("abfss://"):
                # URI format: abfss://container@account.dfs.core.windows.net/path
                path_without_prefix = ABSObjectStore.strip_prefix(table_data.table_path)
                parts = path_without_prefix.split("@", 1)
                if len(parts) < 2:
                    return None

                container_name = parts[0]
                account_parts = parts[1].split("/", 1)
                account_domain = account_parts[0]
                account_name = account_domain.split(".")[0]
            else:
                # Handle HTTPS format: https://account.blob.core.windows.net/container/path
                container_name = ABSObjectStore.get_bucket_name(table_data.table_path)
                if "blob.core.windows.net" in table_data.table_path:
                    account_name = table_data.table_path.split("//")[1].split(".")[0]
                else:
                    return None

            # Construct Azure portal URL
            return f"https://portal.azure.com/#blade/Microsoft_Azure_Storage/ContainerMenuBlade/overview/storageAccountId/{account_name}/containerName/{container_name}"
        except Exception:
            # If any parsing error occurs, return None
            return None

    def __init__(
        self,
        platform: str,
        platform_name: str,
        aws_region: Optional[str] = None,
        azure_storage_account: Optional[str] = None,
    ):
        """
        Initialize the adapter with platform-specific configurations.

        Args:
            platform: The platform identifier (e.g., "s3", "gcs", "abs")
            platform_name: The human-readable platform name
            aws_region: AWS region for S3 URLs, defaults to us-east-1 if not specified
            azure_storage_account: Azure Storage account name
        """
        self.platform = platform
        self.platform_name = platform_name
        self.aws_region = aws_region
        self.azure_storage_account = azure_storage_account
        self.customizations: Dict[str, Callable[..., Any]] = {}

        # Register default customizations based on platform
        if platform == "gcs":
            self.register_customization("is_s3_platform", lambda: True)
            self.register_customization("create_s3_path", self.create_gcs_path)
            self.register_customization(
                "get_external_url",
                lambda table_data: self.get_gcs_external_url(table_data),
            )
            # Fix URI mismatch issue in pattern matching
            self.register_customization(
                "_normalize_uri_for_pattern_matching",
                self._normalize_gcs_uri_for_pattern_matching,
            )
            # Fix URI handling in schema extraction - override strip_s3_prefix for GCS
            self.register_customization("strip_s3_prefix", self._strip_gcs_prefix)
        elif platform == "s3":
            self.register_customization("is_s3_platform", lambda: True)
            self.register_customization("create_s3_path", self.create_s3_path)
            self.register_customization(
                "get_external_url",
                lambda table_data: self.get_s3_external_url(
                    table_data, self.aws_region
                ),
            )
        elif platform == "abs":
            self.register_customization("is_s3_platform", lambda: True)
            # If we have an Azure storage account, create a specialized path creation function
            if self.azure_storage_account:
                storage_account = (
                    self.azure_storage_account
                )  # Create a local non-optional variable
                self.register_customization(
                    "create_s3_path",
                    lambda bucket, key: self.create_abs_path(
                        bucket, key, storage_account
                    ),
                )
            else:
                # Fall back to a simpler implementation if no account provided
                self.register_customization(
                    "create_s3_path", lambda bucket, key: f"abfss://{bucket}@{key}"
                )
            self.register_customization("get_external_url", self.get_abs_external_url)

    def register_customization(
        self, method_name: str, implementation: Callable[..., Any]
    ) -> None:
        """
        Register a customization for a specific method.

        Args:
            method_name: The name of the method to customize
            implementation: The implementation to use
        """
        self.customizations[method_name] = implementation

    def apply_customizations(self, source: Any) -> Any:
        """
        Apply all registered customizations to the source instance.

        Args:
            source: The source instance to customize

        Returns:
            The customized source instance
        """
        # Set the platform
        if hasattr(source, "source_config") and hasattr(
            source.source_config, "platform"
        ):
            source.source_config.platform = self.platform

        # Apply method customizations
        for method_name, implementation in self.customizations.items():
            # For instance methods that use self, we need to bind them to the source
            if (
                hasattr(implementation, "__self__")
                and implementation.__self__ is not None
            ):
                # This is already a bound method, use __get__ to rebind it to the source
                setattr(source, method_name, implementation.__get__(source))
            else:
                # This is a regular function or static/class method
                setattr(source, method_name, implementation)

        return source

    # Add a direct method for tests that may call this directly
    def get_external_url(self, table_data: Any) -> Optional[str]:
        """
        Get the external URL for a table based on the platform type.

        This method routes to the appropriate implementation based on the platform.

        Args:
            table_data: Table data containing path information

        Returns:
            An external URL or None if not applicable
        """
        if self.platform == "s3":
            return self.get_s3_external_url(table_data, self.aws_region)
        elif self.platform == "gcs":
            return self.get_gcs_external_url(table_data)
        elif self.platform == "abs":
            return self.get_abs_external_url(table_data)
        return None

    def _normalize_gcs_uri_for_pattern_matching(self, uri: str) -> str:
        """
        Normalize GCS URI for pattern matching.

        This method converts gs:// URIs to s3:// URIs for pattern matching purposes,
        fixing the URI mismatch issue in GCS ingestion.

        Args:
            uri: The URI to normalize

        Returns:
            The normalized URI for pattern matching
        """
        if uri.startswith("gs://"):
            return uri.replace("gs://", "s3://", 1)
        return uri

    def _strip_gcs_prefix(self, uri: str) -> str:
        """
        Strip GCS prefix from URI.

        This method removes the gs:// prefix from GCS URIs for path processing.

        Args:
            uri: The URI to strip the prefix from

        Returns:
            The URI without the gs:// prefix
        """
        if uri.startswith("gs://"):
            return uri[5:]  # Remove "gs://" prefix
        return uri


# Factory function to create an adapter for a specific platform
def create_object_store_adapter(
    platform: str,
    aws_region: Optional[str] = None,
    azure_storage_account: Optional[str] = None,
) -> ObjectStoreSourceAdapter:
    """
    Create an adapter for a specific object store platform.

    Args:
        platform: The platform identifier (e.g., "s3", "gcs", "abs")
        aws_region: AWS region for S3 URLs, defaults to us-east-1 if not specified
        azure_storage_account: Azure Storage account name

    Returns:
        An adapter configured for the specified platform
    """
    platform_names = {
        "s3": "Amazon S3",
        "gcs": "Google Cloud Storage",
        "abs": "Azure Blob Storage",
    }

    return ObjectStoreSourceAdapter(
        platform=platform,
        platform_name=platform_names.get(platform, f"Unknown ({platform})"),
        aws_region=aws_region,
        azure_storage_account=azure_storage_account,
    )
