from __future__ import annotations

import asyncio
import enum
import inspect
import json
import logging
import os
import sys
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Callable, Coroutine, NamedTuple, overload

from pydantic import BaseModel, Field, field_serializer
from redis.asyncio.lock import Lock

from fluid import settings
from fluid.utils import kernel, log
from fluid.utils.data import compact_dict
from fluid.utils.dates import as_utc
from fluid.utils.text import create_uid, trim_docstring

from .crontab import Scheduler
from .errors import TaskDecoratorError, TaskRunError

if TYPE_CHECKING:
    from .consumer import TaskManager


TaskExecutor = Callable[["TaskRun"], Coroutine[Any, Any, Any]]
RandomizeType = Callable[[], float | int]


class TaskPriority(enum.StrEnum):
    high = enum.auto()
    medium = enum.auto()
    low = enum.auto()


class TaskState(enum.StrEnum):
    init = enum.auto()
    queued = enum.auto()
    running = enum.auto()
    success = enum.auto()
    failure = enum.auto()
    aborted = enum.auto()
    rate_limited = enum.auto()

    @property
    def is_failure(self) -> bool:
        return self is TaskState.failure

    @property
    def is_done(self) -> bool:
        return self in FINISHED_STATES


FINISHED_STATES = frozenset(
    (TaskState.success, TaskState.failure, TaskState.aborted, TaskState.rate_limited)
)


class TaskManagerConfig(BaseModel):
    schedule_tasks: bool = True
    consume_tasks: bool = True
    max_concurrent_tasks: int = settings.MAX_CONCURRENT_TASKS
    """number of coroutine workers"""
    sleep: float = 0.1
    """amount to sleep after completion of a task"""
    broker_url: str = ""


class TaskInfoBase(BaseModel):
    name: str = Field(description="Task name")
    description: str = Field(description="Task description")
    module: str = Field(description="Task module")
    priority: TaskPriority = Field(description="Task priority")
    schedule: str | None = Field(default=None, description="Task schedule")


class TaskInfoUpdate(BaseModel):
    enabled: bool = Field(default=True, description="Task enabled")
    last_run_end: datetime | None = Field(
        default=None, description="Task last run end as milliseconds since epoch"
    )
    last_run_duration: timedelta | None = Field(
        default=None, description="Task last run duration in milliseconds"
    )
    last_run_state: TaskState | None = Field(
        default=None, description="State of last task run"
    )


class TaskInfo(TaskInfoBase, TaskInfoUpdate):
    pass


class QueuedTask(BaseModel):
    """A task to be queued"""

    run_id: str = Field(description="Task run id")
    task: str = Field(description="Task name")
    params: dict[str, Any] = Field(description="Task parameters")
    priority: TaskPriority | None = Field(default=None, description="Task priority")


class Task(NamedTuple):
    """A Task executes any time it is invoked"""

    name: str
    """Task name - unique identifier"""
    executor: TaskExecutor
    """Task executor function"""
    logger: logging.Logger
    """Task logger"""
    module: str = ""
    """Task python module"""
    short_description: str = ""
    """Short task description - one line"""
    description: str = ""
    """Task description - obtained from the executor docstring if not provided"""
    schedule: Scheduler | None = None
    """Task schedule - None means the task is not scheduled"""
    randomize: RandomizeType | None = None
    params_model: type[BaseModel] | None = None
    max_concurrency: int = 0
    """how many tasks can run in each consumer concurrently - 0 means no limit"""
    timeout_seconds: int = 60
    """Task timeout in seconds - how long the task can run before being aborted"""
    priority: TaskPriority = TaskPriority.medium
    """Task priority - high, medium, low"""

    @property
    def cpu_bound(self) -> bool:
        """True if the task is CPU bound"""
        return self.executor is run_in_subprocess

    def params_dump_json(self, params: dict[str, Any]) -> str:
        """Dumps task parameters as JSON string"""
        if params_model := self.params_model:
            return params_model(**params).model_dump_json()
        return json.dumps(params)

    def info(self, **params: Any) -> TaskInfo:
        params.update(
            name=self.name,
            description=self.description,
            module=self.module,
            priority=self.priority,
            schedule=str(self.schedule) if self.schedule else None,
        )
        return TaskInfo(**compact_dict(params))


class TaskRun(BaseModel, arbitrary_types_allowed=True):
    """A TaskRun contains all the data generated by a Task run

    This model is never initialized directly, it is created by the TaskManager
    """

    id: str
    task: Task
    priority: TaskPriority
    params: dict[str, Any]
    state: TaskState = TaskState.init
    task_manager: TaskManager = Field(exclude=True, repr=False)
    queued: datetime | None = None
    start: datetime | None = None
    end: datetime | None = None

    async def execute(self) -> None:
        try:
            self.set_state(TaskState.running)
            await self.task.executor(self)
        except Exception:
            self.set_state(TaskState.failure)
            raise
        else:
            self.set_state(TaskState.success)

    @field_serializer("task")
    def serialize_task(self, task: Task, _info: Any) -> str:
        return task.name

    @property
    def logger(self) -> logging.Logger:
        return self.task.logger

    @property
    def in_queue(self) -> timedelta | None:
        if self.queued and self.start:
            return self.start - self.queued
        return None

    @property
    def duration(self) -> timedelta | None:
        if self.start and self.end:
            return self.end - self.start
        return None

    @property
    def duration_ms(self) -> float | None:
        duration = self.duration
        if duration is not None:
            return round(1000 * duration.total_seconds(), 2)
        return None

    @property
    def total(self) -> timedelta | None:
        if self.queued and self.end:
            return self.end - self.queued
        return None

    @property
    def name(self) -> str:
        return self.task.name

    @property
    def name_id(self) -> str:
        return f"{self.task.name}.{self.id}"

    @property
    def is_done(self) -> bool:
        return self.state.is_done

    @property
    def is_failure(self) -> bool:
        return self.state.is_failure

    def params_dump_json(self) -> str:
        return self.task.params_dump_json(self.params)

    def set_state(
        self,
        state: TaskState,
        state_time: datetime | None = None,
    ) -> None:
        if self.state == state:
            return
        state_time = as_utc(state_time)
        match (self.state, state):
            case (TaskState.init, TaskState.queued):
                self.queued = state_time
                self.state = state
                self._dispatch()
            case (TaskState.init, _):
                self.set_state(TaskState.queued, state_time)
                self.set_state(state, state_time)
            case (TaskState.queued, TaskState.running):
                self.start = state_time
                self.state = state
                self._dispatch()
            case (
                TaskState.queued,
                TaskState.success
                | TaskState.aborted
                | TaskState.rate_limited
                | TaskState.failure,
            ):
                self.set_state(TaskState.running, state_time)
                self.set_state(state, state_time)
            case (
                TaskState.running,
                TaskState.success
                | TaskState.aborted
                | TaskState.rate_limited
                | TaskState.failure,
            ):
                self.end = state_time
                self.state = state
                self._dispatch()
            case _:
                raise TaskRunError(f"invalid state transition {self.state} -> {state}")

    def lock(self, timeout: float | None) -> Lock:
        return self.task_manager.broker.lock(self.name, timeout=timeout)

    def _dispatch(self) -> None:
        self.task_manager.dispatcher.dispatch(self.model_copy())


@dataclass
class TaskRunWaiter:
    task_manager: TaskManager
    uid: str = field(default_factory=create_uid)
    _runs: dict[str, TaskRun] = field(default_factory=dict)

    def __enter__(self) -> TaskRunWaiter:
        for state in FINISHED_STATES:
            self.task_manager.dispatcher.register_handler(f"{state}.{self.uid}", self)
        return self

    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        for state in FINISHED_STATES:
            self.task_manager.dispatcher.unregister_handler(f"{state}.{self.uid}")

    def __call__(self, task_run: TaskRun) -> None:
        self._runs[task_run.id] = task_run

    async def wait(self, task_run: TaskRun, *, timeout: int = 2) -> TaskRun:
        async with asyncio.timeout(timeout):
            while task_run.id not in self._runs:
                await asyncio.sleep(0.01)
        return self._runs[task_run.id]


@overload
def task(executor: TaskExecutor) -> Task: ...


@overload
def task(
    *,
    name: str | None = None,
    schedule: Scheduler | None = None,
    short_description: str | None = None,
    description: str | None = None,
    randomize: RandomizeType | None = None,
    max_concurrency: int = 1,
    priority: TaskPriority = TaskPriority.medium,
    params_model: BaseModel | None = None,
    cpu_bound: bool = False,
    timeout_seconds: int = 60,
) -> TaskConstructor: ...


# implementation of the task decorator
def task(executor: TaskExecutor | None = None, **kwargs: Any) -> Task | TaskConstructor:
    """Decorator to create a Task

    This decorator can be used in two ways:

    - As a simple decorator of the executor function
    - As a function with keyword arguments
    """
    if kwargs and executor:
        raise TaskDecoratorError("cannot use positional parameters")
    elif kwargs:
        return TaskConstructor(**kwargs)
    elif not executor:
        raise TaskDecoratorError("this is a decorator cannot be invoked in this way")
    else:
        return TaskConstructor()(executor)


class TaskConstructor:
    def __init__(self, *, cpu_bound: bool = False, **kwargs: Any) -> None:
        self.cpu_bound = cpu_bound
        self.kwargs = kwargs

    def __call__(self, executor: TaskExecutor) -> Task:
        if self.cpu_bound:
            return self.cpu_bound_task(executor)
        else:
            return self.create_task(executor)

    def create_task(
        self,
        executor: TaskExecutor,
        defaults: dict[str, Any] | None = None,
    ) -> Task:
        kwargs: dict[str, Any] = self.kwargs_defaults(executor)
        if defaults:
            kwargs.update(defaults)
        kwargs.update(self.kwargs)
        name = kwargs["name"]
        kwargs.update(
            executor=executor,
            logger=log.get_logger(f"task.{name}", prefix=True),
        )
        return Task(**kwargs)

    def cpu_bound_task(self, executor: TaskExecutor) -> Task:
        if is_in_subprocess():
            return self.create_task(executor)
        else:
            return self.create_task(run_in_subprocess, self.kwargs_defaults(executor))

    def kwargs_defaults(self, executor: TaskExecutor) -> dict[str, Any]:
        module = inspect.getmodule(executor)
        description = trim_docstring(inspect.getdoc(executor) or "")
        short_description = description.split("\n")[0].strip()
        return dict(
            name=get_name(executor),
            module=module.__name__ if module else "",
            description=description,
            short_description=short_description,
            executor=executor,
        )


def get_name(o: Any) -> str:
    if hasattr(o, "__name__"):
        return str(o.__name__)
    elif hasattr(o, "__class__"):
        return str(o.__class__.__name__)
    else:
        return str(o)


def is_in_subprocess() -> bool:
    return os.getenv("TASK_MANAGER_SPAWN") == "true"


class RemoteLog:
    def __init__(self, out: Any) -> None:
        self.out = out

    def __call__(self, data: bytes) -> None:
        self.out.write(data.decode("utf-8"))


async def run_in_subprocess(ctx: TaskRun) -> None:
    env = dict(os.environ)
    env["TASK_MANAGER_SPAWN"] = "true"
    result = await kernel.run_python(
        "-W",
        "ignore",
        "-m",
        "fluid.scheduler.cpubound",
        ctx.name,
        ctx.task.module,
        ctx.id,
        ctx.task.params_dump_json(ctx.params),
        result_callback=RemoteLog(sys.stdout),
        error_callback=RemoteLog(sys.stderr),
        env=env,
        stream_output=True,
        stream_error=True,
    )
    if result:
        raise TaskRunError(result)
