# 导入 sleep，用于模拟模型训练所需的时间
from time import sleep

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

# 导入 PythonOperator，用于执行 Python 函数
from airflow.operators.python import PythonOperator

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

# 导入 random，用于生成随机准确率
import random


# ==================== 1. 模拟模型训练任务 ====================
def train_model(model_name, **context):
    """
    模拟训练一个机器学习模型，并将模型准确率存入 XCom。

    model_name：模型名称，例如 model_A、model_B、model_C
    **context：Airflow 传入的运行上下文，其中包含 ti 等信息
    """

    # 模拟模型训练过程，让当前任务等待 10 秒
    sleep(10)

    # 随机生成一个 0.7 到 1.0 之间的浮点数，模拟模型准确率
    accuracy = random.uniform(0.7, 1.0)

    # 将模型准确率手动推送到 XCom
    # key='model_accuracy'：这条数据在 XCom 中的键名
    # value=accuracy：实际保存的准确率
    context['ti'].xcom_push(
        key='model_accuracy',
        value=accuracy
    )

    # 将模型名称和准确率输出到任务日志中
    print(f'{model_name} accuracy: {accuracy}')

    # 注意：这里没有 return accuracy
    # 因此不会把 accuracy 作为函数返回值保存到默认的 return_value 中


# ==================== 2. 选择最佳模型 ====================
def choose_best_model(**context):
    """
    从三个训练任务的 XCom 中读取准确率，并找出准确率最高的模型。
    """

    # 获取当前任务实例 TaskInstance
    ti = context['ti']

    # 一次性读取三个上游任务的准确率
    accuracies = ti.xcom_pull(
        task_ids=['model_A', 'model_B', 'model_C'],  # 指定要读取的任务
        key='model_accuracy'                         # 指定要读取的 XCom 键
    )

    # 输出读取到的准确率列表
    # 例如：[0.82, 0.95, 0.88]
    print(f'accuracies: {accuracies}')

    # 找出列表中的最大准确率
    # max(accuracies) 返回最大值，例如 0.95
    # accuracies.index(...) 返回该值第一次出现的索引，例如 1
    best_model_index = accuracies.index(max(accuracies))

    # 输出最高准确率
    # 注意：这里原代码通过索引拼出了 model_1、model_2 或 model_3，
    # 并不是实际任务 ID model_A、model_B 或 model_C。
    print(
        f'Best model accuracy: {accuracies[best_model_index]} '
        f'from model_{best_model_index + 1}'
    )


# ==================== 3. 定义 DAG 工作流 ====================
with DAG(
    # DAG 的唯一标识符
    dag_id='xcom_example2',

    # DAG 的开始日期；不代表当天一定会自动执行
    start_date=datetime(2024, 10, 1),

    # 不设置定时调度，通常需要手动触发 DAG
    schedule_interval=None,

    # 不补跑历史上未执行的调度任务
    catchup=False,

    # 为 DAG 添加标签，便于在 Airflow 页面中筛选
    tags=["demo"]
) as dag:

    # ==================== 4. 定义模型 A 训练任务 ====================
    model_A = PythonOperator(
        # Airflow 中的任务 ID
        task_id='model_A',

        # 指定要执行的 Python 函数
        python_callable=train_model,

        # 向 train_model 函数传递额外参数
        # 相当于调用 train_model(model_name='model_A', ...)
        op_kwargs={'model_name': 'model_A'}
    )

    # ==================== 5. 定义模型 B 训练任务 ====================
    model_B = PythonOperator(
        task_id='model_B',
        python_callable=train_model,
        op_kwargs={'model_name': 'model_B'}
    )

    # ==================== 6. 定义模型 C 训练任务 ====================
    model_C = PythonOperator(
        task_id='model_C',
        python_callable=train_model,
        op_kwargs={'model_name': 'model_C'}
    )

    # ==================== 7. 定义最佳模型选择任务 ====================
    choose_task = PythonOperator(
        # 任务 ID
        task_id='choose_best_model',

        # 指定执行选择最佳模型的函数
        python_callable=choose_best_model
    )

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

    # 三个模型任务互相没有依赖关系，可以并行执行
    # choose_task 必须等待三个模型任务全部成功完成后才会执行
    [model_A, model_B, model_C] >> choose_task


# 脚本被直接执行时，调用 DAG 的 CLI 入口
# 这是旧式的命令行调试写法，不是部署 DAG 必需的内容；
# 正常使用 Airflow 调度器时，通常可以删除这段代码。
if __name__ == '__main__':
    dag.cli()