import logging
import os
import warnings
from functools import partial
from typing import List, Optional, Tuple, Union

import orjson
from beartype import beartype
from beartype.roar import BeartypeDecorHintPep585DeprecationWarning
from PIL import Image

import picsellia.pxl_multithreading as mlt
from picsellia import exceptions
from picsellia.colors import Colors
from picsellia.decorators import exception_handler
from picsellia.exceptions import NoConnectorFound, NoDataError, PicselliaError
from picsellia.sdk.connexion import Connexion
from picsellia.sdk.dao import Dao
from picsellia.sdk.data import Data, MultiData
from picsellia.sdk.datasource import DataSource
from picsellia.sdk.tag import Tag
from picsellia.types.enums import DataType, TagTarget
from picsellia.types.schemas import DatalakeSchema
from picsellia.utils import convert_tag_list_to_query_language, filter_payload

logger = logging.getLogger("picsellia")
warnings.filterwarnings("ignore", category=BeartypeDecorHintPep585DeprecationWarning)


class Datalake(Dao):
    def __init__(self, connexion: Connexion, data: dict):
        Dao.__init__(self, connexion, data)

    def __str__(self):
        return f"{Colors.GREEN}Datalake '{self.name}'{Colors.ENDC} (id: {self.id})"

    @property
    def name(self):
        return self._name

    @property
    def connector_id(self):
        if self._connector_id is None:
            raise NoConnectorFound(
                "This datalake has no connector. You cannot retrieve and upload data into this datalake."
            )
        return self._connector_id

    @exception_handler
    @beartype
    def refresh(self, data: dict):
        schema = DatalakeSchema(**data)
        self._name = schema.name
        if schema.connector_id is not None:
            self._connector_id = schema.connector_id
        return schema

    @exception_handler
    @beartype
    def sync(self) -> dict:
        r = self.connexion.get(f"/sdk/datalake/{self.id}").json()
        self.refresh(r)
        return r

    @exception_handler
    @beartype
    def get_resource_url_on_platform(self) -> str:
        """Get platform url of this resource.

        Examples:
            ```python
                print(foo_dataset.get_resource_url_on_platform())
                >>> "https://app.picsellia.com/datalake/62cffb84-b92c-450c-bc37-8c4dd4d0f590"
            ```

        Returns:
            Url on Platform for this resource
        """

        return f"{self.connexion.host}/datalake/{self.id}"

    @exception_handler
    @beartype
    def upload_data(
        self,
        filepaths: Union[str, List[str]],
        tags: Optional[List[Union[str, Tag]]] = None,
        source: Union[str, DataSource, None] = None,
        max_workers: Optional[int] = None,
    ) -> Union[Data, MultiData]:
        """Upload data into this datalake.

        Upload files representing data, into a datalake.
        You can give some tags as a list.
        You can give a source for your data.

        Examples:
            ```python
                tag_car = client.get_data_tag("car")
                tag_huge_car = client.get_data_tag("huge-car")
                source_camera_one = client.get_datasource("camera-one")
                source_camera_two = client.get_datasource("camera-two")

                lake = client.get_datalake()
                lake.upload_data(filepaths=["porsche.png", "ferrari.png"], tags=[tag_car], source=source_camera_one)
                lake.upload_data(filepaths="truck.png", tags=[tag_huge_car], source=source_camera_two)
            ```
        Arguments:
            filepaths (str or List[str]): Filepaths of your data
            tags (List[Tag], optional): Data Tags that will be given to data. Defaults to [].
            source (DataSource, optional): Source of your data.
            max_workers (int): Max workers to use when multithreading upload

        Returns:
            A (Data) object or a (MultiData) object that wraps a list of Data.
        """
        computed_tags_ids = []
        if tags:
            for tag in tags:
                if isinstance(tag, str):
                    computed_tag = self.get_or_create_data_tag(tag)
                else:
                    if tag.target_type == TagTarget.DATA:
                        computed_tag = tag
                    else:
                        computed_tag = self.create_data_tag(tag.name)

                computed_tags_ids.append(computed_tag.id)

        def upload_filepath(path: str):
            try:
                # Retrieve presigned url from platform
                filename = os.path.split(path)[-1]
                object_name = self.connexion.generate_data_object_name(
                    filename, self.connector_id
                )
                _, _, content_type = self.connexion.upload_file(
                    object_name, path, connector_id=self.connector_id
                )

                # Post new image into platform
                with Image.open(path) as image:
                    width, height = image.size
                payload = {
                    "type": DataType.IMAGE,
                    "filename": filename,
                    "object_name": object_name,
                    "content_type": content_type,
                    "meta": {
                        "height": height,
                        "width": width,
                    },
                    "tags": computed_tags_ids,
                    "data_source_id": source.id if source else None,
                }

                r = self.connexion.post(
                    f"/sdk/datalake/{self.id}/datas",
                    data=orjson.dumps(payload),
                ).json()
                return Data(self.connexion, self.id, r)
            except Exception as e:
                logger.error(f"\nCould not upload {path} because {str(e)}")
                return None

        if isinstance(filepaths, str):
            filepaths = [filepaths]

        logger.info("🌎 Starting upload..")

        results = mlt.do_mlt_function(
            filepaths, upload_filepath, max_workers=max_workers
        )
        data_list: List[Data] = [v for _, v in results.items() if v is not None]

        if len(data_list) != len(filepaths):
            logger.error(
                f"❌ {len(filepaths) - len(data_list)} data not uploaded. Check readability / check paths"
            )

        if len(data_list) == 0:
            raise PicselliaError("Nothing has been uploaded.")
        elif len(data_list) == 1:
            logger.info(f"✅ {data_list[0].filename} data uploaded in {self}")
            return data_list[0]
        else:
            logger.info(f"✅ {len(data_list)} data uploaded in {self}")
            return MultiData(self.connexion, self.id, data_list)

    @exception_handler
    @beartype
    def find_data(
        self, filename: Optional[str] = None, object_name: Optional[str] = None
    ) -> Data:
        """Find a data into this datalake

        You can find it by giving its filename or its object name

        Examples:
            ```python
                my_data = my_datalake.find_asset(filename="test.png")
            ```
        Arguments:
            filename (str, optional): filename of the data. Defaults to None.
            object_name (str, optional): object name in the storage S3. Defaults to None.

        Raises:
            If no data match the query, it will raise a NotFoundError.
            In some case, it can raise an InvalidQueryError,
                it might be because platform stores 2 data matching this query (for example if filename is duplicated)

        Returns:
            The (Data) found
        """
        assert not (
            filename is None and object_name is None
        ), "Select at least one criteria to find a data"

        params = {}
        if filename is not None:
            params["filename"] = filename

        if object_name is not None:
            params["object_name"] = object_name

        r = self.connexion.get(
            f"/sdk/datalake/{self.id}/datas/find", params=params
        ).json()
        return Data(self.connexion, self.id, r)

    @exception_handler
    @beartype
    def list_data(
        self,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        page_size: Optional[int] = None,
        order_by: Optional[List[str]] = None,
        tags: Union[str, Tag, List[Union[str, Tag]], None] = None,
        filenames: Optional[List[str]] = None,
    ) -> MultiData:
        """List all data of this datalake.

        If there is no data, raise a NoDataError exception.

        Returned object is a MultiData. An object that allows manipulation of a bunch of data.
        You can add tags on them or feed a dataset with them.

        Examples:
            ```python
                lake = client.get_datalake()
                data = lake.list_data()
            ```

        Raises:
            NoDataError: When datalake has no data, raise this exception.

        Returns:
            A (MultiData) object that wraps a list of (Data).
        """
        q = convert_tag_list_to_query_language(tags)

        datas = mlt.do_paginate(
            limit,
            offset,
            page_size,
            partial(self._do_list_data, q, order_by, filenames),
        )

        if len(datas) == 0:
            raise NoDataError("No data found in this datalake")

        return MultiData(self.connexion, self.id, datas)

    @exception_handler
    @beartype
    def _do_list_data(
        self,
        q: Optional[str],
        order_by: Optional[List[str]],
        filenames: Optional[List[str]],
        limit: int,
        offset: int,
    ) -> Tuple[List[Data], int]:
        params = {"limit": limit, "offset": offset}

        if order_by is not None:
            params["order_by"] = order_by
        if q is not None:
            params["q"] = q
        if filenames is not None:
            params["filenames"] = filenames

        r = self.connexion.get(f"/sdk/datalake/{self.id}/datas", params=params).json()
        return list(map(partial(Data, self.connexion, self.id), r["items"])), r["count"]

    @exception_handler
    @beartype
    def create_data_tag(self, name: str) -> Tag:
        """Create a data tag used in this datalake

        Examples:
            ```python
                tag_car = lake.create_data_tag("car")
            ```
        Arguments:
            name (str): Name of this tag

        Returns:
            A (Tag) object
        """
        payload = {"name": name}
        r = self.connexion.post(
            f"/sdk/datalake/{self.id}/tags", data=orjson.dumps(payload)
        ).json()
        return Tag(self.connexion, r)

    @exception_handler
    @beartype
    def get_data_tag(self, name: str) -> Tag:
        """Retrieve a data tag used in this datalake.

        Examples:
            ```python
                tag_car = lake.get_data_tag("car")
            ```
        Arguments:
            name (str): Name of the tag you're looking for

        Returns:
            A (Tag) object
        """
        params = {"name": name}
        r = self.connexion.get(
            f"/sdk/datalake/{self.id}/tags/find", params=params
        ).json()
        return Tag(self.connexion, r)

    @exception_handler
    @beartype
    def get_or_create_data_tag(self, name: str) -> Tag:
        """Retrieve a data tag used in this datalake by its name.
        If tag does not exist, create it and return it.

        Examples:
            ```python
                tag = self.get_or_create_data_tag("new_tag")
            ```
        Arguments:
            name (str): Tag to retrieve or create

        Returns:
            A (Tag) object
        """
        try:
            return self.get_data_tag(name)
        except exceptions.ResourceNotFoundError:
            return self.create_data_tag(name)

    @exception_handler
    @beartype
    def list_data_tags(
        self,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        order_by: Optional[List[str]] = None,
    ) -> List[Tag]:
        """List all tags of this datalake

        Examples:
            ```python
                tags = lake.list_data_tags()
                assert tag_car in tags
            ```

        Returns:
            A List of (Tag)
        """
        params = {"limit": limit, "offset": offset, "order_by": order_by}
        params = filter_payload(params)
        r = self.connexion.get(f"/sdk/datalake/{self.id}/tags", params=params).json()
        return list(map(partial(Tag, self.connexion), r["items"]))

    @exception_handler
    @beartype
    def find_all_datas(self, object_names: List[str]) -> MultiData:
        payload = {"object_names": object_names}
        r = self.connexion.xget(
            f"/sdk/datalake/{self.id}/datas", data=orjson.dumps(payload)
        ).json()
        datas = list(map(partial(Data, self.connexion, self.id), r["items"]))

        if len(datas) == 0:
            raise NoDataError("No data found")

        return MultiData(self.connexion, self.id, datas)
