# coding: utf-8
import os
import re
import tempfile
from typing import Dict, List, Literal, Optional
from uuid import UUID

from numpy import uint8
from requests_toolbelt import MultipartDecoder, MultipartEncoder

import supervisely.volume_annotation.constants as constants
from supervisely._utils import batched
from supervisely.api.entity_annotation.figure_api import FigureApi, FigureInfo
from supervisely.api.module_api import ApiField
from supervisely.geometry.closed_surface_mesh import ClosedSurfaceMesh
from supervisely.geometry.mask_3d import Mask3D
from supervisely.io.fs import ensure_base_path, file_exists
from supervisely.video_annotation.key_id_map import KeyIdMap
from supervisely.volume.nrrd_encoder import encode
from supervisely.volume_annotation.plane import Plane
from supervisely.volume_annotation.volume_figure import VolumeFigure


class VolumeFigureApi(FigureApi):
    """
    :class:`VolumeFigure<supervisely.volume_annotation.volume_figure.VolumeFigure>` for a single volume.
    """

    def create(
        self,
        volume_id: int,
        object_id: int,
        plane_name: str,
        slice_index: int,
        geometry_json: dict,
        geometry_type: str,
        custom_data: Optional[dict] = None,
        # track_id=None,
    ):
        """
        Create new VolumeFigure for given slice in given volume ID.

        :param volume_id: Volume ID in Supervisely.
        :type volume_id: int
        :param object_id: ID of the object to which the VolumeFigure belongs.
        :type object_id: int
        :param plane_name: :py:class:`Plane<supervisely.volume_annotation.plane.Plane>` of the slice in volume.
        :type plane_name: str
        :param slice_index: Number of the slice to add VolumeFigure.
        :type slice_index: int
        :param geometry_json: Parameters of geometry for VolumeFigure.
        :type geometry_json: dict
        :param geometry_type: Type of VolumeFigure geometry.
        :type geometry_type: str
        :return: New figure ID
        :rtype: :class:`int`
        :Usage example:

         .. code-block:: python

            import supervisely as sly

            from supervisely.volume_annotation.plane import Plane

            os.environ['SERVER_ADDRESS'] = 'https://app.supervisely.com'
            os.environ['API_TOKEN'] = 'Your Supervisely API Token'
            api = sly.Api.from_env()

            volume_id = 19581134
            object_id = 5565016
            slice_index = 0
            plane_name = Plane.AXIAL
            geometry_json = {'points': {'exterior': [[500, 500], [1555, 1500]], 'interior': []}}
            geometry_type = 'rectangle'

            figure_id = api.volume.figure.create(
                volume_id,
                object_id,
                plane_name,
                slice_index,
                geometry_json,
                geometry_type
            ) # 87821207
        """

        Plane.validate_name(plane_name)

        return super().create(
            volume_id,
            object_id,
            {
                constants.SLICE_INDEX: slice_index,
                constants.NORMAL: Plane.get_normal(plane_name),
                # for backward compatibility
                ApiField.META: {
                    constants.SLICE_INDEX: slice_index,
                    constants.NORMAL: Plane.get_normal(plane_name),
                },
            },
            geometry_json,
            geometry_type,
            custom_data=custom_data,
            # track_id,
        )

    def append_bulk(self, volume_id: int, figures: List[VolumeFigure], key_id_map: KeyIdMap):
        """
        Add VolumeFigures to given Volume by ID.

        :param volume_id: Volume ID in Supervisely.
        :type volume_id: int
        :param key_id_map: KeyIdMap object.
        :type key_id_map: KeyIdMap
        :param figures: List of VolumeFigure objects.
        :type figures: List[VolumeFigure]
        :return: None
        :rtype: :class:`NoneType`
        :Usage example:

         .. code-block:: python

            import supervisely as sly

            from supervisely.volume_annotation.plane import Plane

            os.environ['SERVER_ADDRESS'] = 'https://app.supervisely.com'
            os.environ['API_TOKEN'] = 'Your Supervisely API Token'
            api = sly.Api.from_env()

            project_id = 19370
            volume_id = 19617444

            key_id_map = sly.KeyIdMap()

            project_meta_json = api.project.get_meta(project_id)
            project_meta = sly.ProjectMeta.from_json(project_meta_json)

            vol_ann_json = api.volume.annotation.download(volume_id)
            vol_ann = sly.VolumeAnnotation.from_json(vol_ann_json, project_meta, key_id_map)
            volume_obj_collection = vol_ann.objects.to_json()
            vol_obj = sly.VolumeObject.from_json(volume_obj_collection[1], project_meta)

            figure = sly.VolumeFigure(
                vol_obj,
                sly.Rectangle(20, 20, 129, 200),
                sly.Plane.AXIAL,
                45,
            )

            api.volume.figure.append_bulk(volume_id, [figure], key_id_map)
        """

        if len(figures) == 0:
            return

        keys = []
        mask3d_keys = []
        figures_json = []
        mask3d_figures = []

        for figure in figures:
            if figure.geometry.name() == Mask3D.name():
                mask3d_keys.append(figure.key())
                mask3d_figures.append(figure)
            else:
                keys.append(figure.key())
                figures_json.append(figure.to_json(key_id_map, save_meta=True))
        # Figure is missing required field \"meta.normal\"","index":0}}
        self._append_bulk(volume_id, figures_json, keys, key_id_map)
        if len(mask3d_figures) != 0:
            self._append_bulk_mask3d(volume_id, mask3d_figures, mask3d_keys, key_id_map)

    def _download_geometries_batch(self, ids: List[int]):
        """
        Private method. Download figures geometries with given IDs from storage.
        """

        for batch_ids in batched(ids):
            response = self._api.post("figures.bulk.download.geometry", {ApiField.IDS: batch_ids})
            decoder = MultipartDecoder.from_response(response)
            for part in decoder.parts:
                content_utf8 = part.headers[b"Content-Disposition"].decode("utf-8")
                # Find name="1245" preceded by a whitespace, semicolon or beginning of line.
                # The regex has 2 capture group: one for the prefix and one for the actual name value.
                figure_id = int(re.findall(r'(^|[\s;])name="(\d*)"', content_utf8)[0][1])
                yield figure_id, part

    def download_stl_meshes(self, ids: List[int], paths: List[str]):
        """
        Download STL meshes for the specified figure IDs and saves them to the specified paths.

        :param ids: VolumeFigure ID in Supervisely.
        :type ids: int
        :param paths: List of paths to download.
        :type paths: List[str]
        :return: None
        :rtype: :class:`NoneType`
        :Usage example:

         .. code-block:: python

            import supervisely as sly

            os.environ['SERVER_ADDRESS'] = 'https://app.supervisely.com'
            os.environ['API_TOKEN'] = 'Your Supervisely API Token'
            api = sly.Api.from_env()

            STORAGE_DIR = sly.app.get_data_dir()

            volume_id = 19371414
            project_id = 17215

            volume = api.volume.get_info_by_id(volume_id)

            key_id_map = sly.KeyIdMap()
            project_meta_json = api.project.get_meta(project_id)
            project_meta = sly.ProjectMeta.from_json(project_meta_json)

            vol_ann_json = api.volume.annotation.download(volume_id)
            id_to_paths = {}
            vol_ann = sly.VolumeAnnotation.from_json(vol_ann_json, project_meta, key_id_map)

            for sp_figure in vol_ann.spatial_figures:
                figure_id = key_id_map.get_figure_id(sp_figure.key())
                id_to_paths[figure_id] = f"{STORAGE_DIR}/{sp_figure.key().hex}.stl"
            if id_to_paths:
                api.volume.figure.download_stl_meshes(*zip(*id_to_paths.items()))
        """

        if len(ids) == 0:
            return
        if len(ids) != len(paths):
            raise RuntimeError('Can not match "ids" and "paths" lists, len(ids) != len(paths)')

        id_to_path = {id: path for id, path in zip(ids, paths)}
        for img_id, resp_part in self._download_geometries_batch(ids):
            ensure_base_path(id_to_path[img_id])
            with open(id_to_path[img_id], "wb") as w:
                w.write(resp_part.content)

    def interpolate(self, volume_id: int, spatial_figure: VolumeFigure, key_id_map: KeyIdMap):
        """
        Interpolate a spatial figure with a ClosedSurfaceMesh geometry.

        :param volume_id: VolumeFigure ID in Supervisely.
        :type volume_id: int
        :param spatial_figure: Spatial figure to interpolate.
        :type spatial_figure: VolumeFigure
        :param key_id_map: KeyIdMap object.
        :type key_id_map: KeyIdMap
        :return: None
        :rtype: :class:`NoneType`
        :Usage example:

         .. code-block:: python

            import supervisely as sly

            os.environ['SERVER_ADDRESS'] = 'https://app.supervisely.com'
            os.environ['API_TOKEN'] = 'Your Supervisely API Token'
            api = sly.Api.from_env()

            volume_id = 19371414
            project_id = 17215

            volume = api.volume.get_info_by_id(volume_id)

            key_id_map = sly.KeyIdMap()
            project_meta_json = api.project.get_meta(project_id)
            project_meta = sly.ProjectMeta.from_json(project_meta_json)

            vol_ann_json = api.volume.annotation.download(volume_id)
            id_to_paths = {}
            vol_ann = sly.VolumeAnnotation.from_json(vol_ann_json, project_meta, key_id_map)

            for sp_figure in vol_ann.spatial_figures:
                res = volume_figure_api.interpolate(volume_id, sp_figure, key_id_map)
        """

        if type(spatial_figure._geometry) != ClosedSurfaceMesh:
            raise TypeError(
                "Interpolation can be created only for figures with geometry ClosedSurfaceMesh"
            )
        object_id = key_id_map.get_object_id(spatial_figure.volume_object.key())
        response = self._api.post(
            "figures.volumetric_interpolation",
            {ApiField.VOLUME_ID: volume_id, ApiField.OBJECT_ID: object_id},
        )
        return response.content

    # def interpolate_batch(
    #     self, volume_id, spatial_figures: List[VolumeFigure], key_id_map: KeyIdMap
    # ):
    #     # raise NotImplementedError()
    #     # STL mesh interpolations can not be uploaded:
    #     # 400 Client Error: Bad Request for url: public/api/v3/figures.bulk.add ({"error":"Please, use \"figures.bulk.upload.geometry\" method to update figures with \"geometryType\" closed_surface_mesh","details":{"figures":[14]}})
    #     # figures.bulk.upload.geometry

    #     meshes = []
    #     for mesh in spatial_figures:
    #         object_id = key_id_map.get_object_id(mesh.volume_object.key())
    #         response = self._api.post(
    #             "figures.volumetric_interpolation",
    #             {ApiField.VOLUME_ID: volume_id, ApiField.OBJECT_ID: object_id},
    #         )
    #         figure_id = key_id_map.get_figure_id(mesh.key())

    #         # @TODO: load from disk or get from server
    #         meshes.append(response.json())

    #     # for batch in batched(items_to_upload):
    #     #     content_dict = {}
    #     #     for idx, item in enumerate(batch):
    #     #         content_dict["{}-file".format(idx)] = (str(idx), func_item_to_byte_stream(item), 'nrrd/*')
    #     #     encoder = MultipartEncoder(fields=content_dict)
    #     #     self._api.post('import-storage.bulk.upload', encoder)
    #     #     if progress_cb is not None:
    #     #         progress_cb(len(batch))

    #     # content_dict = {}
    #     # for idx, item in enumerate(meshes):
    #     #     content_dict[f"{idx}-mesh"] = (
    #     #         str(idx),
    #     #         func_item_to_byte_stream(item),
    #     #     )
    #     # encoder = MultipartEncoder(fields=content_dict)

    #     # self._api.post("figures.bulk.upload.geometry", encoder)

    #     # # if progress_cb is not None:
    #     # #     progress_cb(len(batch))

    #     # #     self._api.post(
    #     # #         "figures.bulk.upload.geometry",
    #     # #         {ApiField.FIGURE_ID: figure_id, ApiField.GEOMETRY: response.json()},
    #     # #     )

    #     # #     results.append(response.json())
    #     # return results

    def _upload_meshes_batch(self, figure2bytes):
        """
        Private method. Upload figures geometry by given ID to storage.

        :param figure2bytes: Dictionary with figures IDs and geometries.
        :type figure2bytes: dict
        :rtype: :class:`NoneType`
        :Usage example:
        """

        for figure_id, figure_bytes in figure2bytes.items():
            content_dict = {
                ApiField.FIGURE_ID: str(figure_id),
                ApiField.GEOMETRY: (str(figure_id), figure_bytes, "application/sla"),
            }
            encoder = MultipartEncoder(fields=content_dict)
            resp = self._api.post("figures.bulk.upload.geometry", encoder)

    def upload_stl_meshes(
        self,
        volume_id: int,
        spatial_figures: List[VolumeFigure],
        key_id_map: KeyIdMap,
        interpolation_dir=None,
    ):
        """
        Upload existing interpolations or create on the fly and and add them to empty mesh figures.

        :param volume_id: VolumeFigure ID in Supervisely.
        :type volume_id: int
        :param spatial_figures: List of spatial figures to upload.
        :type spatial_figures: List[VolumeFigure]
        :param key_id_map: KeyIdMap object.
        :type key_id_map: KeyIdMap
        :return: None
        :rtype: :class:`NoneType`
        :Usage example:

         .. code-block:: python

            import supervisely as sly

            os.environ['SERVER_ADDRESS'] = 'https://app.supervisely.com'
            os.environ['API_TOKEN'] = 'Your Supervisely API Token'
            api = sly.Api.from_env()

            volume_id = 19371414
            project_id = 17215

            volume = api.volume.get_info_by_id(volume_id)

            key_id_map = sly.KeyIdMap()
            project_meta_json = api.project.get_meta(project_id)
            project_meta = sly.ProjectMeta.from_json(project_meta_json)

            vol_ann_json = api.volume.annotation.download(volume_id)
            id_to_paths = {}
            vol_ann = sly.VolumeAnnotation.from_json(vol_ann_json, project_meta, key_id_map)
            sp_figures = vol_ann.spatial_figures

            res = volume_figure_api.upload_stl_meshes(volume_id, sp_figures, key_id_map)
        """

        if len(spatial_figures) == 0:
            return

        figure2bytes = {}
        for sp in spatial_figures:
            figure_id = key_id_map.get_figure_id(sp.key())
            if interpolation_dir is not None:
                meth_path = os.path.join(interpolation_dir, sp.key().hex + ".stl")
                if file_exists(meth_path):
                    with open(meth_path, "rb") as in_file:
                        meth_bytes = in_file.read()
                    figure2bytes[figure_id] = meth_bytes
            # else - no stl file
            if figure_id not in figure2bytes:
                meth_bytes = self.interpolate(volume_id, sp, key_id_map)
                figure2bytes[figure_id] = meth_bytes
        self._upload_meshes_batch(figure2bytes)

    def _append_bulk_mask3d(
        self,
        entity_id: int,
        figures: List[VolumeFigure],
        figures_keys: List,
        key_id_map: KeyIdMap,
        field_name=ApiField.ENTITY_ID,
    ):
        """The same method as _append_bulk but for spatial figures. Uploads figures to given Volume by ID.
        You need to upload the geometry right after figures will be created

        :param entity_id: Volume ID.
        :type entity_id: int
        :param figures_json: List of figure dicts.
        :type figures_json: list
        :param figures_keys: List of figure keys as UUID.
        :type figures_keys: list
        :param key_id_map: KeyIdMap object (dict with bidict values)
        :type key_id_map: KeyIdMap
        :param field_name: field name for request body
        :type field_name: str
        :rtype: :class:`NoneType`
        :Usage example:
        """

        if len(figures) == 0:
            return

        empty_figures = []
        for figure in figures:
            empty_figures.append(
                {
                    ApiField.OBJECT_ID: key_id_map.get_object_id(figure.volume_object.key()),
                    ApiField.GEOMETRY_TYPE: Mask3D.name(),
                    ApiField.LABELING_TOOL: Mask3D.name(),
                    ApiField.ENTITY_ID: entity_id,
                    ApiField.CUSTOM_DATA: figure.custom_data
                }
            )
        for batch_keys, batch_jsons in zip(
            batched(figures_keys, batch_size=100), batched(empty_figures, batch_size=100)
        ):
            resp = self._api.post(
                "figures.bulk.add",
                {field_name: entity_id, ApiField.FIGURES: batch_jsons},
            )
            for key, resp_obj in zip(batch_keys, resp.json()):
                figure_id = resp_obj[ApiField.ID]
                key_id_map.add_figure(key, figure_id)

                for figure in figures:
                    if figure.key() == key:
                        geometry_data = figure.geometry.data
                        header = figure.geometry.create_header()
                        geometry_bytes = encode(geometry_data.astype(uint8), header)
                        self.upload_sf_geometries([key], {key: geometry_bytes}, key_id_map)

    def upload_sf_geometries(
        self,
        spatial_figures: List[UUID],
        geometries: Dict[UUID, bytes],
        key_id_map: KeyIdMap,
    ):
        """
        Upload geometries as bytes into spatial figures in the project using their UUID keys.

        :param spatial_figures: List of UUID keys representing spatial figures.
        :type spatial_figures: List[UUID]
        :param geometries: Dictionary where keys are UUIDs of spatial figures, and values are geometries represented as NRRD files in byte format.
        :type geometries: Dict[UUID, bytes]
        :param key_id_map: The KeyIdMap object (a dictionary with bidict values).
        :type key_id_map: KeyIdMap
        :return: None
        :rtype: NoneType

        :Usage example:
        .. code-block:: python

            import numpy as np
            import supervisely as sly
            from supervisely.volume.nrrd_encoder import encode

            os.environ['SERVER_ADDRESS'] = 'https://app.supervisely.com'
            os.environ['API_TOKEN'] = 'Your Supervisely API Token'
            api = sly.Api.from_env()

            volume_id = 23772225
            project_id = 28159
            geometries = {}
            key_id_map = sly.KeyIdMap()
            geometry_bytes = encode(np.random.randint(2, size=(20, 20, 20), dtype=np.uint8))

            vol_ann_json = api.volume.annotation.download(volume_id)
            project_meta = sly.ProjectMeta.from_json(api.project.get_meta(project_id))
            ann = sly.VolumeAnnotation.from_json(vol_ann_json, project_meta, key_id_map)
            spatial_figures = [sp_figure.key() for sp_figure in ann.spatial_figures]
            for figure in spatial_figures:
                geometries[figure] = geometry_bytes
            api.volume.figure.upload_sf_geometries(spatial_figures, geometries, key_id_map)
        """

        if len(spatial_figures) == 0:
            return

        for sf in spatial_figures:
            figure_id = key_id_map.get_figure_id(sf)
            geometry_bytes = geometries.get(sf)
            content_dict = {
                ApiField.FIGURE_ID: str(figure_id),
                ApiField.GEOMETRY: (str(figure_id), geometry_bytes, "application/sla"),
            }
            encoder = MultipartEncoder(fields=content_dict)
            self._api.post("figures.bulk.upload.geometry", encoder)

    def upload_sf_geometry(self, spatial_figures: Dict, geometries: List, key_id_map: KeyIdMap):
        """
        Upload spatial figures geometry as bytes to storage by given ID.

        :param spatial_figures: Dictionary with figure IDs.
        :type spatial_figures: dict
        :param geometries: Dictionary with geometries, which represented as content of NRRD files in byte format.
        :type geometries: list
        :param key_id_map: KeyIdMap object (dict with bidict values)
        :type key_id_map: KeyIdMap
        :rtype: :class:`NoneType`
        :Usage example:
        """

        for sf, geometry_bytes in zip(spatial_figures, geometries):
            figure_id = key_id_map.get_figure_id(sf.key())
            content_dict = {
                ApiField.FIGURE_ID: str(figure_id),
                ApiField.GEOMETRY: (str(figure_id), geometry_bytes, "application/sla"),
            }
            encoder = MultipartEncoder(fields=content_dict)
            self._api.post("figures.bulk.upload.geometry", encoder)

    def download_sf_geometries(self, ids: List[int], paths: List[str]):
        """
        Download spatial figure geometries for the specified figure IDs and save them to the specified paths.

        :param ids: List of VolumeFigure IDs in Supervisely.
        :type ids: List[int]
        :param paths: List of paths to save the downloaded geometries.
        :type paths: List[str]
        :return: None
        :rtype: NoneType

        :Usage example:

         .. code-block:: python

            import supervisely as sly

            os.environ['SERVER_ADDRESS'] = 'https://app.supervisely.com'
            os.environ['API_TOKEN'] = 'Your Supervisely API Token'
            api = sly.Api.from_env()

            STORAGE_DIR = sly.app.get_data_dir()

            volume_id = 19371414
            project_id = 17215

            volume = api.volume.get_info_by_id(volume_id)

            key_id_map = sly.KeyIdMap()
            project_meta_json = api.project.get_meta(project_id)
            project_meta = sly.ProjectMeta.from_json(project_meta_json)

            vol_ann_json = api.volume.annotation.download(volume_id)
            id_to_paths = {}
            vol_ann = sly.VolumeAnnotation.from_json(vol_ann_json, project_meta, key_id_map)

            for sp_figure in vol_ann.spatial_figures:
                figure_id = key_id_map.get_figure_id(sp_figure.key())
                id_to_paths[figure_id] = f"{STORAGE_DIR}/{sp_figure.key().hex}.stl"
            if id_to_paths:
                api.volume.figure.download_sf_geometries(*zip(*id_to_paths.items()))
        """

        if len(ids) == 0:
            return
        if len(ids) != len(paths):
            raise RuntimeError('Can not match "ids" and "paths" lists, len(ids) != len(paths)')

        id_to_path = {id: path for id, path in zip(ids, paths)}
        for figure_id, resp_part in self._download_geometries_batch(ids):
            ensure_base_path(id_to_path[figure_id])
            with open(id_to_path[figure_id], "wb") as w:
                w.write(resp_part.content)

    def load_sf_geometry(self, spatial_figure: VolumeFigure, key_id_map: KeyIdMap):
        """
        Download geometry of an existing in Supervisely figure
        and load this data into the VolumeFigure object to complete this figure representation.

        :param spatial_figure: The spatial figure object from VolumeAnnotation.
        :type spatial_figure: VolumeFigure
        :param key_id_map: The mapped keys and IDs.
        :type key_id_map: KeyIdMap object
        """
        with tempfile.TemporaryDirectory() as temp_dir:
            figure_id = key_id_map.get_figure_id(spatial_figure.key())
            figure_path = f"{temp_dir}/{spatial_figure.key().hex}.nrrd"
            self.download_sf_geometries([figure_id], [figure_path])
            geometry = Mask3D.create_from_file(figure_path)
            spatial_figure._set_3d_geometry(geometry)

    def download(
        self, dataset_id: int, volume_ids: List[int] = None, skip_geometry: bool = False, **kwargs
    ) -> Dict[int, List[FigureInfo]]:
        """
        Method returns a dictionary with pairs of volume ID and list of FigureInfo for the given dataset ID. Can be filtered by volume IDs.

        :param dataset_id: Dataset ID in Supervisely.
        :type dataset_id: int
        :param volume_ids: Specify the list of volume IDs within the given dataset ID. If volume_ids is None, the method returns all possible pairs of images with figures. Note: Consider using `sly.batched()` to ensure that no figures are lost in the response.
        :type volume_ids: List[int], optional
        :param skip_geometry: Skip the download of figure geometry. May be useful for a significant api request speed increase in the large datasets.
        :type skip_geometry: bool

        :return: A dictionary where keys are volume IDs and values are lists of figures.
        :rtype: :class: `Dict[int, List[FigureInfo]]`
        """
        if kwargs.get("image_ids", False) is not False:
            volume_ids = kwargs["image_ids"]  # backward compatibility
        return super().download(dataset_id, volume_ids, skip_geometry)

    def update_custom_data(
        self,
        figure_id: int,
        custom_data: Dict[str, str],
        update_strategy: Literal["replace", "merge"] = "merge",
    ) -> None:
        """
        Update custom data for a specific figure in a volume.

        :param figure_id: ID of the figure to update.
        :type figure_id: int
        :param custom_data: Custom data to update.
        :type custom_data: Dict[str, str]
        :param update_strategy: Strategy to apply, either "replace" or "merge".
        :type update_strategy: Literal["replace", "merge"]
        :return: None
        :rtype: :class:`NoneType`

        """
        data = {
            ApiField.ID: figure_id,
            ApiField.CUSTOM_DATA: custom_data,
            ApiField.UPDATE_STRATEGY: update_strategy,
        }
        self._api.post("figures.custom-data.update", data)
