"""@task decorator and Task class."""
from __future__ import annotations

import functools
from dataclasses import dataclass, field
from datetime import timedelta
from typing import Any, Callable

from crondag.sdk.trigger_rule import TriggerRule


@dataclass
class TaskRef:
    """Reference to a task call within a DAG definition.

    Tracks which task was called and what arguments were passed,
    enabling dependency graph construction.
    """
    task_id: str
    operator: "Task"
    args: tuple = ()
    kwargs: dict = field(default_factory=dict)
    upstream: list["TaskRef"] = field(default_factory=list)

    def __rshift__(self, other: "TaskRef") -> "TaskRef":
        """task_a >> task_b sets task_a as upstream of task_b."""
        if isinstance(other, TaskRef):
            other.upstream.append(self)
        elif isinstance(other, list):
            for item in other:
                if isinstance(item, TaskRef):
                    item.upstream.append(self)
        return other

    def __rrshift__(self, other):
        """Support list >> task_ref."""
        if isinstance(other, list):
            for item in other:
                if isinstance(item, TaskRef):
                    self.upstream.append(item)
        return self


@dataclass
class Task:
    """A task definition (the decorated function + metadata)."""
    task_id: str
    python_callable: Callable
    retries: int = 0
    retry_delay: timedelta = field(default_factory=lambda: timedelta(minutes=5))
    trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS
    timeout: int | None = None  # seconds
    _dag_context: Any = field(default=None, repr=False)

    def __call__(self, *args, **kwargs) -> TaskRef:
        """When called inside a @dag function, record dependency instead of executing."""
        ref = TaskRef(
            task_id=self.task_id,
            operator=self,
            args=args,
            kwargs=kwargs,
        )
        # Track upstream dependencies from arguments that are TaskRefs
        for arg in args:
            if isinstance(arg, TaskRef):
                ref.upstream.append(arg)
        for v in kwargs.values():
            if isinstance(v, TaskRef):
                ref.upstream.append(v)
        # Register with active DAG context
        if _active_dag_context.dag is not None:
            _active_dag_context.dag._add_task_ref(ref)
        return ref

    def execute(self, *args, **kwargs):
        """Actually execute the underlying function."""
        return self.python_callable(*args, **kwargs)


class _DagContext:
    """Thread-local-like context for tracking which DAG is being defined."""
    def __init__(self):
        self.dag = None


_active_dag_context = _DagContext()


def task(
    _func: Callable | None = None,
    *,
    task_id: str | None = None,
    retries: int = 0,
    retry_delay: timedelta = timedelta(minutes=5),
    trigger_rule: TriggerRule = TriggerRule.ALL_SUCCESS,
    timeout: int | None = None,
) -> Task | Callable:
    """Decorator to define a task.

    Usage:
        @task
        def my_task(): ...

        @task(retries=3)
        def my_task(): ...
    """
    def decorator(func: Callable) -> Task:
        tid = task_id or func.__name__
        return Task(
            task_id=tid,
            python_callable=func,
            retries=retries,
            retry_delay=retry_delay,
            trigger_rule=trigger_rule,
            timeout=timeout,
        )

    if _func is not None:
        return decorator(_func)
    return decorator
