# Copyright 2025 The MelowRAG Author @erfanzar (Erfan Zare Chavoshi).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import re

try:
    from libcloud.storage.providers import get_driver  # type:ignore

    LIBCLOUD = True
except ImportError:
    LIBCLOUD = False

from .base import Task


class StorageTask(Task):
    """
    Task that processes object storage buckets. Supports local and cloud providers in Apache libcloud.
    """

    PREFIX = r"(\w+):\/\/.*"
    PATH = r"\w+:\/\/(.*)"

    def register(self, key=None, secret=None, host=None, port=None, token=None, region=None):
        """
        Checks if required dependencies are installed. Reads in cloud storage parameters.

        Args:
            key: provider-specific access key
            secret: provider-specific access secret
            host: server host name
            port: server port
            token: temporary session token
            region: storage region
        """

        if not LIBCLOUD:
            raise ImportError('StorageTask is not available - install "workflow" extra to enable')

        # pylint: disable=W0201
        self.key = key
        self.secret = secret
        self.host = host
        self.port = port
        self.token = token
        self.region = region

    def __call__(self, elements, executor=None):
        outputs = []
        for element in elements:
            if self.matches(element):
                outputs.extend(super().__call__(self.list(element), executor))
            else:
                outputs.append(element)

        return outputs

    def matches(self, element):
        """
        Determines if this element is a storage element.

        Args:
            element: input storage element

        Returns:
            True if this is a storage element
        """

        return re.match(StorageTask.PREFIX, self.upack(element, True).lower())

    def list(self, element):
        """
        Gets a list of urls for a object container.

        Args:
            element: object container

        Returns:
            list of urls
        """

        provider = re.sub(StorageTask.PREFIX, r"\1", element.lower())
        path = re.sub(StorageTask.PATH, r"\1", element)

        key = self.key if self.key is not None else os.environ.get("ACCESS_KEY")
        secret = self.secret if self.secret is not None else os.environ.get("ACCESS_SECRET")

        key, container = (os.path.dirname(path), os.path.basename(path)) if key is None else (key, path)

        parts = container.split("/", 1)
        container, prefix = (parts[0], parts[1]) if len(parts) > 1 else (container, None)

        driver = get_driver(provider)

        client = driver(
            key,
            secret,
            **{field: getattr(self, field) for field in ["host", "port", "region", "token"] if getattr(self, field)},
        )

        container = client.get_container(container_name=container)
        return [
            client.get_object_cdn_url(obj) for obj in client.list_container_objects(container=container, prefix=prefix)
        ]
