# Adapted from FramePack (https://github.com/lllyasviel/FramePack), Copyright (c) Lvmin Zhang (lllyasviel),
# Apache License 2.0 (see docs/third_party_licenses/Apache-2.0.txt). Modified for WanGP.
# Modifications Copyright (c) 2025-2026 DeepBeepMeep - part of WanGP, WanGP Community License 2.0 (see LICENSE.txt)
# based on FramePack https://github.com/lllyasviel/FramePack

import time
import traceback

from threading import Thread, Lock


class _TaskRunner:
    def __init__(self, runner_name="default"):
        self.runner_name = str(runner_name or "default")
        self.thread_name = self.runner_name.replace("_", " ").strip().title() or "Default"
        self.task_queue = []
        self.lock = Lock()
        self.thread = None

    def _process_tasks(self):
        while True:
            task = None
            with self.lock:
                if self.task_queue:
                    task = self.task_queue.pop(0)
                    
            if task is None:
                time.sleep(0.001)
                continue
                
            func, args, kwargs, thread_name = task
            current_name = None
            thread = self.thread
            try:
                if thread_name and thread is not None:
                    current_name = thread.name
                    thread.name = thread_name
                func(*args, **kwargs)
            except Exception as e:
                tb = traceback.format_exc().split('\n')[:-1] 
                print('\n'.join(tb))

                # print(f"Error in listener thread: {e}")
            finally:
                if current_name is not None and thread is not None:
                    thread.name = current_name

    def add_task(self, func, *args, thread_name=None, **kwargs):
        with self.lock:
            self.task_queue.append((func, args, kwargs, thread_name))
            thread = None if self.thread is not None else Thread(target=self._process_tasks, daemon=True, name=self.thread_name)
            if thread is not None:
                self.thread = thread

        if thread is not None:
            thread.start()
        return func

    def promote_task(self, task) -> bool:
        return self.promote_tasks([task])

    def promote_tasks(self, tasks) -> bool:
        targets = list(tasks)
        if not targets:
            return False
        with self.lock:
            remaining = list(self.task_queue)
            promoted = []
            for task in targets:
                index = next((index for index, queued_task in enumerate(remaining) if queued_task[0] is task), None)
                if index is None:
                    return False
                promoted.append(remaining.pop(index))
            self.task_queue[:] = promoted + remaining
            return True


class Listener:
    runners = {}
    lock = Lock()

    @classmethod
    def _get_runner(cls, runner_name="default"):
        runner_name = str(runner_name or "default")
        with cls.lock:
            runner = cls.runners.get(runner_name, None)
            if runner is None:
                runner = _TaskRunner(runner_name)
                cls.runners[runner_name] = runner
            return runner

    @classmethod
    def add_task(cls, func, *args, runner_name="default", thread_name=None, **kwargs):
        return cls._get_runner(runner_name).add_task(func, *args, thread_name=thread_name, **kwargs)

    @classmethod
    def promote_task(cls, task, runner_name="default"):
        return cls._get_runner(runner_name).promote_task(task)

    @classmethod
    def promote_tasks(cls, tasks, runner_name="default"):
        return cls._get_runner(runner_name).promote_tasks(tasks)


def async_run(func, *args, thread_name=None, **kwargs):
    Listener.add_task(func, *args, thread_name=thread_name, **kwargs)


def async_run_in(runner_name, func, *args, thread_name=None, **kwargs):
    return Listener.add_task(func, *args, runner_name=runner_name, thread_name=thread_name, **kwargs)


def promote_async_task(runner_name, task):
    return Listener.promote_task(task, runner_name=runner_name)


def promote_async_tasks(runner_name, tasks):
    return Listener.promote_tasks(tasks, runner_name=runner_name)


class FIFOQueue:
    def __init__(self):
        self.queue = []
        self.lock = Lock()

    def push(self, cmd, data = None):
        with self.lock:
            self.queue.append( (cmd, data) )

    def pop(self):
        with self.lock:
            if self.queue:
                return self.queue.pop(0)
            return None

    def top(self):
        with self.lock:
            if self.queue:
                return self.queue[0]
            return None

    def next(self):
        while True:
            with self.lock:
                if self.queue:
                    return self.queue.pop(0)

            time.sleep(0.001)


class AsyncStream:
    def __init__(self):
        self.input_queue = FIFOQueue()
        self.output_queue = FIFOQueue()