""" object store functions """

import logging
from typing import Any
from typing_extensions import Literal

import pandas as pd
import urllib3
from box import Box
from minio import Minio
from prefect import Task
from prefect.utilities.tasks import defaults_from_attrs

# pylint: disable=arguments-differ, too-many-arguments


class ObjstrClient(Task):
    """setup object storage client via minio"""

    def __init__(
        self,
        config_box: Box = None,
        http_client: urllib3.poolmanager.PoolManager = None,
        **kwargs: Any
    ):
        self.config_box = config_box
        self.http_client = http_client
        super().__init__(**kwargs)

    @defaults_from_attrs("config_box", "http_client")
    def run(
        self, config_box: Box, http_client: urllib3.poolmanager.PoolManager = None
    ) -> Minio:

        logging.info(
            "Creating object store client for endpoint %s.", config_box.endpoint
        )
        client = Minio(
            endpoint=config_box.endpoint,
            access_key=config_box.key,
            secret_key=config_box.secret,
            secure=True,
            http_client=http_client,
        )

        return client


class ObjstrMakeBucket(Task):
    """attempts to make a bucket if not already available"""

    def __init__(self, client: Minio = None, bucket_name: str = None, **kwargs: Any):
        self.client = client
        self.bucket_name = bucket_name

        super().__init__(**kwargs)

    @defaults_from_attrs("client", "bucket_name")
    def run(self, client: Minio, bucket_name: str):

        # make bucket if not exists
        if not client.bucket_exists(bucket_name):
            client.make_bucket(bucket_name)
            logging.info("Created bucket %s", bucket_name)
        else:
            logging.info("Bucket %s already exists, taking no actions.", bucket_name)


class ObjstrPut(Task):
    """put data as object in object store"""

    def __init__(
        self,
        client: Minio = None,
        bucket_name: str = None,
        object_name: str = None,
        data: object = None,
        length: int = -1,
        part_size: int = 5 * 1024 * 1024,
        **kwargs: Any
    ):

        self.client = client
        self.bucket_name = bucket_name
        self.object_name = object_name
        self.data = data
        self.length = length
        self.part_size = part_size

        super().__init__(**kwargs)

    @defaults_from_attrs(
        "bucket_name", "object_name", "data", "length", "part_size", "client"
    )
    def run(
        self,
        client: Minio,
        bucket_name: str,
        object_name: str,
        data: object,
        length: int = None,
        part_size: int = None,
    ):
        if not length:
            length = self.length

        if not part_size:
            part_size = self.part_size

        # upload file as object
        client.put_object(
            bucket_name=bucket_name,
            object_name=object_name,
            data=data,
            length=length,
            part_size=part_size,
        )

        logging.info("Put data under %s as %s", bucket_name, object_name)


class ObjstrGet(Task):
    """get object as data from object store"""

    def __init__(
        self,
        client: Minio = None,
        bucket_name: str = None,
        object_name: str = None,
        **kwargs: Any
    ):
        self.client = client
        self.bucket_name = bucket_name
        self.object_name = object_name

        super().__init__(**kwargs)

    @defaults_from_attrs(
        "client",
        "bucket_name",
        "object_name",
    )
    def run(self, client: Minio, bucket_name: str, object_name: str) -> Any:

        # get object as file
        data = client.get_object(
            bucket_name=bucket_name,
            object_name=object_name,
        )

        logging.info(
            "Retrieved object %s under %s",
            object_name,
            bucket_name,
        )
        return data


class ObjstrGetAsDF(Task):
    """get object and return as pandas.dataframe from object store"""

    def __init__(
        self,
        client: Minio = None,
        bucket_name: str = None,
        object_name: str = None,
        dftype: Literal["csv", "parquet", "excel"] = None,
        **kwargs: Any
    ):
        self.client = client
        self.bucket_name = bucket_name
        self.object_name = object_name
        self.dftype = dftype

        super().__init__(**kwargs)

    @defaults_from_attrs("client", "bucket_name", "object_name", "dftype")
    def run(
        self,
        client: Minio,
        bucket_name: str,
        object_name: str,
        dftype: Literal["csv", "parquet", "excel"],
        **kwargs: Any
    ) -> Any:

        # get object as file
        data = client.get_object(
            bucket_name=bucket_name,
            object_name=object_name,
        )

        logging.info(
            "Retrieved object %s under %s",
            object_name,
            bucket_name,
        )

        if dftype == "csv":
            pd_dataframe = pd.read_csv(filepath_or_buffer=data, **kwargs)
        elif dftype == "parquet":
            pd_dataframe = pd.read_parquet(filepath_or_buffer=data, **kwargs)

        return pd_dataframe


class ObjstrFPut(Task):
    """put file as object in object store"""

    def __init__(
        self,
        client: Minio = None,
        bucket_name: str = None,
        object_name: str = None,
        file_path: str = None,
        **kwargs: Any
    ):

        self.client = client
        self.bucket_name = bucket_name
        self.object_name = object_name
        self.file_path = file_path

        super().__init__(**kwargs)

    @defaults_from_attrs("client", "bucket_name", "object_name", "file_path")
    def run(
        self,
        client: Minio,
        bucket_name: str,
        object_name: str,
        file_path: str,
    ):

        # upload file as object
        client.fput_object(
            bucket_name=bucket_name,
            object_name=object_name,
            file_path=file_path,
        )

        logging.info("Put file %s under %s as %s", file_path, bucket_name, object_name)


class ObjstrFGet(Task):
    """get object as file from object store"""

    def __init__(
        self,
        client: Minio = None,
        bucket_name: str = None,
        object_name: str = None,
        file_path: str = None,
        **kwargs: Any
    ):
        self.client = client
        self.bucket_name = bucket_name
        self.object_name = object_name
        self.file_path = file_path

        super().__init__(**kwargs)

    @defaults_from_attrs("client", "bucket_name", "object_name", "file_path")
    def run(self, client: Minio, bucket_name: str, object_name: str, file_path: str):

        # get object as file
        client.fget_object(
            bucket_name=bucket_name,
            object_name=object_name,
            file_path=file_path,
        )

        logging.info(
            "Retrieved object %s under %s as file %s",
            object_name,
            bucket_name,
            file_path,
        )
