# Copyright (c) 2026 DeepBeepMeep - part of WanGP, WanGP Community License 2.0 (see LICENSE.txt)
"""Forward chat pause to the media worker's existing cooperative checkpoint."""
from uuid import uuid4
import time
import threading

from shared.utils.process_locks import gen_lock
from shared.utils.generation_state import pause_state, set_generation_state


def postprocessing_pause_checkpoint(gen):
    """Compatibility entry point for cooperative media workers."""
    from shared.utils.media_control import current_control, pause_checkpoint
    control = current_control()
    if control is not None and control.gen is gen:
        control.checkpoint()
    else:
        pause_checkpoint(gen)


def generation_pause_controls(state, previous):
    import gradio as gr
    current = pending, paused = pause_state(state.get('gen', {}))
    if current == previous:
        return gr.skip(), gr.skip(), current
    return gr.update(visible=not paused, interactive=not pending), gr.update(visible=paused), current


class MediaToolPause:
    def __init__(self, session, gen, client_ids, send_cmd):
        self.session = session
        self.gen = gen
        self.client_ids = set(client_ids)
        self.send_cmd = send_cmd
        self.owner = 'deepy_pause_' + uuid4().hex
        self.request = 'request:' + self.owner
        self.process = 'process:' + self.owner
        self.manual_pause = False
        self.announced_paused = False
        self.closed = threading.Event()
        self.watcher = None

    def _owns_active_task(self):
        task = self.gen.get('api_active_queue_task')
        if task is None:
            queue = self.gen.get('queue') or []
            task = queue[0] if queue else {}
        params = task.get('params', task.get('settings', task))
        return params.get('client_id') in self.client_ids

    def __enter__(self):
        if self.session is not None:
            self.session.media_tool_active = True
            def watch():
                while not self.closed.wait(.1):
                    self.poll()
            self.watcher = threading.Thread(target=watch, name='Deepy media pause', daemon=True)
            self.watcher.start()
        return self

    def _release_locked(self):
        if self.manual_pause:
            self.gen['resume'] = True
        if self.gen.get('process_status') in (self.request, self.process):
            set_generation_state(self.gen, process_status='process:main' if self.gen.get('main_process_running') else None)
        self.gen.get('process_names', {}).pop(self.owner, None)

    def poll(self):
        from shared.deepy.engine import mark_assistant_paused, request_assistant_interrupt, request_assistant_pause
        from shared.deepy.chat import build_status_event

        session = self.session
        if session is None:
            return
        acknowledged = False
        resuming = False
        interrupted = False
        with session.turn_lock:
            with gen_lock:
                status = self.gen.get('process_status')
                queue_errors = self.gen.get('queue_errors', {})
                cancelled = any(client_id in queue_errors and queue_errors[client_id][1] for client_id in self.client_ids)
                if not session.interrupt_requested and (cancelled or (self.gen.get('abort') and self._owns_active_task())):
                    request_assistant_interrupt(session)
                    interrupted = True
                if status not in ('request:pause', 'process:pause'):
                    self.manual_pause = False
                if status in ('request:pause', 'process:pause') and self._owns_active_task() and not self.manual_pause:
                    self.manual_pause = True
                    request_assistant_pause(session)
                if not session.pause_requested or session.interrupt_requested or session.drop_state_requested or self.gen.get('abort'):
                    self._release_locked()
                    if self.announced_paused and not session.interrupt_requested and not session.drop_state_requested:
                        self.announced_paused = False
                        resuming = True
                elif status == self.process or (status == 'process:pause' and self.manual_pause):
                    acknowledged = not session.paused
                elif status == 'process:main' and self.gen.get('main_process_running'):
                    # Only suspend this tool's active task, never another chat's
                    # generation or an unrelated GPU extension.
                    if self._owns_active_task():
                        self.gen.setdefault('process_names', {})[self.owner] = 'Deepy Pause'
                        self.gen['pause_msg'] = 'Media Processing Paused - Resume in Deepy'
                        set_generation_state(self.gen, process_status=self.request)
            if interrupted:
                self.send_cmd('chat_output', build_status_event('Stopping generation...', kind='stop_pending', session=session))
            if resuming:
                self.send_cmd('chat_output', build_status_event('Resuming media processing...', kind='tool', session=session))
            if acknowledged and mark_assistant_paused(session):
                self.announced_paused = True
                self.send_cmd('chat_output', build_status_event('Deepy is paused.', kind='paused', session=session))

    def __exit__(self, *exc):
        self.closed.set()
        if self.watcher is not None:
            self.watcher.join()
        with gen_lock:
            self._release_locked()
        if self.session is not None:
            with self.session.turn_lock:
                self.session.media_tool_active = False
                # If a tool ends while paused (completion/error), hand the
                # pending request back to the assistant's normal pause gate.
                if self.session.pause_requested:
                    self.session.paused = False