import asyncio
from typing import Callable

from examples.cambricon_custom_func.vllm.mlu_hijack.common import init_logger

logger = init_logger(__name__)


class PriorityLock:
    """
    A lock class that prioritizes tasks based on their priority level and supports dynamic
    updating of priority thresholds after each lock release.

    Attributes:
    -----------
    _lock : asyncio.Lock
        An internal asyncio lock used to ensure mutual exclusion.
    _queue : asyncio.PriorityQueue
        A priority queue to store tasks by their priority. Tasks with lower numerical priority
        values have higher priority.
    _condition : asyncio.Condition
        A condition variable to manage the waiting and notification of tasks.
    _active_task : asyncio.Task or None
        Tracks the task currently holding the lock, or None if the lock is not held.
    _current_priority_threshold : int
        The current priority threshold for tasks allowed to acquire the lock.
    _priority_interval : int
        The value by which the priority threshold is incremented after a lock release when
        `update_priority_threshold` is enabled.
    """

    def __init__(self, init_priority_threshold: int, priority_interval: int):
        """
        Initializes a PriorityLock with an initial priority threshold and interval.

        Parameters:
        -----------
        init_priority_threshold : int
            The initial threshold for task priorities that can acquire the lock.
        priority_interval : int
            The interval by which the priority threshold increases after each lock release.
        """
        self._lock = asyncio.Lock()  # Internal asyncio lock
        self._queue = asyncio.PriorityQueue()  # Priority queue to manage tasks by priority
        self._condition = asyncio.Condition()  # Condition variable to manage waiting tasks
        self._active_task = None  # Keep track of the current active task holding the lock
        self._current_priority_threshold = init_priority_threshold
        self._priority_interval = priority_interval

    async def acquire(self, priority):
        """
        Acquires the lock for a task based on its priority.

        Parameters:
        -----------
        priority : int
            The priority level of the task attempting to acquire the lock.

        Behavior:
        ---------
        - The task is enqueued based on its priority.
        - The task waits until it is the highest-priority task in the queue, has a priority
          below the current threshold, and the lock is available.
        """
        queue_item = (priority, asyncio.current_task())
        async with self._condition:
            await self._queue.put(queue_item)

            # Wait until the current task is the one with the highest priority and the lock is available
            while True:
                # Check if the current task is at the front of the queue and the lock is available
                current_priority, current_task = self._queue._queue[0]  # Peek at the highest priority task
                if current_priority < self._current_priority_threshold and current_task is asyncio.current_task() and not self._lock.locked():
                    await self._lock.acquire()  # Acquire the lock
                    self._active_task = current_task  # Mark the current task as holding the lock
                    await self._queue.get()  # Remove the task from the queue
                    break
                # If not the highest priority task, wait until notified
                await self._condition.wait()

    async def release(self, update_priority_threshold):
        """
        Releases the lock, optionally updating the priority threshold.

        Parameters:
        -----------
        update_priority_threshold : bool
            If True, increments the priority threshold by the configured interval.
        """
        # Notify waiting tasks that the lock has been released
        async with self._condition:
            self._active_task = None  # Clear the reference to the current task
            self._lock.release()

            if update_priority_threshold:
                self._current_priority_threshold += self._priority_interval
            self._condition.notify_all()  # Wake up all waiting tasks to recheck their priority

    async def __aenter__(self, priority):
        """
        Async context manager entry. Acquires the lock with the specified priority.

        Parameters:
        -----------
        priority : int
            The priority level of the task acquiring the lock.

        Returns:
        --------
        self : PriorityLock
            The lock instance.
        """
        await self.acquire(priority)
        return self

    async def __aexit__(self, exc_type, exc, tb, update_priority_threshold):
        """
        Async context manager exit. Releases the lock and optionally updates the priority threshold.

        Parameters:
        -----------
        exc_type : Exception or None
            The exception type, if any, raised in the 'async with' block.
        exc : Exception or None
            The exception instance, if any, raised in the 'async with' block.
        tb : traceback or None
            The traceback object, if any, associated with the exception.
        update_priority_threshold : bool
            If True, increments the priority threshold after releasing the lock.
        """
        await self.release(update_priority_threshold)  # Now release is async


class PriorityLockManager:
    """
    A helper class to manage the acquisition and release of a PriorityLock using an 'async with' block.

    Attributes:
    -----------
    _lock : PriorityLock
        The PriorityLock instance to be managed.
    _priority : int
        The priority level for the current task.
    _update_priority_threshold : bool
        Whether to update the priority threshold after the lock is released.
    """

    def __init__(self, lock, priority, update_priority_threshold):
        """
        Initializes a PriorityLockManager with a PriorityLock and task-specific parameters.

        Parameters:
        -----------
        lock : PriorityLock
            The lock instance to manage.
        priority : int
            The priority level for the current task.
        update_priority_threshold : bool
            Whether to update the priority threshold after releasing the lock.
        """
        self._lock = lock  # The lock being managed
        self._priority = priority  # The priority level for the current task
        self._update_priority_threshold = update_priority_threshold

    async def __aenter__(self):
        """
        Async context manager entry. Acquires the lock with the specified priority.

        Returns:
        --------
        lock : PriorityLock
            The lock instance that was acquired.
        """
        await self._lock.acquire(self._priority)  # Acquire the lock with priority
        return self._lock

    async def __aexit__(self, exc_type, exc, tb):
        """
        Async context manager exit. Releases the lock and optionally updates the priority threshold.

        Parameters:
        -----------
        exc_type : Exception or None
            The exception type, if any, raised in the 'async with' block.
        exc : Exception or None
            The exception instance, if any, raised in the 'async with' block.
        tb : traceback or None
            The traceback object, if any, associated with the exception.
        """
        await self._lock.__aexit__(exc_type, exc, tb, self._update_priority_threshold)  # Release the lock


async def _run_task_with_priority_lock(
    task: Callable, lock: asyncio.Lock, priority: int,
    update_priority_threshold: bool, *args, **kwargs):
    """
    Runs a task within the context of a PriorityLock, ensuring proper acquisition and release.

    Parameters:
    -----------
    task : Callable
        The async function representing the task to be executed.
    lock : PriorityLock
        The PriorityLock instance managing access.
    priority : int
        The priority level for the task.
    update_priority_threshold : bool
        Whether to update the priority threshold after releasing the lock.
    *args, **kwargs:
        Additional arguments to pass to the task function.

    Returns:
    --------
    result : Any
        The result of the task execution.
    """
    async with PriorityLockManager(lock, priority, update_priority_threshold):  # Acquire the lock based on priority
        return await task(*args, **kwargs)
