# 导入 DAG 类，用于定义工作流
from airflow import DAG

# 导入 PythonOperator：执行普通的 Python 函数
# 导入 BranchPythonOperator：根据 Python 函数的返回值选择执行分支
from airflow.operators.python import PythonOperator, BranchPythonOperator

# 导入 DummyOperator：空操作任务，用于组织任务依赖关系
from airflow.operators.dummy import DummyOperator

# 导入 TriggerRule：用于设置任务的触发规则
from airflow.utils.trigger_rule import TriggerRule

# 导入 datetime，用于设置 DAG 的开始日期
from datetime import datetime

# 导入 random，用于生成随机数
import random


# ==================== 1. 定义分支选择函数 ====================

def choose_branch(**kwargs):
    """
    分支选择函数。

    BranchPythonOperator 执行该函数，
    根据函数的返回值决定执行哪个下游任务。

    **kwargs：接收 Airflow 自动传入的任务上下文参数。
    """

    # 随机生成 1～3 之间的整数，包含 1 和 3
    random_number = random.randint(1, 3)

    # 将生成的随机数输出到任务日志中，方便调试
    print(f'Random number generated: {random_number}')

    # 根据随机数选择执行的分支
    # 注意：返回值必须是下游任务的 task_id

    if random_number == 1:
        return 'task_a'  # 随机数为 1，选择 A 分支

    elif random_number == 2:
        return 'task_b'  # 随机数为 2，选择 B 分支

    else:
        return 'task_c'  # 随机数为 3，选择 C 分支


# ==================== 2. 定义普通任务函数 ====================

def print_task(task_name):
    """
    普通任务函数。

    task_name：任务名称，用于区分不同的分支任务。
    """

    # 输出当前执行的任务名称
    print(f'Executing {task_name}')


# ==================== 3. 创建 DAG 工作流 ====================

with DAG(
    # DAG 的唯一标识符
    dag_id='branch_example',

    # 设置调度周期为每天执行一次
    schedule_interval='@daily',

    # 设置 DAG 开始日期
    start_date=datetime(2026, 10, 1),

    # 开启历史补跑
    # Airflow 会根据调度规则处理开始日期之后符合条件的历史调度周期
    catchup=True,

    # 设置 DAG 标签，方便在 Airflow 页面中查找
    tags=['demo']

) as dag:

    # ==================== 4. 定义分支选择任务 ====================

    branching_task = BranchPythonOperator(
        # 任务 ID
        task_id='branching_task',

        # 指定执行的 Python 函数
        # 函数的返回值决定执行哪个下游分支
        python_callable=choose_branch,
    )

    # ==================== 5. 定义 A 分支任务 ====================

    task_a = PythonOperator(
        # 任务 ID，必须与分支函数返回的名称一致
        task_id='task_a',

        # 使用 lambda 创建一个无参数函数
        # 执行任务时，调用 print_task('Task A')
        python_callable=lambda: print_task('Task A'),
    )

    # ==================== 6. 定义 B 分支任务 ====================

    task_b = PythonOperator(
        # B 分支的任务 ID
        task_id='task_b',

        # 执行时打印 Task B
        python_callable=lambda: print_task('Task B'),
    )

    # ==================== 7. 定义 C 分支任务 ====================

    task_c = PythonOperator(
        # C 分支的任务 ID
        task_id='task_c',

        # 执行时打印 Task C
        python_callable=lambda: print_task('Task C'),
    )

    # ==================== 8. 定义结束任务 ====================

    end_task = DummyOperator(
        # 结束任务的 ID
        task_id='end_task',

        # 设置触发规则：
        # 所有上游任务中没有失败的任务，并且至少有一个上游任务成功，
        # 当前任务就可以执行。
        #
        # 因为分支选择会把未选中的任务标记为 skipped，
        # 所以不能简单使用默认的 all_success，
        # 否则 end_task 可能因为存在 skipped 上游任务而被跳过。
        trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
    )

    # ==================== 9. 设置任务依赖关系 ====================

    # 分支任务执行完成后，根据返回值选择一个分支：
    #
    # 返回 'task_a' -> 执行 task_a
    # 返回 'task_b' -> 执行 task_b
    # 返回 'task_c' -> 执行 task_c
    #
    # 未被选中的分支会被标记为 skipped。
    #
    # 被选中的分支执行完成后，end_task 在满足触发规则时执行。

    branching_task >> [task_a, task_b, task_c] >> end_task

    # 上面的依赖关系等价于：
    #
    # branching_task >> task_a
    # branching_task >> task_b
    # branching_task >> task_c
    #
    # task_a >> end_task
    # task_b >> end_task
    # task_c >> end_task