# Copyright cocotb contributors
# Licensed under the Revised BSD License, see LICENSE for details.
# SPDX-License-Identifier: BSD-3-Clause
from __future__ import annotations

import pdb  # noqa: T100
import sys
from asyncio import CancelledError
from collections.abc import Awaitable, Coroutine
from typing import (
    Any,
    Callable,
    NoReturn,
    TypeVar,
    overload,
)

import cocotb
import cocotb._event_loop
from cocotb._base_triggers import Event, NullTrigger, Trigger, TriggerCallback
from cocotb._deprecation import deprecated
from cocotb._extended_awaitables import SimTimeoutError
from cocotb._gpi_triggers import Timer
from cocotb.simtime import TimeUnit
from cocotb.task import ResultType, Task
from cocotb_tools import _env

if sys.version_info < (3, 11):
    from exceptiongroup import BaseExceptionGroup

_pdb_on_exception = _env.get_bool("COCOTB_PDB_ON_EXCEPTION")


class TestManager:
    """State of the currently executing Test."""

    # TODO
    # Make the tasks list a TaskManager.
    # Make shutdown errors and outcome be an ExceptionGroup from that TaskManager.
    # Replace result() with passing the outcome to the done callback.
    # Make this and Task the same object which is a Coroutine.
    # Reimplement the logic in the body of an async function.
    # Make RunningTest a normal Task that the RegressionManager runs and registers a
    #  done callback with.

    def __init__(
        self,
        test_coro: Coroutine[Trigger, None, Any],
        *,
        name: str,
        timeout: tuple[float, TimeUnit] | None = None,
        test_complete_cb: Callable[[], None],
    ) -> None:
        self._timeout = timeout
        self._timeout_cb: TriggerCallback | None = None
        self._test_complete_cb: Callable[[], None] = test_complete_cb

        # We create the main task here so that the init checks can be done immediately.
        self._main_task = Task[None](test_coro, name=f"Test {name}")
        self._tasks: dict[Task[Any], None] = {}
        self._excs: list[BaseException] = []
        self._finishing: bool = False
        self._complete: bool = False
        self._done = Event()

    def _test_done_callback(self, task: Task[None]) -> None:
        self.remove_task(task)
        # If cancelled, end the Test without additional error. This case would only
        # occur if a child threw a CancelledError or if the Test was forced to shutdown.
        if task.cancelled():
            self._abort()
            return
        # Handle outcome appropriately and shut down the Test.
        e = task.exception()
        if e is None:
            self._abort()
        elif isinstance(e, EndTest):
            msg_str = f": {e.msg}" if e.msg else ""
            task._log.info("Test stopped early by this task%s", msg_str)
            self._abort()
        else:
            self._abort(e)

    def start(self) -> None:
        # set global current test manager
        global _current_test
        _current_test = self

        # start main task
        self._main_task._add_done_callback(self._test_done_callback)
        self._main_task._start_soon()
        self._tasks[self._main_task] = None

        # start timeout if specified
        if self._timeout is not None:
            self._timeout_cb = Timer(*self._timeout)._register(self._on_timeout)

        cocotb._event_loop._inst.run()

    def _on_timeout(self) -> None:
        self._timeout_cb = None
        self._abort(SimTimeoutError())

    def exception(self) -> BaseException | None:
        if self._excs:
            if len(self._excs) == 1:
                return self._excs[0]
            else:
                return BaseExceptionGroup("Multiple exceptions in test", self._excs)
        return None

    def result(self) -> None:
        exc = self.exception()
        if exc:
            raise exc

    def done(self) -> bool:
        """Return whether the test has completed."""
        return self._complete

    def cancel(self, msg: str | None = None) -> None:
        """End the test prematurely."""
        if msg is not None:
            exc = CancelledError(msg)
        else:
            exc = CancelledError()
        self._abort(exc)

    def _abort(self, exc: BaseException | None = None) -> None:
        """Shutdown the test"""

        if exc is not None:
            self._excs.append(exc)

        if self._finishing:
            return
        self._finishing = True

        # Break into pdb on test end before all Tasks are killed.
        if _pdb_on_exception and exc is not None:
            try:
                pdb.post_mortem(exc.__traceback__)
            except BaseException:  # noqa: BLE001
                pdb.set_trace()  # noqa: T100

        # Cancel the timeout if it is running.
        if self._timeout_cb is not None:
            self._timeout_cb.cancel()

        # Cancel Tasks.
        for task in list(self._tasks):
            task._cancel_now()

        # Register function to clean up global state once all Tasks have ended.
        if self._done.is_set():
            self._on_complete()
        else:
            self._done.wait()._register(self._on_complete)

    def _on_complete(self) -> None:
        # Clear the current global current test manager
        global _current_test
        _current_test = None

        self._complete = True

        # Tell RegressionManager the test is complete
        self._test_complete_cb()

    def add_task(self, task: Task[Any]) -> None:
        task._add_done_callback(self._task_done_callback)
        self._tasks[task] = None
        self._done.clear()

    def remove_task(self, task: Task[Any]) -> None:
        del self._tasks[task]
        if not self._tasks:
            self._done.set()

    def _task_done_callback(self, task: Task[Any]) -> None:
        self.remove_task(task)
        # if cancelled, do nothing
        if task.cancelled():
            return
        # if there's a Task awaiting this one, don't fail
        if task.complete._callbacks:
            return
        # if no failure, do nothing
        e = task.exception()
        if e is None:
            return
        elif isinstance(e, EndTest):
            msg_str = f": {e.msg}" if e.msg else ""
            task._log.info("Test stopped early by this task%s", msg_str)
            self._abort()
        else:
            self._abort(e)


TriggerT = TypeVar("TriggerT", bound=Trigger)


@overload
def start_soon(
    coro: Task[ResultType]
    | Coroutine[Trigger, None, ResultType]
    | Awaitable[ResultType],
    *,
    name: str | None = None,
) -> Task[ResultType]: ...


@overload
def start_soon(
    coro: TriggerT,
    *,
    name: str | None = None,
) -> Task[TriggerT]: ...


def start_soon(
    coro: Awaitable[Any],
    *,
    name: str | None = None,
) -> Task[ResultType]:
    """
    Schedule a :term:`coroutine` to be run concurrently in a :class:`~cocotb.task.Task`.

    Note that this is not an :keyword:`async` function,
    and the new task will not execute until the calling task yields control.

    Args:
        coro: A :class:`!Task` or :term:`!coroutine` to be run concurrently.
        name:
            The task's name.

            .. versionadded:: 2.0

    Raises:
        RuntimeError: If a :class:`!Task` was given and it has already completed.

    Returns:
        The :class:`~cocotb.task.Task` that is scheduled to be run.

    .. versionadded:: 1.6

    .. versionchanged:: 2.1
        This function now accepts any :class:`~collections.abc.Awaitable` object, not just :term:`coroutines <coroutine>`.

        .. warning::
            If you are passing a bespoke :class:`~collections.abc.Awaitable` object to this function,
            read the :ref:`design note <awaitable-design-note>` for important information about how to use it correctly.
    """
    task = create_task(coro, name=name)
    if task._unstarted():
        task._start_soon()
    elif task.done():
        raise RuntimeError("Cannot schedule a Task that has already completed.")
    return task


@overload
async def start(
    coro: Task[ResultType]
    | Coroutine[Trigger, None, ResultType]
    | Awaitable[ResultType],
    *,
    name: str | None = None,
) -> Task[ResultType]: ...


@overload
async def start(
    coro: TriggerT,
    *,
    name: str | None = None,
) -> Task[TriggerT]: ...


@deprecated("Use ``cocotb.start_soon`` instead.")
async def start(
    coro: Awaitable[Any],
    *,
    name: str | None = None,
) -> Task[Any]:
    """
    Schedule a :term:`coroutine` to be run concurrently, then yield control to allow pending tasks to execute.

    The calling task will resume execution before control is returned to the simulator.

    When the calling task resumes, the newly scheduled task may have completed,
    raised an Exception, or be pending on a :class:`~cocotb.triggers.Trigger`.

    Args:
        coro: A :class:`!Task` or :term:`!coroutine` to be run concurrently.
        name:
            The task's name.

            .. versionadded:: 2.0

    Returns:
        The :class:`~cocotb.task.Task` that has been scheduled and allowed to execute.

    .. versionadded:: 1.6

    .. deprecated:: 2.0
        Use :func:`cocotb.start_soon` instead.
        If you need the scheduled Task to start before continuing the current Task,
        use an :class:`.Event` to block the current Task until the scheduled Task starts,
        like so:

        .. code-block:: python

            async def coro(started: Event) -> None:
                started.set()
                # Do stuff...


            task_started = Event()
            task = cocotb.start_soon(coro(task_started))
            await task_started.wait()

    .. versionchanged:: 2.1
        This function now accepts any :class:`~collections.abc.Awaitable` object, not just :term:`coroutines <coroutine>`.

        .. warning::
            If you are passing a bespoke :class:`~collections.abc.Awaitable` object to this function,
            read the :ref:`design note <awaitable-design-note>` for important information about how to use it correctly.
    """
    task = start_soon(coro, name=name)
    await NullTrigger()
    return task


@overload
def create_task(
    coro: Task[ResultType]
    | Coroutine[Trigger, None, ResultType]
    | Awaitable[ResultType],
    *,
    name: str | None = None,
) -> Task[ResultType]: ...


@overload
def create_task(
    coro: TriggerT,
    *,
    name: str | None = None,
) -> Task[TriggerT]: ...


def create_task(
    coro: Awaitable[Any],
    *,
    name: str | None = None,
) -> Task[Any]:
    """
    Construct a :term:`!coroutine` into a :class:`~cocotb.task.Task` without scheduling the task.

    The task can later be scheduled with :func:`cocotb.start` or :func:`cocotb.start_soon`.

    Args:
        coro: A :class:`!Task` or a :term:`!coroutine` to be turned into a :class:`!Task`.
        name:
            The task's name.

            .. versionadded:: 2.0

    Returns:
        Either the provided :class:`~cocotb.task.Task` or a new Task wrapping the coroutine.

    .. versionadded:: 1.6

    .. versionchanged:: 2.1
        This function now accepts any :class:`~collections.abc.Awaitable` object, not just :term:`coroutines <coroutine>`.

        .. warning::
            If you are passing a bespoke :class:`~collections.abc.Awaitable` object to this function,
            read the :ref:`design note <awaitable-design-note>` for important information about how to use it correctly.
    """
    if _current_test is None:
        raise RuntimeError("No test is currently running; cannot schedule new Task.")

    if isinstance(coro, Task):
        # We do not add the done callback here, so that custom Tasks are not considered
        # "toplevel" tasks and end the test when they fail.
        if name is not None:
            coro.set_name(name)
        return coro

    task = Task(coro, name=name)
    _current_test.add_task(task)

    return task


class TestSuccess(BaseException):
    """Implementation of :func:`pass_test`.

    Users are *not* intended to catch or raise this exception type.
    """

    def __init__(self, msg: str | None) -> None:
        super().__init__(msg)
        self.msg = msg


@deprecated(
    "Use `cocotb.skip` to skip a test, or use `cocotb.end_test` to end a test while respecting expected failures."
)
def pass_test(msg: str | None = None) -> NoReturn:
    """Force a test to pass.

    The test will end after this function is called.
    The outcome of the test will be forced to pass,
    even if a :deco:`cocotb.xfail` decorator is used,
    or ``expect_error`` and ``expect_fail`` arguments to :deco:`cocotb.test` are passed.

    Args:
        msg: The message to display when the test passes.

    .. versionadded:: 2.0

    .. deprecated:: 2.1
        Forcing a test to pass is considered bad practice.
        Use :func:`cocotb.skip` if you want to end the test immediately and force the outcome to 'skipped'.
        Or use :func:`cocotb.end_test` if you want to end the test immediately while respecting expected failures.
    """
    raise TestSuccess(msg)


class EndTest(BaseException):
    """Implementation of :func:`cocotb.end_test`.

    Users are *not* intended to catch or raise this exception type.
    """

    def __init__(self, msg: str | None) -> None:
        super().__init__(msg)
        self.msg = msg


def end_test(msg: str | None = None) -> NoReturn:
    """End a test.

    The test will end after this function is called as if it passed.
    However, any :deco:`cocotb.xfail` decorators
    or ``expect_error`` and ``expect_fail`` arguments to :func:`cocotb.test` will be respected.
    Meaning, if they are used, the test will fail, as a failure was expected.

    Args:
        msg: The message to display when the test ends.
    """
    raise EndTest(msg)


_current_test: TestManager | None = None
"""The currently executing test's state."""
