import numpy as np
import functools
from fastapi import Request, BackgroundTasks
from typing import Any, Dict, List, Optional, Union
from pathlib import Path
import supervisely as sly
from supervisely.annotation.label import Label
from supervisely.nn.prediction_dto import PredictionSegmentation
from supervisely.nn.inference.tracking.tracker_interface import TrackerInterface
from supervisely.nn.inference import Inference
import supervisely.nn.inference.tracking.functional as F
import traceback


class MaskTracking(Inference):
    def __init__(
        self,
        model_dir: Optional[str] = None,
        custom_inference_settings: Optional[Union[Dict[str, Any], str]] = None,
    ):
        super().__init__(
            model_dir,
            custom_inference_settings,
            sliding_window_mode=None,
            use_gui=False,
        )

        try:
            self.load_on_device(model_dir, "cuda")
        except RuntimeError:
            self.load_on_device(model_dir, "cpu")

    def get_info(self):
        info = super().get_info()
        info["task type"] = "tracking"
        return info

    def serve(self):
        super().serve()
        server = self._app.get_server()

        @server.post("/track")
        def start_track(request: Request, task: BackgroundTasks):
            task.add_task(track, request)
            return {"message": "Track task started."}

        def send_error_data(func):
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                value = None
                try:
                    value = func(*args, **kwargs)
                except Exception as exc:
                    # print("An error occured:")
                    # print(traceback.format_exc())
                    request: Request = args[0]
                    context = request.state.context
                    api: sly.Api = request.state.api
                    track_id = context["trackId"]
                    api.logger.error("An error occured:")
                    api.logger.exception(exc)

                    api.post(
                        "videos.notify-annotation-tool",
                        data={
                            "type": "videos:tracking-error",
                            "data": {
                                "trackId": track_id,
                                "error": {"message": repr(exc)},
                            },
                        },
                    )
                return value

            return wrapper

        @send_error_data
        def track(request: Request):
            context = request.state.context
            api: sly.Api = request.state.api

            self.video_interface = TrackerInterface(
                context=context,
                api=api,
                load_all_frames=True,
                notify_in_predict=True,
                per_point_polygon_tracking=False,
            )
            api.logger.info("Starting tracking process")
            # load frames
            frames = self.video_interface.frames
            # combine several binary masks into one multilabel mask
            i = 0
            label2id = {}
            frames_generator = [_ for _ in self.video_interface.frames_loader_generator()]
            for (fig_id, geometry), obj_id in zip(
                self.video_interface.geometries.items(),
                self.video_interface.object_ids,
            ):
                original_geometry = geometry.clone()
                if not isinstance(geometry, sly.Bitmap) and not isinstance(geometry, sly.Polygon):
                    raise TypeError(
                        f"This app does not support {geometry.geometry_name()} tracking"
                    )
                # convert polygon to bitmap
                if isinstance(geometry, sly.Polygon):
                    polygon_obj_class = sly.ObjClass("polygon", sly.Polygon)
                    polygon_label = sly.Label(geometry, polygon_obj_class)
                    bitmap_obj_class = sly.ObjClass("bitmap", sly.Bitmap)
                    bitmap_label = polygon_label.convert(bitmap_obj_class)[0]
                    geometry = bitmap_label.geometry
                if i == 0:
                    multilabel_mask = geometry.data.astype(int)
                    multilabel_mask = np.zeros(frames[0].shape, dtype=np.uint8)
                    geometry.draw(bitmap=multilabel_mask, color=[1, 1, 1])
                    i += 1
                else:
                    i += 1
                    geometry.draw(bitmap=multilabel_mask, color=[i, i, i])
                label2id[i] = {
                    "fig_id": fig_id,
                    "obj_id": obj_id,
                    "original_geometry": original_geometry.geometry_name(),
                }
            # run tracker
            tracked_multilabel_masks = self.predict(
                frames=frames, input_mask=multilabel_mask[:, :, 0]
            )
            tracked_multilabel_masks = np.array(tracked_multilabel_masks)
            # decompose multilabel masks into binary masks
            for i in np.unique(tracked_multilabel_masks):
                if i != 0:
                    binary_masks = tracked_multilabel_masks == i
                    fig_id = label2id[i]["fig_id"]
                    obj_id = label2id[i]["obj_id"]
                    geometry_type = label2id[i]["original_geometry"]
                    for j, mask in enumerate(binary_masks[1:]):
                        if geometry_type == "polygon":
                            bitmap_geometry = sly.Bitmap(mask)
                            bitmap_obj_class = sly.ObjClass("bitmap", sly.Bitmap)
                            bitmap_label = sly.Label(bitmap_geometry, bitmap_obj_class)
                            polygon_obj_class = sly.ObjClass("polygon", sly.Polygon)
                            polygon_labels = bitmap_label.convert(polygon_obj_class)
                            geometries = [label.geometry for label in polygon_labels]
                        else:
                            geometries = [sly.Bitmap(mask)]
                        for l, geometry in enumerate(geometries):
                            if l == 0:
                                notify = True
                            else:
                                notify = False
                            self.video_interface.add_object_geometry_on_frame(
                                geometry,
                                obj_id,
                                self.video_interface._cur_frames_indexes[j + 1],
                                notify=notify,
                            )
                    if self.video_interface.global_stop_indicatior:
                        return
                    api.logger.info(f"Figure with id {fig_id} was successfully tracked")
