from __future__ import annotations

from collections import OrderedDict
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field
from typing import Any, Generic, TypeVar

from datapizza.core.executors.async_executor import AsyncExecutor

from .errors import WorkflowDefinitionError
from .models import INTERNAL_START_EVENT, RouteSpec, TaskSpec, WorkflowDescription
from .policies import RetryPolicy, RunPolicies
from .state import DefaultStateAdapter, StateAdapter

StateT = TypeVar("StateT")
TaskFunc = TypeVar("TaskFunc", bound=Callable[..., Any])

TASK_METADATA_ATTR = "__workflow_task_definition__"


def _normalize_on(value: str | Sequence[str] | None) -> tuple[str, ...]:
    if value is None:
        return ()
    if isinstance(value, str):
        events = (value,)
    elif isinstance(value, Sequence):
        events = tuple(value)
    else:
        raise WorkflowDefinitionError("on must be a string or sequence of strings")

    if not events:
        raise WorkflowDefinitionError("on must not be empty")
    for item in events:
        if not isinstance(item, str) or not item:
            raise WorkflowDefinitionError("event subscriptions must be non-empty strings")
        if item == INTERNAL_START_EVENT:
            raise WorkflowDefinitionError(
                f"{INTERNAL_START_EVENT} is reserved for internal workflow use"
            )
    return events


@dataclass(frozen=True)
class TaskDefinition:
    name: str | None = None
    on: tuple[str, ...] = field(default_factory=tuple)
    start: bool = False
    retry: RetryPolicy = field(default_factory=RetryPolicy)
    timeout: float | None = None
    failure_event: str | None = None
    input_type: type[Any] | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_spec(self, method_name: str) -> TaskSpec:
        subscriptions = self.on
        if self.start:
            subscriptions = (INTERNAL_START_EVENT, *subscriptions)

        return TaskSpec(
            name=self.name or method_name,
            method_name=method_name,
            subscriptions=subscriptions,
            retry=self.retry,
            timeout=self.timeout,
            failure_event=self.failure_event,
            input_type=self.input_type,
            metadata=dict(self.metadata),
        )


def task(
    func: TaskFunc | str | None = None,
    *,
    name: str | None = None,
    on: str | Sequence[str] | None = None,
    start: bool = False,
    retry: RetryPolicy | dict | None = None,
    timeout: float | None = None,
    failure_event: str | None = None,
    input_type: type[Any] | None = None,
    metadata: dict[str, Any] | None = None,
):
    if isinstance(func, str):
        name = func
        func = None

    definition = TaskDefinition(
        name=name,
        on=_normalize_on(on),
        start=start,
        retry=RetryPolicy.from_value(retry),
        timeout=timeout,
        failure_event=failure_event,
        input_type=input_type,
        metadata=metadata or {},
    )

    if definition.failure_event == INTERNAL_START_EVENT:
        raise WorkflowDefinitionError(
            f"{INTERNAL_START_EVENT} is reserved for internal workflow use"
        )

    def decorator(f: TaskFunc) -> TaskFunc:
        setattr(f, TASK_METADATA_ATTR, definition)
        return f

    if func is None:
        return decorator
    return decorator(func)


def _build_routes(task_specs: tuple[TaskSpec, ...]) -> tuple[RouteSpec, ...]:
    routes: list[RouteSpec] = []
    for spec in task_specs:
        for event_name in spec.subscriptions:
            routes.append(RouteSpec(event=event_name, task=spec.name))
    return tuple(routes)


def _compile_workflow_class(
    *,
    bases: tuple[type[Any], ...],
    namespace: Mapping[str, Any],
) -> tuple[tuple[TaskSpec, ...], tuple[RouteSpec, ...]]:
    inherited: "OrderedDict[str, TaskSpec]" = OrderedDict()
    for base in bases:
        for spec in getattr(base, "__workflow_tasks__", ()):
            inherited[spec.method_name] = spec

    for attr_name, attr_value in namespace.items():
        metadata = getattr(attr_value, TASK_METADATA_ATTR, None)
        if metadata is None:
            inherited.pop(attr_name, None)
            continue

        inherited[attr_name] = metadata.to_spec(attr_name)

    task_specs = tuple(inherited.values())
    task_names = [spec.name for spec in task_specs]
    if len(task_names) != len(set(task_names)):
        raise WorkflowDefinitionError("Task names must be unique within a workflow")

    return task_specs, _build_routes(task_specs)


class Workflow(Generic[StateT]):
    __workflow_tasks__: tuple[TaskSpec, ...] = ()
    __workflow_routes__: tuple[RouteSpec, ...] = ()
    state_factory: Callable[[], StateT] = dict  # type: ignore[assignment]
    run_policies: RunPolicies = RunPolicies()
    state_adapter: StateAdapter[StateT] | None = None

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        cls.__workflow_tasks__, cls.__workflow_routes__ = _compile_workflow_class(
            bases=cls.__bases__,
            namespace=cls.__dict__,
        )

    def __init__(
        self,
        *,
        run_policies: RunPolicies | None = None,
        state_adapter: StateAdapter[StateT] | None = None,
    ):
        self._workflow_name = self.__class__.__name__
        self._run_policies = run_policies or self.run_policies
        self._state_adapter = state_adapter or self._build_state_adapter()
        self._validate_definition()

    def _validate_definition(self) -> None:
        if not self.__workflow_tasks__:
            raise WorkflowDefinitionError("Workflow must declare at least one task")
        if not any(
            route.event == INTERNAL_START_EVENT for route in self.__workflow_routes__
        ):
            raise WorkflowDefinitionError("Workflow must declare at least one start task")

    def _build_state_adapter(self) -> StateAdapter[StateT]:
        factory = getattr(self, "state_factory", dict)
        inferred_type = factory if isinstance(factory, type) else None
        return DefaultStateAdapter(inferred_type)

    @property
    def workflow_name(self) -> str:
        return self._workflow_name

    @property
    def policies(self) -> RunPolicies:
        return self._run_policies

    def _make_initial_state(self) -> StateT:
        state = self.state_factory()  # type: ignore[misc]
        return self._state_adapter.validate(state)

    def routes(self) -> list[RouteSpec]:
        return list(self.__workflow_routes__)

    def describe(self) -> WorkflowDescription:
        return WorkflowDescription(
            name=self.workflow_name,
            tasks=list(self.__workflow_tasks__),
            routes=self.routes(),
            run_policies=self.policies,
        )

    async def a_run(
        self,
        payload: Any = None,
        *,
        state: StateT | None = None,
        run_id: str | None = None,
    ):
        from .engine import WorkflowEngine

        engine = WorkflowEngine(self)
        return await engine.run(payload=payload, state=state, run_id=run_id)

    def run(
        self,
        payload: Any = None,
        *,
        state: StateT | None = None,
        run_id: str | None = None,
    ):
        return AsyncExecutor.get_instance().run(
            self.a_run(payload=payload, state=state, run_id=run_id)
        )
