import os
from typing import Dict, List, Optional, Set, Tuple
from uuid import UUID

import magic
import numpy as np

from supervisely import (
    Api,
    PointcloudAnnotation,
    batched,
    generate_free_name,
    is_development,
    logger,
)
from supervisely.api.module_api import ApiField
from supervisely.convert.base_converter import BaseConverter
from supervisely.io.fs import get_file_ext, get_file_name_with_ext
from supervisely.io.json import load_json_file
from supervisely.pointcloud.pointcloud import ALLOWED_POINTCLOUD_EXTENSIONS
from supervisely.pointcloud.pointcloud import validate_ext as validate_pcd_ext
from supervisely.pointcloud_annotation.constants import OBJECT_KEY
from supervisely.video_annotation.key_id_map import KeyIdMap

CONVERTIBLE_EXTENSIONS = [".bin"]


class PointcloudConverter(BaseConverter):
    """Base converter for pointcloud projects (uploads pointclouds, annotations, and optional related images)."""

    allowed_exts = ALLOWED_POINTCLOUD_EXTENSIONS + CONVERTIBLE_EXTENSIONS
    modality = "pointclouds"

    class Item(BaseConverter.BaseItem):
        """Base pointcloud item: path + optional annotation data and related images bundle."""

        def __init__(
            self,
            item_path,
            ann_data=None,
            related_images: Optional[list] = None,
            custom_data: Optional[dict] = None,
        ):
            """
            :param item_path: Path to pointcloud file.
            :type item_path: str
            :param ann_data: Annotation path or data.
            :type ann_data: str, optional
            :param related_images: List of related image tuples.
            :type related_images: list, optional
            :param custom_data: Extra per-item data.
            :type custom_data: dict, optional
            """
            self._name: str = None
            self._path = item_path
            self._ann_data = ann_data
            self._type = "point_cloud"
            self._related_images = related_images if related_images is not None else []
            self._custom_data = custom_data if custom_data is not None else {}

        def create_empty_annotation(self) -> PointcloudAnnotation:
            return PointcloudAnnotation()

        def set_related_images(self, related_images: Tuple[str, str, Optional[str]]) -> None:
            """Adds related image to the item.

            related_images tuple:
                - path to image
                - path to .json with image metadata
                - path to .figures.json (can be None if no figures)
            """
            self._related_images.append(related_images)

    @property
    def format(self):
        return self._converter.format

    @property
    def ann_ext(self):
        return None

    @property
    def key_file_ext(self):
        return None

    @staticmethod
    def validate_ann_file(ann_path, meta=None):
        return False

    def upload_dataset(
        self,
        api: Api,
        dataset_id: int,
        batch_size: int = 1,
        log_progress=True,
    ):
        """Upload converted data to Supervisely"""

        meta, renamed_classes, renamed_tags = self.merge_metas_with_conflicts(api, dataset_id)

        existing_pcd_infos = api.pointcloud.get_list(dataset_id)
        existing_pointcloud_names = set([pcd.name for pcd in existing_pcd_infos])
        used_related_image_names: Set[str] = set()
        try:
            existing_pcd_ids = [pcd.id for pcd in existing_pcd_infos]
            for pcd_ids_batch in batched(existing_pcd_ids, batch_size=200):
                related_images = api.pointcloud.get_list_related_images_batch(
                    dataset_id, pcd_ids_batch
                )
                for related_image in related_images:
                    related_image_name = related_image.get(ApiField.NAME)
                    if related_image_name is not None:
                        used_related_image_names.add(related_image_name)
        except Exception as e:
            logger.debug(
                f"Failed to fetch existing related image names for dataset ID:{dataset_id}: {repr(e)}"
            )

        if log_progress:
            progress, progress_cb = self.get_progress(self.items_count, "Uploading pointclouds...")
        else:
            progress_cb = None

        upload_fn = api.pointcloud.upload_paths
        if self.upload_as_links and self.supports_links:
            upload_fn = api.pointcloud.upload_links

        key_id_map = KeyIdMap()
        for batch in batched(self._items, batch_size=batch_size):
            item_names = []
            item_paths = []
            anns = []
            for item in batch:
                item.name = generate_free_name(
                    existing_pointcloud_names, item.name, with_ext=True, extend_used_names=True
                )
                item_names.append(item.name)
                if self.upload_as_links:
                    item_paths.append(
                        self.remote_files_map.get(os.path.abspath(item.path), item.path)
                    )
                else:
                    item_paths.append(item.path)

                ann = self.to_supervisely(item, meta, renamed_classes, renamed_tags)
                anns.append(ann)

            pcd_infos = upload_fn(
                dataset_id,
                item_names,
                item_paths,
            )
            pcd_ids = [pcd_info.id for pcd_info in pcd_infos]
            pcl_to_rimg_figures: Dict[int, Dict[str, List[Dict]]] = {}
            pcl_to_rimg_to_id: Dict[int, Dict[str, int]] = {}
            for pcd_id, ann, item in zip(pcd_ids, anns, batch):
                if ann is not None:
                    api.pointcloud.annotation.append(pcd_id, ann, key_id_map)

                rimg_infos = []
                camera_names = []
                for img_ind, rel_tuple in enumerate(item._related_images):
                    img_path = rel_tuple[0]
                    rimg_ann_path = rel_tuple[1]
                    fig_path = rel_tuple[2] if len(rel_tuple) > 2 else None
                    meta_json = load_json_file(rimg_ann_path)
                    try:
                        if ApiField.META not in meta_json:
                            raise ValueError("Related image meta not found in json file.")
                        if ApiField.NAME not in meta_json:
                            raise ValueError("Related image name not found in json file.")

                        img_hash = None
                        if not self.upload_as_links:
                            img_hash = api.pointcloud.upload_related_image(img_path)

                        if "deviceId" not in meta_json[ApiField.META].keys():
                            camera_names.append(f"CAM_{str(img_ind).zfill(2)}")
                        else:
                            camera_names.append(meta_json[ApiField.META]["deviceId"])

                        related_image_name = generate_free_name(
                            used_related_image_names,
                            meta_json[ApiField.NAME],
                            with_ext=True,
                            extend_used_names=True,
                        )

                        rimage_dict = {
                            ApiField.ENTITY_ID: pcd_id,
                            ApiField.NAME: related_image_name,
                            ApiField.META: meta_json[ApiField.META],
                        }
                        link = None
                        if img_hash is not None:
                            rimage_dict[ApiField.HASH] = img_hash
                        else:
                            link = self.remote_files_map.get(os.path.abspath(img_path), img_path)
                            rimage_dict[ApiField.LINK] = link
                        img_data = img_hash or link
                        rimg_infos.append(rimage_dict)

                        if fig_path is not None and os.path.isfile(fig_path):
                            try:
                                figs_json = load_json_file(fig_path)
                                pcl_to_rimg_figures.setdefault(pcd_id, {})[img_data] = figs_json
                            except Exception as e:
                                logger.debug(f"Failed to read figures json '{fig_path}': {repr(e)}")

                    except Exception as e:
                        logger.warning(
                            f"Failed to upload related image or add it to pointcloud: {repr(e)}"
                        )
                        continue

                # add images for this point cloud
                if len(rimg_infos) > 0:
                    try:
                        uploaded_rimgs = api.pointcloud.add_related_images(rimg_infos, camera_names)
                        # build mapping hash->id
                        for info, uploaded in zip(rimg_infos, uploaded_rimgs):
                            img_data = info.get(ApiField.HASH, info.get(ApiField.LINK, None))
                            img_id = (
                                uploaded.get(ApiField.ID)
                                if isinstance(uploaded, dict)
                                else getattr(uploaded, "id", None)
                            )
                            if img_data is not None and img_id is not None:
                                pcl_to_rimg_to_id.setdefault(pcd_id, {})[img_data] = img_id
                    except Exception as e:
                        logger.debug(f"Failed to add related images to pointcloud: {repr(e)}")

            # ---- upload figures for processed batch ----
            if len(pcl_to_rimg_figures) > 0:
                try:
                    dataset_info = api.dataset.get_info_by_id(dataset_id)
                    project_id = dataset_info.project_id

                    figures_payload: List[Dict] = []

                    for pcl_id, rimg_data_to_figs in pcl_to_rimg_figures.items():
                        rimg_data_to_id = pcl_to_rimg_to_id.get(pcl_id, {})
                        if len(rimg_data_to_id) == 0:
                            continue

                        for img_data, figs_json in rimg_data_to_figs.items():
                            if img_data not in rimg_data_to_id:
                                continue
                            rimg_id = rimg_data_to_id[img_data]
                            for fig in figs_json:
                                try:
                                    fig[ApiField.ENTITY_ID] = rimg_id
                                    fig[ApiField.DATASET_ID] = dataset_id
                                    fig[ApiField.PROJECT_ID] = project_id
                                    if OBJECT_KEY in fig:
                                        fig[ApiField.OBJECT_ID] = key_id_map.get_object_id(
                                            UUID(fig[OBJECT_KEY])
                                        )
                                except Exception as e:
                                    logger.debug(
                                        f"Failed to process figure json for current image. Hash / link = {img_data}: {repr(e)}"
                                    )
                                    continue

                            figures_payload.extend(figs_json)

                        if len(figures_payload) > 0:
                            try:
                                api.image.figure.create_bulk(
                                    figures_json=figures_payload, dataset_id=dataset_id
                                )
                            except Exception as e:
                                logger.debug(
                                    f"Failed to upload figures for related images: {repr(e)}"
                                )
                except Exception as e:
                    logger.debug(f"Unexpected error during related image figures upload: {repr(e)}")

            if log_progress:
                progress_cb(len(batch))

        if log_progress:
            if is_development():
                progress.close()
        logger.info(f"Dataset ID:{dataset_id} has been successfully uploaded.")

    def _collect_items_if_format_not_detected(self) -> Tuple[List[Item], bool, Set[str]]:
        only_modality_items = True
        unsupported_exts = set()
        pcd_list, rimg_dict, rimg_ann_dict, rimg_fig_dict = [], {}, {}, {}
        for root, _, files in os.walk(self._input_data):
            for file in files:
                full_path = os.path.join(root, file)
                if file in ["key_id_map.json", "meta.json"]:
                    continue

                ext = get_file_ext(full_path)
                if file.endswith(".figures.json"):
                    rimg_fig_dict[file] = full_path
                elif ext == ".json":
                    dir_name = os.path.basename(root)
                    parent_dir_name = os.path.basename(os.path.dirname(root))
                    if any(
                        p.replace("_", " ") in ["images", "related images", "photo context"]
                        for p in [dir_name, parent_dir_name]
                    ) or dir_name.endswith("_pcd"):
                        rimg_ann_dict[file] = full_path
                elif self._is_image_file(full_path):
                    dir_name = os.path.basename(root)
                    if dir_name not in rimg_dict:
                        rimg_dict[dir_name] = []
                    rimg_dict[dir_name].append(full_path)
                elif ext.lower() in self.allowed_exts:
                    try:
                        path = self._convert_to_pcd_if_needed(full_path, ext)
                        if not path:
                            continue
                        pcd_list.append(path)
                    except:
                        pass
                else:
                    only_modality_items = False
                    unsupported_exts.add(ext)

        # create Items
        items = []
        for pcd_path in pcd_list:
            item = self.Item(pcd_path)
            rimg_dir_name = item.name.replace(".pcd", "_pcd")
            rimgs = rimg_dict.get(rimg_dir_name, [])
            for rimg_path in rimgs:
                rimg_ann_name = f"{get_file_name_with_ext(rimg_path)}.json"
                if rimg_ann_name in rimg_ann_dict:
                    rimg_ann_path = rimg_ann_dict[rimg_ann_name]
                    rimg_fig_name = f"{get_file_name_with_ext(rimg_path)}.figures.json"
                    rimg_fig_path = rimg_fig_dict.get(rimg_fig_name, None)
                    if rimg_fig_path is not None and not os.path.exists(rimg_fig_path):
                        rimg_fig_path = None
                    item.set_related_images((rimg_path, rimg_ann_path, rimg_fig_path))
            items.append(item)
        return items, only_modality_items, unsupported_exts

    def _convert_to_pcd_if_needed(self, pcd_path: str, ext: str) -> Optional[str]:
        """Convert point cloud to .pcd format if it is in another supported format."""
        if ext == ".pcd":
            return pcd_path
        elif ext == ".bin":
            if self.upload_as_links:
                logger.warning(
                    f"Uploading .bin pointcloud files as links is not supported. "
                    f"Please convert .bin files to .pcd format before uploading. "
                    f"Skipping conversion for: {pcd_path}"
                )
                return None
            if not self._validate_bin_pointcloud(pcd_path):
                logger.warning(
                    f"The .bin pointcloud file '{pcd_path}' is not valid. Skipping conversion to .pcd."
                )
                return None
            pcd_output_path = pcd_path.replace(".bin", ".pcd")
            self._convert_bin_to_pcd(pcd_path, pcd_output_path)
            return pcd_output_path

    def _validate_bin_pointcloud(self, bin_file: str) -> bool:
        try:
            data = np.fromfile(bin_file, dtype=np.float32)
            if data.size == 0:
                return False

            if data.size % 5 == 0:
                points = data.reshape(-1, 5)
                return bool(np.isfinite(points).all())

            if data.size % 4 == 0:
                points = data.reshape(-1, 4)
                return bool(np.isfinite(points).all())

            return False
        except Exception:
            return False

    def _convert_bin_to_pcd(self, bin_file: str, pcd_file: str) -> None:
        """Convert .bin point cloud file to .pcd format using open3d library (supports 4 or 5 floats per point)."""
        try:
            import open3d as o3d
        except ImportError:
            raise ImportError(
                "open3d is not installed. Cannot convert .bin point cloud to .pcd format. "
                "Please install open3d package to enable this feature."
            )

        try:
            raw = np.fromfile(bin_file, dtype=np.float32)
            if raw.size % 5 == 0:
                arr = raw.reshape(-1, 5)
                points = arr[:, 0:3]
                intensity = arr[:, 3]
            elif raw.size % 4 == 0:
                arr = raw.reshape(-1, 4)
                points = arr[:, 0:3]
                intensity = arr[:, 3]
            else:
                raise ValueError("The number of float32 values is not a multiple of 4 or 5.")
        except ValueError as e:
            raise Exception(
                f"Incorrect data in the BIN pointcloud file: {bin_file}. "
                f"The number of float32 values must be a multiple of 4 (x,y,z,intensity) "
                f"or 5 (x,y,z,intensity,ring_index). Details: {e} "
                "Please ensure that the binary file contains a valid number of elements to be "
                "successfully reshaped into a (N, 4) or (N, 5) array.\n"
            )
        # normalize intensity to [0, 1] for Open3D colors
        intensity = np.nan_to_num(intensity, nan=0.0, posinf=0.0, neginf=0.0)
        if intensity.size > 0:
            i_min = float(np.min(intensity))
            i_max = float(np.max(intensity))
            if i_max > i_min:
                norm_intensity = (intensity - i_min) / (i_max - i_min)
            else:
                norm_intensity = np.zeros_like(intensity)
        else:
            norm_intensity = intensity
        intensity_fake_rgb = np.zeros((norm_intensity.shape[0], 3))
        intensity_fake_rgb[:, 0] = norm_intensity
        pc = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(points))
        pc.colors = o3d.utility.Vector3dVector(intensity_fake_rgb)
        o3d.io.write_point_cloud(pcd_file, pc)

    @staticmethod
    def _is_image_file(path: str) -> bool:
        try:
            return magic.from_file(path, mime=True).startswith("image/")
        except Exception:
            return False
