"""Taskiq entrypoint."""

import asyncio
import contextlib
from typing import TYPE_CHECKING, Any, Self

from taskiq.api import run_receiver_task, run_scheduler_task

from haolib.components.events import EventEmitter
from haolib.components.plugins.helpers import apply_plugin, apply_preset
from haolib.components.plugins.registry import PluginRegistry
from haolib.entrypoints.abstract import AbstractEntrypoint, EntrypointInconsistencyError
from haolib.entrypoints.events.abstract import EntrypointShutdownEvent, EntrypointStartupEvent
from haolib.entrypoints.plugins.abstract import AbstractEntrypointPlugin, AbstractEntrypointPluginPreset

if TYPE_CHECKING:
    from collections.abc import Mapping, Sequence
    from types import TracebackType

    from taskiq import AsyncBroker, TaskiqScheduler


class TaskiqEntrypointWorker:
    """Taskiq entrypoint worker."""

    def __init__(
        self,
        broker: AsyncBroker,
        scheduler: TaskiqScheduler | None = None,
        should_run_worker: bool = True,
        worker_args: Sequence[Any] | None = None,
        worker_kwargs: Mapping[str, Any] | None = None,
        scheduler_args: Sequence[Any] | None = None,
        scheduler_kwargs: Mapping[str, Any] | None = None,
    ) -> None:
        """Initialize the Taskiq entrypoint worker."""
        self._broker = broker
        self._scheduler = scheduler
        self._should_run_worker = should_run_worker
        self._worker_args = worker_args
        self._worker_kwargs = worker_kwargs
        self._scheduler_args = scheduler_args
        self._scheduler_kwargs = scheduler_kwargs
        self.worker_task: asyncio.Task[Any] | None = None
        self.scheduler_task: asyncio.Task[Any] | None = None

    async def startup(self) -> None:
        """Startup."""
        if self._scheduler is None and not self._should_run_worker:
            raise EntrypointInconsistencyError("Both scheduler and worker are not set.")

        await self._broker.startup()

        if self._should_run_worker:
            self.worker_task = asyncio.create_task(
                run_receiver_task(self._broker, *(self._worker_args or ()), **(self._worker_kwargs or {}))
            )

        if self._scheduler is not None:
            self.scheduler_task = asyncio.create_task(
                run_scheduler_task(self._scheduler, *(self._scheduler_args or ()), **(self._scheduler_kwargs or {}))
            )

    async def cancel_worker_task(self) -> None:
        """Cancel worker task."""
        if self.worker_task is not None:
            with contextlib.suppress(asyncio.CancelledError):
                self.worker_task.cancel()

                await self.worker_task

            self.worker_task = None

    async def cancel_scheduler_task(self) -> None:
        """Cancel scheduler task."""
        if self.scheduler_task is not None:
            with contextlib.suppress(asyncio.CancelledError):
                self.scheduler_task.cancel()

                await self.scheduler_task

            self.scheduler_task = None

    async def shutdown(self) -> None:
        """Shutdown."""

        await self.cancel_worker_task()

        await self.cancel_scheduler_task()

        await self._broker.shutdown()

    async def __aenter__(self) -> Self:
        """Enter context manager."""
        await self.startup()

        return self

    async def __aexit__(
        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
    ) -> None:
        """Exit context manager."""

        await self.shutdown()


class TaskiqEntrypoint(AbstractEntrypoint):
    """Taskiq entrypoint implementation.

    Provides a builder-pattern interface for configuring and running Taskiq
    task queue workers and schedulers with features like dependency injection
    and exception handling.

    Example:
        ```python
        from taskiq import AsyncBroker
        from taskiq_redis import RedisStreamBroker

        broker = RedisStreamBroker(url="redis://localhost:6379")
        entrypoint = (
            TaskiqEntrypoint(broker=broker)
            .setup_dishka(container)
            .setup_worker()
            .setup_scheduler(scheduler)
        )

        await entrypoint.run()
        ```

    Attributes:
        _broker: The Taskiq async broker instance.
        _should_run_worker: Whether to run the worker task.
        _scheduler: Optional Taskiq scheduler instance.
        _worker_args: Positional arguments for worker task.
        _worker_kwargs: Keyword arguments for worker task.
        _scheduler_args: Positional arguments for scheduler task.
        _scheduler_kwargs: Keyword arguments for scheduler task.
        _exception_handlers: Optional exception handlers dictionary.
        _worker: Optional TaskiqEntrypointWorker instance (created during startup).

    """

    def __init__(self, broker: AsyncBroker) -> None:
        """Initialize the Taskiq entrypoint.

        Args:
            broker: The Taskiq async broker instance.

        Raises:
            EntrypointInconsistencyError: If broker is None.

        """
        self._broker = broker
        self._should_run_worker = False
        self._scheduler: TaskiqScheduler | None = None
        self._worker_args: Sequence[Any] | None = None
        self._worker_kwargs: Mapping[str, Any] | None = None
        self._scheduler_args: Sequence[Any] | None = None
        self._scheduler_kwargs: Mapping[str, Any] | None = None
        self._worker: TaskiqEntrypointWorker | None = None
        self._events = EventEmitter[Self]()
        self._plugin_registry = PluginRegistry[Self]()

    def use_plugin(self, plugin: AbstractEntrypointPlugin[Self]) -> Self:
        """Add and apply a plugin to the entrypoint.

        Plugins are applied immediately and stored for lifecycle hooks.

        Args:
            plugin: The plugin to add and apply.

        Returns:
            Self for method chaining.

        """
        return apply_plugin(self, plugin, self._plugin_registry)

    def use_preset(
        self,
        preset: AbstractEntrypointPluginPreset[
            Self,
            AbstractEntrypointPlugin[Self],
        ],
    ) -> Self:
        """Add and apply a plugin preset to the entrypoint.

        Args:
            preset: The plugin preset to apply.

        Returns:
            Self for method chaining.

        """
        return apply_preset(self, preset, self._plugin_registry)

    @property
    def version(self) -> str:
        """Get the version of the Taskiq entrypoint."""
        return "1.0.0"

    @property
    def events(self) -> EventEmitter[Self]:
        """Get the event emitter for the entrypoint."""
        return self._events

    async def startup(self) -> None:
        """Startup the Taskiq entrypoint.

        Creates and initializes the TaskiqEntrypointWorker. Calls startup hooks for all plugins.
        This method is called before run() and should be idempotent.

        Raises:
            EntrypointInconsistencyError: If configuration is invalid.

        """
        self._worker = TaskiqEntrypointWorker(
            self._broker,
            self._scheduler,
            self._should_run_worker,
            self._worker_args,
            self._worker_kwargs,
            self._scheduler_args,
            self._scheduler_kwargs,
        )

        await self._worker.startup()

        await self.events.emit(EntrypointStartupEvent(component=self))

    async def shutdown(self) -> None:
        """Shutdown the Taskiq entrypoint.

        Cleans up resources and stops the worker and scheduler. Calls shutdown hooks for all plugins.
        This method is called after run() completes or is cancelled. Should be idempotent
        and safe to call multiple times.

        """
        await self.events.emit(EntrypointShutdownEvent(component=self))

        if self._worker is not None:
            await self._worker.shutdown()
            self._worker = None

    def get_broker(self) -> AsyncBroker:
        """Get the Taskiq broker instance.

        Returns:
            The Taskiq broker instance.

        """
        return self._broker

    @property
    def plugin_registry(self) -> PluginRegistry:
        """Get the plugin registry for plugin discovery.

        Provides read-only access to the plugin registry, allowing plugins
        to discover other plugins without accessing private attributes.

        Returns:
            The plugin registry instance.

        """
        return self._plugin_registry

    def setup_worker(self, *args: Any, **kwargs: Any) -> Self:
        """Setup worker task.

        Configures the entrypoint to run a worker that processes tasks from
        the broker queue.

        Args:
            *args: The arguments to pass to the run_receiver_task function.
            **kwargs: The keyword arguments to pass to the run_receiver_task function.

        Returns:
            Self for method chaining.

        Example:
            ```python
            entrypoint.setup_worker(max_async_tasks=10)
            ```

        """
        self._should_run_worker = True

        self._worker_args = args
        self._worker_kwargs = kwargs

        return self

    def setup_scheduler(self, taskiq_scheduler: TaskiqScheduler, *args: Any, **kwargs: Any) -> Self:
        """Setup scheduler task.

        Configures the entrypoint to run a scheduler that executes scheduled tasks.

        Args:
            taskiq_scheduler: The Taskiq scheduler instance.
            *args: The arguments to pass to the run_scheduler_task function.
            **kwargs: The keyword arguments to pass to the run_scheduler_task function.

        Returns:
            Self for method chaining.

        Example:
            ```python
            from taskiq import TaskiqScheduler
            from taskiq_redis import ListRedisScheduleSource

            scheduler = TaskiqScheduler(
                broker=broker,
                sources=[ListRedisScheduleSource(url="redis://localhost:6379")]
            )
            entrypoint.setup_scheduler(scheduler)
            ```

        """
        self._scheduler = taskiq_scheduler
        self._scheduler_args = args
        self._scheduler_kwargs = kwargs

        return self

    async def run(self) -> None:
        """Run the Taskiq entrypoint.

        Waits for the worker and/or scheduler tasks to complete. This method
        runs indefinitely until cancelled or an error occurs. The tasks are
        already started during startup().

        Raises:
            EntrypointInconsistencyError: If entrypoint was not started via startup().
            Exception: Any exception that occurs during execution (unless handled).

        """
        if self._worker is None:
            raise EntrypointInconsistencyError("Taskiq entrypoint must be started via startup() before run()")

        # Collect tasks to wait on
        tasks_to_wait: list[asyncio.Task[Any]] = []
        if self._worker.worker_task is not None:
            tasks_to_wait.append(self._worker.worker_task)
        if self._worker.scheduler_task is not None:
            tasks_to_wait.append(self._worker.scheduler_task)

        if not tasks_to_wait:
            raise EntrypointInconsistencyError("No tasks to run in Taskiq entrypoint")

        # Wait for tasks to complete
        with contextlib.suppress(asyncio.CancelledError):
            await asyncio.gather(*tasks_to_wait)
