from __future__ import annotations

import asyncio
import heapq
import logging
from contextlib import AsyncExitStack
from functools import partial
from time import monotonic
from types import ModuleType
from typing import Any, Awaitable, Callable, Self

from inflection import underscore
from starlette.datastructures import State
from typing_extensions import Annotated, Doc

from fluid import settings
from fluid.scheduler.plugin import TaskManagerPlugin
from fluid.utils import log
from fluid.utils.dates import utcnow
from fluid.utils.dispatcher import AsyncDispatcher, Dispatcher, Event
from fluid.utils.worker import AsyncConsumer, Worker, WorkerFunction, Workers

from .broker import TaskBroker, TaskRegistry
from .errors import TaskAbortedError, TaskRunError, UnknownTaskError
from .models import (
    Task,
    TaskManagerConfig,
    TaskPriority,
    TaskRun,
    TaskRunWaiter,
    TaskState,
)

AsyncHandler = Callable[[TaskRun], Awaitable[None]]

logger = log.get_logger(__name__)


class TaskDispatcher(Dispatcher[TaskRun]):
    """The task dispatcher is responsible for dispatching task run messages"""

    def event_type(self, message: TaskRun) -> str:
        """The event type is determined by the state of the task run"""
        return message.state


class AsyncTaskDispatcher(AsyncDispatcher[TaskRun]):
    """The async task dispatcher is responsible for dispatching task
    run messages asynchronously"""

    def event_type(self, message: TaskRun) -> str:
        """The event type is determined by the state of the task run"""
        return message.state


class TaskManager:
    """The task manager is the main class for managing tasks"""

    def __init__(
        self,
        *,
        deps: Any = None,
        config: TaskManagerConfig | None = None,
        **kwargs: Any,
    ) -> None:
        self.deps: Annotated[
            Any,
            Doc("""
                Dependencies for the task manager.

                Production applications requires global dependencies to be
                available to all tasks. This can be achieved by setting
                the `deps` attribute of the task manager to an object
                with the required dependencies.

                Each task can cast the dependencies to the required type.
                """),
        ] = (
            deps if deps is not None else State()
        )
        self.state: Annotated[
            State,
            Doc("""
                State for the task manager.
                This can be used by plugins to store state in the task manager.
                """),
        ] = State()
        self.config: Annotated[
            TaskManagerConfig, Doc("""Task manager configuration""")
        ] = config or TaskManagerConfig(**kwargs)
        self.dispatcher: Annotated[
            TaskDispatcher,
            Doc("""
                A dispatcher of [TaskRun][fluid.scheduler.TaskRun] events.

                Application can register handlers to listen for events
                happening during the lifecycle of a task run.
                """),
        ] = TaskDispatcher()
        self.broker = TaskBroker.from_url(self.config.broker_url)
        self.manager_id: str = self.broker.new_uuid()
        self._async_contexts: list[Any] = []
        self._stack = AsyncExitStack()

    async def __aenter__(self) -> Self:
        for cm in self._async_contexts:
            await self._stack.enter_async_context(cm)
        return self

    async def __aexit__(self, *exc_info: Any) -> None:
        try:
            await self._stack.__aexit__(*exc_info)
        finally:
            await self.broker.close()

    async def on_startup(self) -> None:
        await self.__aenter__()

    async def on_shutdown(self) -> None:
        await self.__aexit__(None, None, None)

    def add_async_context_manager(self, cm: Any) -> None:
        """Add an async context manager to the task manager

        These context managers are entered when the task manager starts
        """
        self._async_contexts.append(cm)

    @property
    def registry(self) -> TaskRegistry:
        """The task registry"""
        return self.broker.registry

    @property
    def type(self) -> str:
        """The type of the task manager"""
        return underscore(self.__class__.__name__)

    def register_task(self, task: Annotated[Task, Doc("Task to register")]) -> None:
        """Register a task with the task manager"""
        self.broker.register_task(task)

    async def execute(
        self,
        task: Annotated[
            str | Task,
            Doc(
                "The task or task name,"
                " if a task name it must be registered with the task manager."
            ),
        ],
        *,
        run_id: Annotated[
            str,
            Doc("Unique ID for the task run. If not provided a new UUID is generated."),
        ] = "",
        priority: Annotated[
            TaskPriority | None, Doc("Override the default task priority if provided")
        ] = None,
        **params: Annotated[
            Any,
            Doc(
                "The optional parameters for the task run. "
                "They must match the task params model"
            ),
        ],
    ) -> TaskRun:
        """Execute a task and wait for it to finish

        This method is an async method that should be used in an asynchronous
        context when one need to wait for the task to finish execution.
        """
        task_run = self.create_task_run(
            task,
            run_id=run_id,
            priority=priority,
            **params,
        )
        try:
            await task_run._execute()
        except TaskAbortedError as exc:
            await self.broker.set_task_aborted(task_run.id, str(exc))
        return task_run

    def execute_sync(
        self,
        task: Annotated[
            str | Task,
            Doc(
                "The task or task name,"
                " if a task name it must be registered with the task manager."
            ),
        ],
        *,
        run_id: Annotated[
            str,
            Doc("Unique ID for the task run. If not provided a new UUID is generated."),
        ] = "",
        priority: Annotated[
            TaskPriority | None, Doc("Override the default task priority if provided")
        ] = None,
        **params: Annotated[
            Any,
            Doc(
                "The optional parameters for the task run. "
                "They must match the task params model"
            ),
        ],
    ) -> TaskRun:
        """Execute a task synchronously

        This method is a blocking method that should be used in a synchronous
        context.
        """
        return asyncio.run(
            self._execute_and_exit(
                task,
                run_id=run_id,
                priority=priority,
                **params,
            )
        )

    async def queue(
        self,
        task: Annotated[
            str | Task,
            Doc(
                "The task or task name,"
                " if a task name it must be registered with the task manager."
            ),
        ],
        *,
        run_id: Annotated[
            str,
            Doc("Unique ID for the task run. If not provided a new UUID is generated."),
        ] = "",
        priority: Annotated[
            TaskPriority | None, Doc("Override the default task priority if provided")
        ] = None,
        **params: Annotated[
            Any,
            Doc(
                "The optional parameters for the task run. "
                "They must match the task params model"
            ),
        ],
    ) -> TaskRun:
        """Queue a task for execution

        This methods fires two events:

        - `init`: when the task run is created
        - `queued`: after the task is queued

        It returns the [TaskRun][fluid.scheduler.TaskRun] object
        """
        task_run = self.create_task_run(
            task,
            run_id=run_id,
            priority=priority,
            **params,
        )
        return await self._queue_task_run(task_run)

    def create_task_run(
        self,
        task: Annotated[
            str | Task,
            Doc(
                "The task or task name,"
                " if a task name it must be registered with the task manager."
            ),
        ],
        *,
        run_id: Annotated[
            str,
            Doc("Unique ID for the task run. If not provided a new UUID is generated."),
        ] = "",
        priority: Annotated[
            TaskPriority | None, Doc("Override the default task priority if provided")
        ] = None,
        **params: Annotated[
            Any,
            Doc(
                "The optional parameters for the task run. "
                "They must match the task params model"
            ),
        ],
    ) -> TaskRun:
        """Create a [TaskRun][fluid.scheduler.TaskRun] in `init` state"""
        task = self.broker.task_from_registry(task)
        run_id = run_id or self.broker.new_uuid()
        return TaskRun(
            id=run_id,
            task=task,
            priority=priority or task.priority,
            params=task.params_model(**params),
            task_manager=self,
        )

    def register_from_module(
        self,
        module: Annotated[
            ModuleType,
            Doc(
                "Python module with tasks implementations "
                "- can contain any object, only instances of Task are registered"
            ),
        ],
    ) -> None:
        """Register tasks from a python module"""
        for name in dir(module):
            if name.startswith("_"):
                continue
            if isinstance(obj := getattr(module, name), Task):
                self.register_task(obj)

    def register_from_dict(
        self,
        data: Annotated[
            dict[str, Any],
            Doc(
                "Python dictionary with tasks implementations "
                "- can contain any object, only instances of Task are registered"
            ),
        ],
    ) -> None:
        """Register tasks from a python dictionary"""
        for name, obj in data.items():
            if name.startswith("_"):
                continue
            if isinstance(obj, Task):
                self.register_task(obj)

    def register_async_handler(
        self,
        event: Annotated[Event | str, Doc("The event to register the handler for")],
        handler: AsyncHandler,
    ) -> None:
        """Register an async handler for a given event

        This method is a no op for a TaskManager that is not a worker
        """

    def unregister_async_handler(
        self,
        event: Annotated[Event | str, Doc("The event to unregister the handler for")],
    ) -> AsyncHandler | None:
        """Unregister an async handler for a given event

        This method is a no op for a TaskManager that is not a worker
        """

    def with_plugin(
        self,
        plugin: Annotated[TaskManagerPlugin, Doc("The plugin to register")],
    ) -> Self:
        """Register a plugin with the task manager"""
        plugin.register(self)
        return self

    async def _execute_and_exit(self, task: Task | str, **params: Any) -> TaskRun:
        async with self:
            return await self.execute(task, **params)

    async def _queue_task_run(self, task_run: TaskRun) -> TaskRun:
        self.dispatcher.dispatch(task_run)
        task_run.set_state(TaskState.queued)
        await self.broker.queue_task(task_run)
        return task_run


class TaskConsumer(TaskManager, Workers):
    """The Task Consumer is a Task Manager responsible for consuming tasks
    from a task queue
    """

    def __init__(
        self,
        *,
        name: Annotated[
            str,
            Doc("Worker's name, if not provided it is evaluated from the class name"),
        ] = "",
        stopping_grace_period: Annotated[
            float,
            Doc(
                "Grace period in seconds to wait for workers to stop running "
                "when this worker is shutdown. "
                "It defaults to the `FLUID_STOPPING_GRACE_PERIOD` "
                "environment variable or 10 seconds."
            ),
        ] = settings.STOPPING_GRACE_PERIOD,
        **config: Any,
    ) -> None:
        super().__init__(**config)
        Workers.__init__(
            self, name=name, stopping_grace_period=2 * stopping_grace_period
        )
        self._async_dispatcher_worker = AsyncConsumer(
            AsyncTaskDispatcher(), stopping_grace_period=stopping_grace_period
        )
        self._in_process_queue = InProcessTaskQueue(
            self, stopping_grace_period=stopping_grace_period
        )
        self.add_async_context_manager(self._async_dispatcher_worker)
        self.add_workers(self._in_process_queue)
        for i in range(self.config.max_concurrent_tasks):
            worker_name = f"task-worker-{i+1}"
            self.add_workers(
                WorkerFunction(
                    partial(self._consume_tasks, worker_name),
                    name=worker_name,
                    stopping_grace_period=stopping_grace_period,
                )
            )
        self.add_workers(
            WorkerFunction(
                self._ping_status,
                heartbeat=1.0,
                name="manager-status",
                stopping_grace_period=stopping_grace_period,
            )
        )

    def sync_queue(self, task: str | Task | TaskRun, delay: float = 0) -> None:
        """Queue a task synchronously"""
        self._in_process_queue.queue(task, delay=delay)

    async def queue_and_wait(
        self,
        task: Annotated[
            str | Task,
            Doc(
                "The task or task name,"
                " if a task name it must be registered with the task manager."
            ),
        ],
        *,
        timeout: Annotated[
            int | None, Doc("Timeout for waiting the task to finish")
        ] = None,
        **params: Annotated[
            Any,
            Doc(
                "The optional parameters for the task run. "
                "They must match the task params model"
            ),
        ],
    ) -> TaskRun:
        """Queue a task and wait for it to finish"""
        with TaskRunWaiter(self) as waiter:
            task_run = await self.queue(task, **params)
            return await waiter.wait(task_run, timeout=timeout)

    def register_async_handler(self, event: Event | str, handler: AsyncHandler) -> None:
        event = Event.from_string_or_event(event)
        self.dispatcher.register_handler(
            f"{event.type}.async_dispatch",
            self._async_dispatcher_worker.send,
        )
        self._async_dispatcher_worker.dispatcher.register_handler(event, handler)

    def unregister_async_handler(self, event: Event | str) -> AsyncHandler | None:
        return self._async_dispatcher_worker.dispatcher.unregister_handler(event)

    async def _ping_status(self) -> None:
        await self.broker.set_manager_status(
            self.manager_id,
            {"kind": self.type, "status": await self.status()},
            ttl=5,
        )

    # Internals

    async def _consume_tasks(self, worker_name: str) -> None:
        if not self.config.consume_tasks:
            await asyncio.sleep(self.config.sleep)
            return
        try:
            task_run = await self.broker.get_task_run(self)
            if task_run is None:
                await asyncio.sleep(self.config.sleep)
                return
        except UnknownTaskError as exc:
            logger.error(
                "%s unknown task %s - it looks like it is not "
                "registered with this consumer",
                worker_name,
                exc,
            )
            return
        # re-queue if the task is not yet ready to execute
        if self._re_queue_if_not_ready(task_run):
            return
        if await self._register_task_run(task_run):
            return
        # run the task
        if not task_run.is_done:
            try:
                params = task_run.params.model_dump_json()
            except Exception:
                task_run.logger.exception("%s - start - params exception", task_run.id)
            else:
                task_run.logger.info("%s - %s - start", task_run.id, params)
            try:
                await task_run._execute()
            except TaskRunError:
                # no logging as this was a controlled exception
                pass
            except TaskAbortedError as exc:
                task_run.logger.info("%s - %s - aborted - %s", task_run.id, params, exc)
            except asyncio.TimeoutError:
                task_run.logger.error("task run %s - %s - timeout", task_run.id, params)
            except asyncio.CancelledError:
                task_run.logger.warning(
                    "task run %s - interrupted by worker shutdown", task_run.id
                )
                task_run.set_state(TaskState.interrupted)
                await self.broker.remove_task_run(task_run)
                await self._update_task_run_status(task_run)
                raise
            except Exception as exc:
                task_run.logger.exception("critical exception while executing")
                if retry_task_run := task_run._maybe_failure_retry(exc):
                    await self._queue_task_run(retry_task_run)
            await self.broker.remove_task_run(task_run)
        await self._update_task_run_status(task_run)

    async def _update_task_run_status(self, task_run: TaskRun) -> None:
        duration_ms = task_run.duration_ms
        if duration_ms is not None:
            await self.broker.update_task(
                task_run.task,
                dict(
                    last_run_end=task_run.end,
                    last_run_duration=task_run.duration,
                    last_run_state=task_run.state,
                ),
            )
            task_run.logger.log(
                logging.WARNING if task_run.is_failure else logging.INFO,
                "end - %s - milliseconds - %s",
                task_run.state,
                duration_ms,
            )

    async def _register_task_run(self, task_run: TaskRun) -> bool:
        """Check enabled status and concurrency, register the task run for execution.

        Returns True if should return early, False otherwise.
        """
        task_info = await self.broker.get_tasks_info(task_run.name)
        if not task_info[0].enabled:
            task_run.set_state(TaskState.aborted)
            return False
        async with task_run.lock(timeout=5):
            if task_run.task.max_concurrency > 0:
                current_runs = await self.broker.current_task_runs(task_run.name)
                if new_task_run := task_run._maybe_rate_limit_retry(current_runs):
                    await self.broker.queue_task(new_task_run)
                    return True
            if task_run.state is not TaskState.rate_limited:
                await self.broker.add_task_run(task_run)
        return False

    def _re_queue_if_not_ready(self, task_run: TaskRun) -> bool:
        if task_run.execute_after is None:
            return False
        current_time = utcnow()
        if current_time >= task_run.execute_after:
            return False
        # wait for this delay before re-queuing
        delay = max((task_run.execute_after - current_time).total_seconds(), 5)
        self.sync_queue(task_run, delay=delay)
        return True


class InProcessTaskQueue(Worker):
    """In-process worker that holds delayed task runs in a min-heap ordered by
    due timestamp and dispatches them to the broker when they become due."""

    def __init__(
        self,
        task_consumer: TaskConsumer,
        *,
        name: Annotated[
            str,
            Doc("Worker's name, if not provided it is evaluated from the class name"),
        ] = "",
        stopping_grace_period: Annotated[
            float,
            Doc(
                "Grace period in seconds to wait for workers to stop running "
                "when this worker is shutdown. "
                "It defaults to the `FLUID_STOPPING_GRACE_PERIOD` "
                "environment variable or 10 seconds."
            ),
        ] = settings.STOPPING_GRACE_PERIOD,
    ) -> None:
        super().__init__(name=name, stopping_grace_period=stopping_grace_period)
        self.task_consumer = task_consumer
        self._sleep = task_consumer.config.sleep_millis / 1000
        self._heap: list[tuple[float, int, str | Task | TaskRun]] = []
        self._counter = 0

    def queue(self, task_run: str | Task | TaskRun, delay: float = 0) -> None:
        """Queue a task run, due in `delay` seconds from now"""
        heapq.heappush(self._heap, (monotonic() + delay, self._counter, task_run))
        self._counter += 1

    def queue_len(self) -> int:
        """Number of pending task runs in the queue"""
        return len(self._heap)

    async def run(self) -> None:
        while self.is_running():
            sleep = self._sleep
            if self._heap and self._heap[0][0] <= monotonic():
                await self._consume_one()
                sleep = 0
            await asyncio.sleep(sleep)
        # when stopping, re-queue all pending tasks to avoid losing them
        while self._heap:
            await self._consume_one()

    async def _consume_one(self) -> None:
        _, _, task = heapq.heappop(self._heap)
        if isinstance(task, TaskRun):
            await self.task_consumer.broker.queue_task(task)
        else:
            await self.task_consumer.queue(task)


# required by pydantic to avoid `Class not fully defined` error
TaskRun.model_rebuild()
