"""Rucio plugin for Snakemake."""

import dataclasses
import inspect
from collections.abc import Iterable, Sequence
from urllib.parse import urlparse

import rucio.client
import rucio.client.baseclient
import rucio.client.downloadclient
import rucio.client.uploadclient
import rucio.common.exception
from snakemake.logging import logger as smk_logger
from snakemake_interface_storage_plugins.io import IOCacheStorageInterface, Mtime
from snakemake_interface_storage_plugins.settings import StorageProviderSettingsBase
from snakemake_interface_storage_plugins.storage_object import (
    StorageObjectGlob,
    StorageObjectRead,
    StorageObjectWrite,
    retry_decorator,
)
from snakemake_interface_storage_plugins.storage_provider import (
    ExampleQuery,
    Operation,
    QueryType,
    StorageProviderBase,
    StorageQueryValidationResult,
)

# Optional:
# Define settings for your storage plugin (e.g. host url, credentials).
# They will occur in the Snakemake CLI as --storage-<storage-plugin-name>-<param-name>
# Make sure that all defined fields are 'Optional' and specify a default value
# of None or anything else that makes sense in your case.
# Note that we allow storage plugin settings to be tagged by the user. That means,
# that each of them can be specified multiple times (an implicit nargs=+), and
# the user can add a tag in front of each value (e.g. tagname1:value1 tagname2:value2).
# This way, a storage plugin can be used multiple times within a workflow with different
# settings.


# Copy the available settings over from the Rucio client.
_RUCIO_CLIENT_CLS = rucio.client.baseclient.BaseClient


def _get_help(arg: str) -> str:
    """Read the help text for a Rucio client argument."""
    doc = _RUCIO_CLIENT_CLS.__init__.__doc__
    for line in doc.split("\n"):
        if arg in line:
            return line.split(f":param {arg}:", 1)[1]
    return ""


StorageProviderSettings = dataclasses.make_dataclass(
    "StorageProviderSettings",
    fields=[
        (
            f.name,
            f.annotation,
            dataclasses.field(
                default=f.default,
                metadata={"help": _get_help(f.name)},
            ),
        )
        for f in inspect.signature(_RUCIO_CLIENT_CLS).parameters.values()
        if f.name not in ("logger",)
    ]
    + [
        (
            "ignore_checksum",
            bool,
            dataclasses.field(
                default=False,
                metadata={
                    "help": "If true, skips the checksum validation between the downloaded file and the rucio catalogue.",
                },
            ),
        ),
        (
            "download_rse",
            str,
            dataclasses.field(
                default=None,
                metadata={
                    "help": "Rucio Storage Element (RSE) expression to download files from.",
                },
            ),
        ),
        (
            "upload_rse",
            str,
            dataclasses.field(
                default=None,
                metadata={
                    "help": "Rucio Storage Element (RSE) expression to upload files to.",
                },
            ),
        ),
        (
            "cache_scope",
            bool,
            dataclasses.field(
                default=False,
                metadata={
                    "help": "If true, minimize the number of server calls by caching the size and creation time of all files in the same scope.",
                },
            ),
        ),
    ],
    bases=(StorageProviderSettingsBase,),
)


class StorageProvider(StorageProviderBase):
    """Rucio storage provider."""

    # For compatibility with future changes, you should not overwrite the __init__
    # method. Instead, use __post_init__ to set additional attributes and initialize
    # further stuff.

    def __post_init__(self) -> None:
        """Initialize the storage provider."""
        # This is optional and can be removed if not needed.
        # Alternatively, you can e.g. prepare a connection to your storage backend here.
        # and set additional attributes.
        valid_client_args = inspect.signature(_RUCIO_CLIENT_CLS).parameters
        client_kwargs = {
            k: v
            for k, v in dataclasses.asdict(self.settings).items()
            if k in valid_client_args
        }

        client = rucio.client.Client(logger=smk_logger, **client_kwargs)
        self.client = client
        self.dclient = rucio.client.downloadclient.DownloadClient(
            client, logger=smk_logger.logger
        )
        self.uclient = rucio.client.uploadclient.UploadClient(
            client, logger=smk_logger.logger
        )

    @classmethod
    def example_queries(cls) -> list[ExampleQuery]:
        """Return example queries with description for this storage provider."""
        return [
            ExampleQuery(
                query="rucio://myscope/myfile.txt",
                type=QueryType.ANY,
                description="A file in a Rucio scope.",
            ),
        ]

    def rate_limiter_key(
        self,
        query: str,  # noqa: ARG002
        operation: Operation,  # noqa: ARG002
    ) -> str:
        """Return a key for identifying a rate limiter given a query and an operation.

        This is used to identify a rate limiter for the query.
        E.g. for a storage provider like http that would be the host name.
        For s3 it might be just the endpoint URL.
        """
        return self.settings.rucio_host

    def default_max_requests_per_second(self) -> float:
        """Return the default maximum number of requests per second."""
        return 100

    def use_rate_limiter(self) -> bool:
        """Return False if no rate limiting is needed for this provider."""
        return False

    @classmethod
    def is_valid_query(cls, query: str) -> StorageQueryValidationResult:
        """Return whether the given query is valid for this storage provider."""
        # Ensure that also queries containing wildcards (e.g. {sample}) are accepted
        # and considered valid. The wildcards will be resolved before the storage
        # object is actually used.
        try:
            parsed = urlparse(query)
        except Exception as exc:  # noqa: BLE001
            return StorageQueryValidationResult(
                query=query,
                valid=False,
                reason=f"cannot be parsed as URL ({exc})",
            )
        if parsed.scheme != "rucio":
            return StorageQueryValidationResult(
                query=query,
                valid=False,
                reason="must start with rucio (rucio://...)",
            )
        return StorageQueryValidationResult(
            query=query,
            valid=True,
        )


class StorageObject(StorageObjectRead, StorageObjectWrite, StorageObjectGlob):
    """Rucio storage object. This represents a single file in Rucio."""

    # For compatibility with future changes, you should not overwrite the __init__
    # method. Instead, use __post_init__ to set additional attributes and initialize
    # further stuff.

    def __post_init__(self) -> None:
        """Initialize the storage object."""
        # This is optional and can be removed if not needed.
        # Alternatively, you can e.g. prepare a connection to your storage backend here.
        # and set additional attributes.
        if not self.is_valid_query():
            raise ValueError(self.query)
        parsed = urlparse(self.query)
        self.scope = parsed.netloc
        self.file = parsed.path.lstrip("/")

    @property
    def client(self) -> rucio.client.Client:
        """Rucio client."""
        return self.provider.client

    async def inventory(self, cache: IOCacheStorageInterface) -> None:
        """Read the information about all files in a scope efficiently.

        From this file, try to find as much existence and modification date
        information as possible. Only retrieve that information that comes for free
        given the current object.
        """
        # This retrieves information about all the files in the file scope at
        # once. This is faster than sending a new request for each file, but
        # If this becomes too slow because scopes contain too many files,
        # we may need to add a setting to disable it.

        # If this is implemented in a storage object, results have to be stored in
        # the given IOCache object, using self.cache_key() as key.
        # Optionally, this can take a custom local suffix, needed e.g. when you want
        # to cache more items than the current query: self.cache_key(local_suffix=...)
        if self.get_inventory_parent() in cache.exists_in_storage:
            # record has been inventorized before
            return

        if self.provider.settings.cache_scope:
            # Cache the entire scope
            if self.scope not in self.client.list_scopes():
                # check if scope exists
                cache.exists_in_storage[self.cache_key()] = False
            else:
                cache.exists_in_storage[self.get_inventory_parent()] = True
                files = self.client.list_dids(
                    scope=self.scope,
                    filters={"type": "file"},
                )
                batch_size = 500
                batch = []
                for i, file in enumerate(files, 1):
                    batch.append(file)
                    if i % batch_size == 0:
                        self._handle(cache, batch)
                        batch.clear()
                self._handle(cache, batch)
        else:
            self._handle(cache, [self.file])

    def _handle(self, cache: IOCacheStorageInterface, files: Sequence[str]) -> None:
        """Add a sequence of files to the cache."""
        dids = [{"scope": self.scope, "name": f} for f in files]
        for file, meta in zip(files, self.client.get_metadata_bulk(dids)):
            key = self.cache_key(f"{self.scope}/{file}")
            cache.mtime[key] = Mtime(storage=meta["updated_at"].timestamp())
            cache.size[key] = meta["bytes"]
            cache.exists_in_storage[key] = True

    def get_inventory_parent(self) -> str:
        """Return the parent directory of this object."""
        return self.scope

    def local_suffix(self) -> str:
        """Return a unique suffix for the local path, determined from self.query."""
        return f"{self.scope}/{self.file}"

    def cleanup(self) -> None:
        """Perform local cleanup of any remainders of the storage object."""
        # self.local_path() should not be removed, as this is taken care of by
        # Snakemake.

    # Fallible methods should implement some retry logic.
    # The easiest way to do this (but not the only one) is to use the retry_decorator
    # provided by snakemake-interface-storage-plugins.
    @retry_decorator
    def exists(self) -> bool:
        """Return True if the object exists."""
        try:
            self.client.get_did(scope=self.scope, name=self.file)
        except rucio.common.exception.DataIdentifierNotFound:
            return False
        return True

    @retry_decorator
    def mtime(self) -> float:
        """Return the modification time."""
        meta = self.client.get_metadata(scope=self.scope, name=self.file)
        return meta["updated_at"].timestamp()

    @retry_decorator
    def size(self) -> int:
        """Return the size in bytes."""
        did = self.client.get_did(scope=self.scope, name=self.file)
        return did["bytes"]

    @retry_decorator
    def retrieve_object(self) -> None:
        """Download the file to self.local_path()."""
        self.provider.dclient.download_dids(
            [
                {
                    "base_dir": self.local_path().parent,
                    "did": f"{self.scope}:{self.file}",
                    "ignore_checksum": self.provider.settings.ignore_checksum,
                    "no_subdir": True,
                    "rse": self.provider.settings.download_rse,
                },
            ],
            num_threads=1,
        )

    @retry_decorator
    def store_object(self) -> None:
        """Upload the file."""
        if self.exists():
            msg = f'File "{self.scope}:{self.file}" already exists on Rucio'
            raise ValueError(msg)
        rse = self.provider.settings.upload_rse
        if rse is None:
            msg = "Please specify the `upload_rse`."
            raise ValueError(msg)
        self.provider.uclient.upload(
            [
                {
                    "path": self.local_path(),
                    "did_scope": self.scope,
                    "rse": self.provider.settings.upload_rse,
                    "register_after_upload": True,
                },
            ],
        )

    @retry_decorator
    def remove(self) -> None:
        """Remove the file from the storage."""
        raise NotImplementedError

    @retry_decorator
    def list_candidate_matches(self) -> Iterable[str]:
        """Return a list of candidate matches in the storage for the query."""
        # This is used by glob_wildcards() to find matches for wildcards in the query.
        # The method has to return concretized queries without any remaining wildcards.
        # Use snakemake_executor_plugins.io.get_constant_prefix(self.query) to get the
        # prefix of the query before the first wildcard.
        return self.client.list_dids(
            scope=self.scope,
            filters={
                "name": self.file,
                "type": "file",
            },
        )
