#  Copyright © 2025 Bentley Systems, Incorporated
#  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 functools
from collections.abc import Callable, Iterator, Sequence
from threading import Lock
from typing import TypeVar

from ..interfaces import IFeedback

__all__ = [
    "NoFeedback",
    "PartialFeedback",
    "create_default_feedback",
    "iter_with_fb",
    "reset_feedback_factory",
    "set_feedback_factory",
    "split_feedback",
]

_N_DIGITS = 4  # Let's maintain a user-friendly number of dp.


class _NoFeedback(IFeedback):
    def progress(self, progress: float, message: str | None = None) -> None:
        pass


NoFeedback = _NoFeedback()
"""A default feedback object that does nothing. Use this when no feedback is needed."""


def _no_feedback_factory(_label: str) -> IFeedback:
    return NoFeedback


_feedback_factory: Callable[[str], IFeedback] = _no_feedback_factory


def set_feedback_factory(factory: Callable[[str], IFeedback]) -> None:
    """Set the factory used by :func:`create_default_feedback`.

    This is intended to be called by the IPython extension (``%load_ext evo.widgets``)
    so that SDK operations automatically display a progress widget when running in a notebook.

    :param factory: A callable that accepts a label string and returns an :class:`IFeedback` instance.
    """
    global _feedback_factory
    _feedback_factory = factory


def reset_feedback_factory() -> None:
    """Reset the feedback factory to the default (returns :data:`NoFeedback`)."""
    global _feedback_factory
    _feedback_factory = _no_feedback_factory


def create_default_feedback(label: str) -> IFeedback:
    """Create a feedback instance using the currently registered factory.

    When no factory has been registered (or after :func:`reset_feedback_factory`),
    this returns :data:`NoFeedback`.

    :param label: A descriptive label for the operation (e.g. ``"Tasks"``, ``"Uploading data"``).
    :returns: An :class:`IFeedback` instance.
    """
    return _feedback_factory(label)


class PartialFeedback(IFeedback):
    """A wrapper for IFeedback objects that subdivides the feedback range."""

    def __init__(self, parent: IFeedback, start: float, end: float) -> None:
        """
        :param parent: The parent feedback object to subdivide.
        :param start: The start of the partial feedback range.
        :param end: The end of the partial feedback range.
        """
        self.__offset = start
        self.__factor = end - start

        self.__lock = Lock()
        self.__parent = parent

    def progress(self, progress: float, message: str | None = None) -> None:
        """Update the progress of the feedback object.

        :param progress: The progress value, between 0 and 1.
        :param message: An optional message to display.
        """
        partial_progress = round(self.__offset + (progress * self.__factor), ndigits=_N_DIGITS)
        with self.__lock:
            self.__parent.progress(partial_progress, message)


T = TypeVar("T")


def iter_with_fb(elements: Sequence[T], feedback: IFeedback | None = None) -> Iterator[tuple[T, IFeedback]]:
    """Iterate over a sequence of elements, dividing feedback uniformly throughout the range.

    :param elements: The sequence of elements to iterate over.
    :param feedback: A feedback object to subdivide.

    :yields: A tuple containing a sequence element, and a new feedback object for that element.
    """
    if len(elements) == 0:
        return

    fb_part_size = 1 / len(elements)
    for i, element in enumerate(elements):
        if feedback is None:
            yield element, NoFeedback
        else:
            start = round(i * fb_part_size, ndigits=_N_DIGITS)
            end = round((i + 1) * fb_part_size, ndigits=_N_DIGITS)
            yield element, PartialFeedback(feedback, start, end)
            feedback.progress(end)


class _ConcurrentFeedback(IFeedback):
    """A feedback object that reports progress to a ConcurrentFeedbackGroup, allowing aggregation of progress from
    multiple sources concurrently.
    """

    def __init__(self, callback: Callable[[float, str | None], None]) -> None:
        self.__callback = callback
        self.__progress = 0.0

    def progress(self, progress: float, message: str | None = None) -> None:
        """Update the progress of the feedback object.

        :param progress: The progress value, between 0 and 1.
        :param message: An optional message to display.
        """
        diff = progress - self.__progress
        self.__progress = progress
        self.__callback(diff, message)


class _ConcurrentFeedbackGroup:
    """A group of feedbacks that can be updated concurrently, aggregating their progress."""

    def __init__(self, parent: IFeedback) -> None:
        self.__lock = Lock()
        self.__parent = parent
        self.__total_weight = 0.0
        self.__progress = 0.0
        self.__n_feedbacks = 0

    def _update(self, message: str | None = None) -> None:
        if self.__total_weight == 0:
            progress = self.__progress / self.__n_feedbacks
        else:
            progress = self.__progress / self.__total_weight
        self.__parent.progress(round(progress, _N_DIGITS), message)

    def _progress(self, weight: float, progress_change: float, message: str | None) -> None:
        with self.__lock:
            if self.__total_weight == 0:
                self.__progress += progress_change
            else:
                self.__progress += weight * progress_change
            self._update(message)

    def create_feedback(self, weight: float) -> IFeedback:
        with self.__lock:
            if self.__progress > 0:
                raise RuntimeError("Cannot create new feedback after progress has been reported.")
            self.__total_weight += weight
            self.__n_feedbacks += 1

        return _ConcurrentFeedback(functools.partial(self._progress, weight))


def split_feedback(feedback: IFeedback, weights: Sequence[float]) -> list[IFeedback]:
    """Split a feedback object into multiple feedback objects based on weights.

    Each of them can be updated concurrently, and their progress will be aggregated into the parent feedback object.

    Note, if all weights are zero, all feedback objects will be uniformly weighted.

    :param feedback: The parent feedback object to subdivide.
    :param weights: Weights for each partial feedback object.
    :returns: A list of ConcurrentFeedback objects.
    """
    group = _ConcurrentFeedbackGroup(feedback)
    return [group.create_feedback(weight) for weight in weights]
