from __future__ import annotations

from collections.abc import Callable, Coroutine
from concurrent.futures import Future
from functools import wraps
from inspect import iscoroutinefunction
from logging import getLogger
from threading import Lock, current_thread
from typing import Any, Literal, Protocol, cast, overload

from PySide6.QtCore import QObject, QRunnable, QThreadPool, Signal, Slot
from PySide6.QtWidgets import QApplication
from vsengine.loops import EventLoop, get_loop

type _CoroutineFunc[**P, R] = Callable[P, Coroutine[Any, Any, R]]
type _Func[**P, R] = Callable[P, R]

_logger = getLogger(__name__)


class QtEventLoop(QObject, EventLoop):
    """Qt event loop adapter for vsview."""

    _invoke = Signal(int)

    def attach(self) -> None:
        self._lock = Lock()
        self._counter = 0
        self._pending = dict[int, Callable[[], None]]()
        self._invoke.connect(self._on_invoke)

    def detach(self) -> None:
        self.wait_for_threads(100)
        self._invoke.disconnect(self._on_invoke)
        self._pending.clear()

    @Slot(int)
    def _on_invoke(self, task_id: int) -> None:
        with self._lock:
            wrapper = self._pending.pop(task_id, None)
        if wrapper:
            wrapper()

    def from_thread[**P, R](self, func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> Future[R]:
        """Schedule func to run on the main Qt thread."""
        fut = Future[R]()

        def wrapper() -> None:
            if not fut.set_running_or_notify_cancel():
                return
            try:
                result = func(*args, **kwargs)
            except BaseException as e:
                _logger.debug(e, exc_info=True)
                fut.set_exception(e)
            else:
                fut.set_result(result)

        with self._lock:
            task_id = self._counter
            self._counter += 1
            self._pending[task_id] = wrapper

        self._invoke.emit(task_id)

        return fut

    def to_thread[**P, R](self, func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> Future[R]:
        """Run func in Qt's global thread pool."""
        fut = Future[R]()

        def wrapper() -> None:
            if not fut.set_running_or_notify_cancel():
                return
            try:
                result = func(*args, **kwargs)
            except BaseException as e:
                _logger.error(e)
                _logger.debug(e, exc_info=True)
                fut.set_exception(e)
            else:
                fut.set_result(result)

        QThreadPool.globalInstance().start(QRunnable.create(wrapper))
        return fut

    def to_thread_named[**P, R](self, name: str, func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> Future[R]:
        """Run func in Qt's global thread pool with a custom thread name."""
        fut = Future[R]()

        def wrapper() -> None:
            current_thread().name = name
            if not fut.set_running_or_notify_cancel():
                return
            try:
                result = func(*args, **kwargs)
            except BaseException as e:
                _logger.error(e)
                _logger.debug(e, exc_info=True)
                fut.set_exception(e)
            else:
                fut.set_result(result)

        QThreadPool.globalInstance().start(QRunnable.create(wrapper))
        return fut

    def wait_for_threads(self, timeout_ms: int = 500) -> None:
        """Wait for all background threads in the global thread pool to finish."""

        pool = QThreadPool.globalInstance()

        for _ in range(max(1, timeout_ms // 10)):
            if pool.activeThreadCount() == 0:
                break
            pool.waitForDone(10)
            QApplication.processEvents()

        # Final flush to process any signals from threads that just finished
        QApplication.processEvents()


class _DecoratorFuture(Protocol):
    """Protocol for decorators that wrap a function and return a Future."""

    @overload
    def __call__[**P, R](self, func: _CoroutineFunc[P, R]) -> Callable[P, Future[R]]: ...
    @overload
    def __call__[**P, R](self, func: _Func[P, R]) -> Callable[P, Future[R]]: ...


class _DecoratorDirect(Protocol):
    """Protocol for decorators that wrap a function and return the result directly."""

    @overload
    def __call__[**P, R](self, func: _CoroutineFunc[P, R]) -> Callable[P, R]: ...
    @overload
    def __call__[**P, R](self, func: _Func[P, R]) -> Callable[P, R]: ...


@overload
def run_in_loop[**P, R](func: _CoroutineFunc[P, R]) -> Callable[P, Future[R]]: ...
@overload
def run_in_loop[**P, R](func: _Func[P, R]) -> Callable[P, Future[R]]: ...
@overload
def run_in_loop(*, return_future: Literal[True]) -> _DecoratorFuture: ...
@overload
def run_in_loop(*, return_future: Literal[False]) -> _DecoratorDirect: ...
@overload
def run_in_loop(*, return_future: bool) -> _DecoratorFuture | _DecoratorDirect: ...
def run_in_loop(func: Any = None, *, return_future: bool = True) -> Any:
    """
    Decorator. Executes the decorated function within the `QtEventLoop` (Main Thread).

    Args:
        func: The function to wrap (when used as `@run_in_loop` without parens)
        return_future: If False, blocks and returns R directly.

    Returns:
        A future object or the result directly, depending on return_future.

    Usage:
    ```python
    @run_in_loop
    def my_func(): ...


    @run_in_loop(return_future=False)
    def my_blocking_func(): ...
    ```
    """

    def decorator(fn: Any) -> Any:
        @wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            loop = cast(QtEventLoop, get_loop())

            if iscoroutinefunction(fn):

                def run_coro() -> Any:
                    import asyncio

                    coro = fn(*args, **kwargs)
                    try:
                        return asyncio.run(coro)
                    except RuntimeError:
                        return asyncio.run_coroutine_threadsafe(coro, asyncio.get_running_loop()).result()

                fut = loop.from_thread(run_coro)
            else:
                # Delegate to from_thread to marshal execution to the main loop
                fut = loop.from_thread(fn, *args, **kwargs)

            return fut if return_future else fut.result()

        return wrapper

    return decorator if func is None else decorator(func)


@overload
def run_in_background[**P, R](func: _CoroutineFunc[P, R]) -> Callable[P, Future[R]]: ...
@overload
def run_in_background[**P, R](func: _Func[P, R]) -> Callable[P, Future[R]]: ...
@overload
def run_in_background(*, name: str) -> _DecoratorFuture: ...
def run_in_background(func: Any = None, *, name: str | None = None) -> Any:
    """
    Executes the decorated function in a background thread (via QThreadPool)
    using the `QtEventLoop`'s `to_thread` logic.

    Args:
        func: The function to wrap (when used as `@run_in_background` without parens)
        name: Optional thread name for logging (when used as `@run_in_background(name="...")`)

    Returns:
        A future object representing the result of the execution.

    Usage:
    ```python
    @run_in_background
    def my_func(): ...


    @run_in_background(name="MyWorker")
    def my_named_func(): ...
    ```
    """

    def decorator(fn: Any) -> Any:
        @wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            loop = cast(QtEventLoop, get_loop())

            if iscoroutinefunction(fn):

                def run_coro() -> Any:
                    import asyncio

                    coro = fn(*args, **kwargs)
                    try:
                        return asyncio.run(coro)
                    except RuntimeError:
                        return asyncio.run_coroutine_threadsafe(coro, asyncio.get_running_loop()).result()

                return loop.to_thread(run_coro) if name is None else loop.to_thread_named(name, run_coro)

            return (
                loop.to_thread(fn, *args, **kwargs) if name is None else loop.to_thread_named(name, fn, *args, **kwargs)
            )

        return wrapper

    return decorator if func is None else decorator(func)
