# Copyright (C) 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions
# and limitations under the License.

import re
from collections.abc import Callable
from importlib import resources
from typing import Any

import numpy as np
from pathvalidate import sanitize_filepath

from geti_sdk.data_models import (
    Annotation,
    Label,
    OptimizedModel,
    Prediction,
    Project,
    ScoredLabel,
    Task,
    TaskType,
)
from geti_sdk.data_models.shapes import Rectangle
from geti_sdk.deployment.data_models import ROI, IntermediateInferenceResult

try:
    OVMS_README_PATH = str(resources.files("geti_sdk.deployment.resources") / "OVMS_README.md")
except AttributeError:
    with resources.path("geti_sdk.deployment.resources", "OVMS_README.md") as data_path:
        OVMS_README_PATH = str(data_path)


def generate_ovms_model_name(project: Project, model: OptimizedModel, omit_version: bool = True) -> str:
    """
    Generate a valid name for a model to be uploaded to the OpenVINO Model Server.

    :param project: Project the model belongs to
    :param model: The model for which to generate the name
    :param omit_version: True not to include the version of the model in the name.
        This is useful if the name needs to be used to create a directory structure,
        since including the version will break this.
    :return: String containing the name of the model
    """
    model_name = sanitize_filepath(project.name + "_" + model.name)
    model_name = model_name.replace(" ", "_")
    model_name = model_name.replace("-", "_")
    model_name = model_name.lower()
    if model.version is not None and not omit_version:
        model_name += f":{model.version}"
    return model_name


def generate_ovms_model_address(ovms_address: str, model_name: str) -> str:
    """
    Generate a string representing a valid model address for the OpenVINO Model Server,
    from a given `ovms_address` and `model_name`.

    The format expected for model addresses by OVMS is:

        <address>:<port>/models/<model_name>[:<model_version>]

    Where `address` must not contain the protocol (http or https)

    :param ovms_address: IP address or URL pointing to the OVMS instance, including
        port specification
    :param model_name: Name of the model
    :return: String containing the model address
    """
    model_address = f"{ovms_address}/models/{model_name}"
    if model_address.startswith("https://"):
        model_address = model_address[8:]
    if model_address.startswith("http://"):
        model_address = model_address[7:]
    return model_address


def target_device_is_ovms(device: str) -> bool:
    """
    Return True if the target `device` specified is a URL or IP address, False otherwise

    :param device: Target device string to check
    :return: True if the string represents a URL or IP address, False otherwise
    """
    # Check if 'device' has been specified as a URL or IP address.
    server_pattern = re.compile(r"^((https?://)|(www.))(?:([a-zA-Z]+)|(\d+\.\d+\.\d+\.\d+)):\d{1,5}?$")
    return server_pattern.match(device) is not None


def rgb_to_hex(rgb: tuple[int, int, int]) -> str:
    """
    Convert an RGB color value to its corresponding hexadecimal representation.

    :param rgb: A tuple representing the RGB color value, where each element is an integer between 0 and 255.
    :return: The hexadecimal representation of the RGB color value.

    _Example:

        >>> rgb_to_hex((255, 0, 0))
        '#ff0000'
    """
    return f"#{rgb[0]:02x}{rgb[1]:02x}{rgb[2]:02x}"


def assign_empty_label(prediction: Prediction, image: np.ndarray, empty_label: Label | None = None):
    """
    Modify the `prediction` in place to assign an empty label, if the prediction
    does not contain any objects.

    :param prediction: Prediction object to assign the empty label to, if needed
    :param image: The original image for which the prediction was generated
    :param empty_label: The Label object representing the empty label to assign
    """
    # Empty label is not generated by prediction postprocessing correctly, append
    # it here if there are no other predictions
    height, width = image.shape[0:2]
    if len(prediction.annotations) == 0 and empty_label is not None:
        prediction.append(
            Annotation(
                shape=Rectangle(x=0, y=0, width=width, height=height),
                labels=[ScoredLabel.from_label(empty_label, probability=1)],
            )
        )


def flow_control(intermediate_result: IntermediateInferenceResult, task: Task) -> tuple[list[ROI], list[np.ndarray]]:
    """
    Logic for the flow control tasks

    :param intermediate_result: IntermediateInferenceResult object, containing the
        predictions so far, as well as the rois.
    :param task: The task to perform the flow control operation for
    :return: Tuple containing:
        - rois -> A list of rois generated by the flow control task, to be passed
            down to the next task in the task chain
        - image_views -> A list of numpy arrays containing the image data in the
            corresponding roi
    """
    # CROP task
    if task.type == TaskType.CROP:
        rois = intermediate_result.filter_rois(label=None)
        image_views = intermediate_result.generate_views(rois)
    else:
        raise NotImplementedError(
            f"Unable to run inference for the pipeline in the deployed "
            f"project: Unsupported task type {task.type} found."
        )
    return rois, image_views


def pipeline_callback_factory(
    is_first_task: bool,
    is_final_task: bool,
    add_flow_control: bool,
    current_task: Task,
    deployment,
    final_callback: Callable[[Prediction, Any | None], None],
    next_trainable_task: Task | None = None,
    next_task: Task | None = None,
) -> Callable[[Prediction, Any], None]:
    """
    Generate callback functions for pipeline inference
    """
    if is_first_task:

        def callback(result: Prediction, runtime_data: tuple[np.ndarray, Any]):
            """
            Process inference results for the first task in a task chain, this
            initializes the intermediate result
            """
            image, runtime_user_data = runtime_data
            rois = [ROI.from_annotation(annotation) for annotation in result.annotations]

            if len(rois) == 0:
                # Empty prediction. Assign empty label and run final callback immediately
                assign_empty_label(result, image, deployment._empty_label)
                final_callback(image, result, runtime_user_data)
                return

            intermediate_result = IntermediateInferenceResult(image=image, prediction=result, rois=rois)

            if add_flow_control:
                rois, image_views = flow_control(intermediate_result, next_task)
                intermediate_result.rois = rois
            else:
                image_views = intermediate_result.generate_views()
            for roi, view in zip(rois, image_views):
                deployment._infer_task_async(
                    image=view,
                    explain=False,
                    task=next_trainable_task,
                    runtime_data=(image, roi, intermediate_result, runtime_user_data),
                )

    else:

        def callback(
            result: Prediction,
            runtime_data: tuple[np.ndarray, ROI, IntermediateInferenceResult, Any],
        ):
            """
            Process inference results for subsequent tasks
            """
            image, roi, intermediate_result, runtime_user_data = runtime_data

            if current_task.is_global:
                # Global tasks add their labels to the existing shape in the ROI
                intermediate_result.extend_annotations(result.annotations, roi=roi)
            else:
                # Local tasks create new shapes in the image coordinate system,
                # and generate ROI's corresponding to the new shapes
                for annotation in result.annotations:
                    intermediate_result.append_annotation(annotation, roi=roi)
                    new_roi = ROI.from_annotation(annotation)
                    intermediate_result.add_to_infer_queue(new_roi.to_absolute_coordinates(parent_roi=roi))
            intermediate_result.increment_infer_counter()
            if intermediate_result.all_rois_inferred():
                # In this case, all ROIS from the preceding task have been inferred.
                # The infer queue now becomes the new ROIs for the intermediate result
                new_rois = intermediate_result.get_infer_queue()
                intermediate_result.reset_infer_counter()
                intermediate_result.rois = new_rois

                if add_flow_control:
                    rois, image_views = flow_control(intermediate_result, next_task)
                    intermediate_result.rois = rois
                else:
                    image_views = intermediate_result.generate_views()

                if not is_final_task:
                    for roi, view in zip(rois, image_views):
                        deployment._infer_task_async(
                            image=view,
                            explain=False,
                            task=next_trainable_task,
                            runtime_data=(
                                image,
                                roi,
                                intermediate_result,
                                runtime_user_data,
                            ),
                        )
                else:
                    prediction = intermediate_result.prediction
                    final_callback(image, prediction, runtime_user_data)

    return callback
