import time
import heapq
import logging
from typing import Optional, Any, Callable


class QueueHookManager:
    """Manages installation/uninstallation of queue hooks.

    Responsible for wrapping prompt queue methods to add persistence and pause behavior.
    """

    def __init__(self, *, is_paused_fn: Callable[[], bool], on_job_started: Callable[[str], None], on_task_done: Callable[[Any, Any, Any], None], should_run_when_paused: Optional[Callable[[str], bool]] = None):
        self._original_queue_get = None
        self._original_task_done = None
        self._installed = False
        self._is_paused = is_paused_fn
        self._on_job_started = on_job_started
        self._on_task_done = on_task_done
        self._should_run_when_paused = should_run_when_paused

    def install(self) -> None:
        """Install hooks into execution.PromptQueue if not already installed."""
        if self._installed:
            return
        import execution

        if self._original_queue_get is None:
            self._original_queue_get = execution.PromptQueue.get

            def get_wrapper(q_self, timeout=None):
                # If paused, only allow items explicitly permitted by the manager (run-selected mode)
                try:
                    if self._is_paused():
                        allowed = False
                        if callable(self._should_run_when_paused):
                            try:
                                top = q_self.queue[0] if getattr(q_self, 'queue', None) else None
                            except Exception:
                                top = None
                            if top is not None:
                                try:
                                    top_pid = str(top[1])
                                    allowed = bool(self._should_run_when_paused(top_pid))
                                except Exception:
                                    allowed = False
                        if not allowed:
                            time.sleep(0.1)
                            return None
                except Exception:
                    # If anything goes wrong during checks, be safe and respect pause
                    if self._is_paused():
                        time.sleep(0.1)
                        return None

                result = self._original_queue_get(q_self, timeout=timeout)
                if result is not None:
                    try:
                        item, _item_id = result
                        prompt_id = item[1]
                        # Helper: sanitize prompt by stripping non-node metadata keys
                        def _sanitize_item(it):
                            try:
                                n, pid, prompt, extra, outs = it
                                if isinstance(prompt, dict):
                                    try:
                                        prompt_clean = {k: v for k, v in prompt.items() if k not in ('workflow', 'name')}
                                    except Exception:
                                        prompt_clean = prompt
                                    if prompt_clean is not prompt:
                                        return (n, pid, prompt_clean, extra, outs)
                            except Exception:
                                pass
                            return it
                        # If paused, ensure the popped item is allowed; otherwise, reinsert and yield None
                        if self._is_paused() and callable(self._should_run_when_paused):
                            try:
                                if not self._should_run_when_paused(str(prompt_id)):
                                    # Undo the side effects of the pop from the original get():
                                    # 1) Remove from currently_running
                                    # 2) Reinsert the item into the heap queue
                                    # 3) Notify the server that the queue/running set changed
                                    try:
                                        with q_self.mutex:
                                            try:
                                                q_self.currently_running.pop(_item_id, None)
                                            except Exception:
                                                pass
                                            try:
                                                safe_item = _sanitize_item(item)
                                                q_self.queue.append(safe_item)
                                                heapq.heapify(q_self.queue)
                                            except Exception:
                                                pass
                                            try:
                                                q_self.server.queue_updated()
                                            except Exception:
                                                pass
                                    except Exception:
                                        pass
                                    time.sleep(0.05)
                                    return None
                            except Exception:
                                time.sleep(0.05)
                                return None
                        # Sanitize both the running copy and the returned item to prevent crashes
                        try:
                            with q_self.mutex:
                                if _item_id in q_self.currently_running:
                                    try:
                                        q_self.currently_running[_item_id] = _sanitize_item(q_self.currently_running[_item_id])
                                    except Exception:
                                        pass
                        except Exception:
                            pass
                        # Replace the returned item with sanitized version
                        try:
                            item = _sanitize_item(item)
                            result = (item, _item_id)
                        except Exception:
                            pass
                        self._on_job_started(prompt_id)
                    except Exception as e:
                        logging.debug(f"QueueHookManager get_wrapper failed: {e}")
                return result

            execution.PromptQueue.get = get_wrapper

        if self._original_task_done is None:
            self._original_task_done = execution.PromptQueue.task_done

            def task_done_wrapper(q_self, item_id, history_result, status):
                try:
                    self._on_task_done((q_self, item_id, history_result, status))
                except Exception as e:
                    logging.debug(f"QueueHookManager task_done on_task_done failed: {e}")
                self._original_task_done(q_self, item_id, history_result, status)

            execution.PromptQueue.task_done = task_done_wrapper

        self._installed = True

    def uninstall(self) -> None:
        """Restore original methods."""
        if not self._installed:
            return
        import execution
        if self._original_queue_get is not None:
            execution.PromptQueue.get = self._original_queue_get
            self._original_queue_get = None
        if self._original_task_done is not None:
            execution.PromptQueue.task_done = self._original_task_done
            self._original_task_done = None
        self._installed = False