"""
A class to manage concurrent [work](/guide/workers).
"""

from __future__ import annotations

import asyncio
import enum
import inspect
from contextvars import ContextVar
from threading import Event
from time import monotonic
from typing import (
    TYPE_CHECKING,
    Awaitable,
    Callable,
    Coroutine,
    Generic,
    TypeVar,
    Union,
    cast,
)

import rich.repr
from typing_extensions import TypeAlias

from .message import Message

if TYPE_CHECKING:
    from .app import App
    from .dom import DOMNode


active_worker: ContextVar[Worker] = ContextVar("active_worker")
"""Currently active worker context var."""


class NoActiveWorker(Exception):
    """There is no active worker."""


class WorkerError(Exception):
    """A worker related error."""


class WorkerFailed(WorkerError):
    """The worker raised an exception and did not complete."""

    def __init__(self, error: BaseException) -> None:
        self.error = error
        super().__init__(f"Worker raised exception: {error!r}")


class DeadlockError(WorkerError):
    """The operation would result in a deadlock."""


class WorkerCancelled(WorkerError):
    """The worker was cancelled and did not complete."""


def get_current_worker() -> Worker:
    """Get the currently active worker.

    Raises:
        NoActiveWorker: If there is no active worker.

    Returns:
        A Worker instance.
    """
    try:
        return active_worker.get()
    except LookupError:
        raise NoActiveWorker(
            "There is no active worker in this task or thread."
        ) from None


class WorkerState(enum.Enum):
    """A description of the worker's current state."""

    PENDING = 1
    """Worker is initialized, but not running."""
    RUNNING = 2
    """Worker is running."""
    CANCELLED = 3
    """Worker is not running, and was cancelled."""
    ERROR = 4
    """Worker is not running, and exited with an error."""
    SUCCESS = 5
    """Worker is not running, and completed successfully."""


ResultType = TypeVar("ResultType")


WorkType: TypeAlias = Union[
    Callable[[], Coroutine[None, None, ResultType]],
    Callable[[], ResultType],
    Awaitable[ResultType],
]
"""Type used for [workers](/guide/workers/)."""


class _ReprText:
    """Shim to insert a word into the Worker's repr."""

    def __init__(self, text: str) -> None:
        self.text = text

    def __repr__(self) -> str:
        return self.text


@rich.repr.auto(angular=True)
class Worker(Generic[ResultType]):
    """A class to manage concurrent work (either a task or a thread)."""

    @rich.repr.auto
    class StateChanged(Message, bubble=False, namespace="worker"):
        """The worker state changed."""

        def __init__(self, worker: Worker, state: WorkerState) -> None:
            """Initialize the StateChanged message.

            Args:
                worker: The worker object.
                state: New state.
            """
            self.worker = worker
            self.state = state
            super().__init__()

        def __rich_repr__(self) -> rich.repr.Result:
            yield self.worker
            yield self.state

    def __init__(
        self,
        node: DOMNode,
        work: WorkType,
        *,
        name: str = "",
        group: str = "default",
        description: str = "",
        exit_on_error: bool = True,
        thread: bool = False,
    ) -> None:
        """Initialize a Worker.

        Args:
            node: The widget, screen, or App that initiated the work.
            work: A callable, coroutine, or other awaitable object to run in the worker.
            name: Name of the worker (short string to help identify when debugging).
            group: The worker group.
            description: Description of the worker (longer string with more details).
            exit_on_error: Exit the app if the worker raises an error. Set to `False` to suppress exceptions.
            thread: Mark the worker as a thread worker.
        """
        self._node = node
        self._work = work
        self.name = name
        self.group = group
        self.description = description
        self.exit_on_error = exit_on_error
        self.cancelled_event: Event = Event()
        """A threading event set when the worker is cancelled."""
        self._thread_worker = thread
        self._state = WorkerState.PENDING
        self.state = self._state
        self._error: BaseException | None = None
        self._completed_steps: int = 0
        self._total_steps: int | None = None
        self._cancelled: bool = False
        self._created_time = monotonic()
        self._result: ResultType | None = None
        self._task: asyncio.Task | None = None
        self._node.post_message(self.StateChanged(self, self._state))

    def __rich_repr__(self) -> rich.repr.Result:
        yield _ReprText(self.state.name)
        yield "name", self.name, ""
        yield "group", self.group, "default"
        yield "description", self.description, ""
        yield "progress", round(self.progress, 1), 0.0

    @property
    def node(self) -> DOMNode:
        """The node where this worker was run from."""
        return self._node

    @property
    def state(self) -> WorkerState:
        """The current state of the worker."""
        return self._state

    @state.setter
    def state(self, state: WorkerState) -> None:
        """Set the state, and send a message."""
        changed = state != self._state
        self._state = state
        if changed:
            self._node.post_message(self.StateChanged(self, state))

    @property
    def is_cancelled(self) -> bool:
        """Has the work been cancelled?

        Note that cancelled work may still be running.
        """
        return self._cancelled

    @property
    def is_running(self) -> bool:
        """Is the task running?"""
        return self.state == WorkerState.RUNNING

    @property
    def is_finished(self) -> bool:
        """Has the task finished (cancelled, error, or success)?"""
        return self.state in (
            WorkerState.CANCELLED,
            WorkerState.ERROR,
            WorkerState.SUCCESS,
        )

    @property
    def completed_steps(self) -> int:
        """The number of completed steps."""
        return self._completed_steps

    @property
    def total_steps(self) -> int | None:
        """The number of total steps, or None if indeterminate."""
        return self._total_steps

    @property
    def progress(self) -> float:
        """Progress as a percentage.

        If the total steps is None, then this will return 0. The percentage will be clamped between 0 and 100.
        """
        if not self._total_steps:
            return 0.0
        return max(0, min(100, (self._completed_steps / self._total_steps) * 100.0))

    @property
    def result(self) -> ResultType | None:
        """The result of the worker, or `None` if there is no result."""
        return self._result

    @property
    def error(self) -> BaseException | None:
        """The exception raised by the worker, or `None` if there was no error."""
        return self._error

    def update(
        self, completed_steps: int | None = None, total_steps: int | None = -1
    ) -> None:
        """Update the number of completed steps.

        Args:
            completed_steps: The number of completed seps, or `None` to not change.
            total_steps: The total number of steps, `None` for indeterminate, or -1 to leave unchanged.
        """
        if completed_steps is not None:
            self._completed_steps += completed_steps
        if total_steps != -1:
            self._total_steps = None if total_steps is None else max(0, total_steps)

    def advance(self, steps: int = 1) -> None:
        """Advance the number of completed steps.

        Args:
            steps: Number of steps to advance.
        """
        self._completed_steps += steps

    async def _run_threaded(self) -> ResultType:
        """Run a threaded worker.

        Returns:
            Return value of the work.
        """

        def run_awaitable(work: Awaitable[ResultType]) -> ResultType:
            """Set the active worker and await the awaitable."""

            async def do_work() -> ResultType:
                active_worker.set(self)
                return await work

            return asyncio.run(do_work())

        def run_coroutine(
            work: Callable[[], Coroutine[None, None, ResultType]],
        ) -> ResultType:
            """Set the active worker and await coroutine."""
            return run_awaitable(work())

        def run_callable(work: Callable[[], ResultType]) -> ResultType:
            """Set the active worker, and call the callable."""
            active_worker.set(self)
            return work()

        if (
            inspect.iscoroutinefunction(self._work)
            or hasattr(self._work, "func")
            and inspect.iscoroutinefunction(self._work.func)
        ):
            runner = run_coroutine
        elif inspect.isawaitable(self._work):
            runner = run_awaitable
        elif callable(self._work):
            runner = run_callable
        else:
            raise WorkerError("Unsupported attempt to run a thread worker")

        loop = asyncio.get_running_loop()
        assert loop is not None
        return await loop.run_in_executor(None, runner, self._work)

    async def _run_async(self) -> ResultType:
        """Run an async worker.

        Returns:
            Return value of the work.
        """
        if (
            inspect.iscoroutinefunction(self._work)
            or hasattr(self._work, "func")
            and inspect.iscoroutinefunction(self._work.func)
        ):
            return await self._work()
        elif inspect.isawaitable(self._work):
            return await self._work
        elif callable(self._work):
            raise WorkerError("Request to run a non-async function as an async worker")
        raise WorkerError("Unsupported attempt to run an async worker")

    async def run(self) -> ResultType:
        """Run the work.

        Implement this method in a subclass, or pass a callable to the constructor.

        Returns:
            Return value of the work.
        """
        return await (
            self._run_threaded() if self._thread_worker else self._run_async()
        )

    async def _run(self, app: App) -> None:
        """Run the worker.

        Args:
            app: App instance.
        """
        app._set_active()
        active_worker.set(self)

        self.state = WorkerState.RUNNING
        app.log.worker(self)
        try:
            self._result = await self.run()
        except asyncio.CancelledError as error:
            self.state = WorkerState.CANCELLED
            self._error = error
            app.log.worker(self)
        except Exception as error:
            self.state = WorkerState.ERROR
            self._error = error
            app.log.worker(self, "failed", repr(error))
            from rich.traceback import Traceback

            app.log.worker(Traceback())
            if self.exit_on_error:
                worker_failed = WorkerFailed(self._error)
                app._handle_exception(worker_failed)
        else:
            self.state = WorkerState.SUCCESS
            app.log.worker(self)

    def _start(
        self, app: App, done_callback: Callable[[Worker], None] | None = None
    ) -> None:
        """Start the worker.

        Args:
            app: An app instance.
            done_callback: A callback to call when the task is done.
        """
        if self._task is not None:
            return
        self.state = WorkerState.RUNNING
        self._task = asyncio.create_task(self._run(app))

        def task_done_callback(_task: asyncio.Task) -> None:
            """Run the callback.

            Called by `Task.add_done_callback`.

            Args:
                The worker's task.
            """
            if done_callback is not None:
                done_callback(self)

        self._task.add_done_callback(task_done_callback)

    def cancel(self) -> None:
        """Cancel the task."""
        self._cancelled = True
        if self._task is not None:
            self._task.cancel()
        self.cancelled_event.set()

    async def wait(self) -> ResultType:
        """Wait for the work to complete.

        Raises:
            WorkerFailed: If the Worker raised an exception.
            WorkerCancelled: If the Worker was cancelled before it completed.

        Returns:
            The return value of the work.
        """
        try:
            if active_worker.get() is self:
                raise DeadlockError(
                    "Can't call worker.wait from within the worker function!"
                )
        except LookupError:
            # Not in a worker
            pass

        if self.state == WorkerState.PENDING:
            raise WorkerError("Worker must be started before calling this method.")
        if self._task is not None:
            try:
                await self._task
            except asyncio.CancelledError as error:
                self.state = WorkerState.CANCELLED
                self._error = error
        if self.state == WorkerState.ERROR:
            assert self._error is not None
            raise WorkerFailed(self._error)
        elif self.state == WorkerState.CANCELLED:
            raise WorkerCancelled("Worker was cancelled, and did not complete.")
        return cast("ResultType", self._result)
