"""@dag decorator and DAG class."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable

from crondag.sdk.schedule import Cron, Interval, Schedule
from crondag.sdk.task import Task, TaskRef, _active_dag_context


# Global registry of all defined DAGs
_dag_registry: dict[str, "DAG"] = {}


@dataclass
class DAG:
    """A Directed Acyclic Graph of tasks."""
    dag_id: str
    description: str = ""
    schedule: Schedule = None
    tags: list[str] = field(default_factory=list)
    max_active_runs: int = 1
    _callable: Callable | None = field(default=None, repr=False)
    _task_refs: list[TaskRef] = field(default_factory=list, repr=False)
    _tasks: dict[str, Task] = field(default_factory=dict, repr=False)
    fileloc: str = ""

    def _add_task_ref(self, ref: TaskRef):
        """Register a task reference during DAG definition."""
        self._task_refs.append(ref)
        self._tasks[ref.task_id] = ref.operator

    def get_task_ids(self) -> list[str]:
        """Return all task IDs in this DAG."""
        return list(self._tasks.keys())

    def get_task(self, task_id: str) -> Task:
        """Get a task by ID."""
        return self._tasks[task_id]

    def get_dependency_map(self) -> dict[str, list[str]]:
        """Return {task_id: [upstream_task_ids]} mapping."""
        deps: dict[str, list[str]] = {}
        for ref in self._task_refs:
            upstream_ids = list({u.task_id for u in ref.upstream})
            deps[ref.task_id] = upstream_ids
        return deps

    def get_downstream_map(self) -> dict[str, list[str]]:
        """Return {task_id: [downstream_task_ids]} mapping."""
        dep_map = self.get_dependency_map()
        downstream: dict[str, list[str]] = {tid: [] for tid in self._tasks}
        for task_id, upstreams in dep_map.items():
            for u in upstreams:
                if u in downstream:
                    downstream[u].append(task_id)
        return downstream

    def __call__(self):
        """Execute the DAG definition function to build the task graph."""
        if self._callable:
            old = _active_dag_context.dag
            _active_dag_context.dag = self
            try:
                self._callable()
            finally:
                _active_dag_context.dag = old
        return self


def dag(
    _func: Callable | None = None,
    *,
    dag_id: str | None = None,
    schedule: str | Schedule | None = None,
    description: str = "",
    tags: list[str] | None = None,
    max_active_runs: int = 1,
) -> DAG | Callable:
    """Decorator to define a DAG.

    Usage:
        @dag(schedule="0 22 * * *")
        def my_dag():
            result = my_task()
            other_task(result)

        my_dag()  # Register the DAG
    """
    def decorator(func: Callable) -> DAG:
        import inspect

        did = dag_id or func.__name__

        # Normalize schedule
        sched: Schedule = None
        if isinstance(schedule, str):
            sched = Cron(schedule)
        elif isinstance(schedule, (Cron, Interval)):
            sched = schedule

        d = DAG(
            dag_id=did,
            description=description or (func.__doc__ or "").strip(),
            schedule=sched,
            tags=tags or [],
            max_active_runs=max_active_runs,
            _callable=func,
            fileloc=inspect.getfile(func),
        )

        # Execute the function to build the task graph
        d()

        # Register globally
        _dag_registry[did] = d
        return d

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


def get_all_dags() -> dict[str, DAG]:
    """Return all registered DAGs."""
    return dict(_dag_registry)


def clear_dag_registry():
    """Clear all registered DAGs (useful for testing)."""
    _dag_registry.clear()
