"""示例 7：DAG 管道 —— 一个源 → 两个分支 → 汇合。

管道把多个任务按依赖关系串起来，由一个调度器推进：节点先停在 ``blocked`` 等上游，
上游满足后才放行；注定跑不了的下游会被 ``skipped``。

运行::

    uv run python examples/07_pipeline.py
"""

from __future__ import annotations

import time
from typing import Any

from chongming_task import Controller, Pipeline, render_pipeline
from chongming_task.policies import AggregationPolicy, JoinMode


def fetch(ctx: Any, day: str) -> str:
    """取数：模拟一次耗时操作。"""
    for _ in range(10):
        ctx.checkpoint()
        time.sleep(0.03)
    return f"{day} 的数据"


def transform(ctx: Any, day: str, shard: int) -> str:
    """分片转换。"""
    for _ in range(8):
        ctx.checkpoint()
        time.sleep(0.02)
    return f"{day}/shard{shard}"


def load(ctx: Any, day: str) -> str:
    """入库：依赖全部分片完成。"""
    del ctx
    return f"{day} 已入库"


def audit(ctx: Any, day: str) -> str:
    """审计：**不管上游成败都要跑**（ALL_SETTLED）。"""
    del ctx
    return f"{day} 审计完成"


def main() -> None:
    controller = Controller(max_parallel=4)

    # ---- 定义 ----
    pipeline = Pipeline("nightly-etl", aggregation=AggregationPolicy.all_settled())

    controller.register("fetch", fetch, tags={"io"})
    controller.register("transform", transform)
    controller.register("load", load, tags={"io"})
    controller.register("audit", audit)

    fetch_node = pipeline.node("fetch", "2026-10-02", node_id="fetch")
    shards = [
        pipeline.node(
            "transform",
            "2026-10-02",
            index,
            node_id=f"shard{index}",
            depends_on=[fetch_node],
        )
        for index in range(3)
    ]
    # 入库：等所有分片成功（默认 ALL_SUCCEEDED 屏障）
    load_node = pipeline.node("load", "2026-10-02", node_id="load", depends_on=shards)
    # 审计：等 load 落终态即可，不管成败
    pipeline.node(
        "audit",
        "2026-10-02",
        node_id="audit",
        depends_on=[load_node],
        join=JoinMode.ALL_SETTLED,
    )

    print(f"管道 {pipeline.name!r} 共 {len(pipeline)} 个节点，拓扑序：")
    print("  " + " -> ".join(node.node_id for node in pipeline.topological_order()))
    print("  注意：位置参数是传给任务的，节点名由 node_id= 指定\n")

    # ---- 跑起来 ----
    run = controller.submit_pipeline(pipeline)
    print(f"已提交，pipeline_run_id = {run.run_id}\n")

    time.sleep(0.05)
    print("--- 刚提交时：下游都停在 blocked ---")
    print(render_pipeline(run))

    print("\n--- 等待收场 ---")
    outcome = run.wait(timeout=60)

    print(render_pipeline(run))
    print(f"\n管道结论：{outcome.status.value}")
    for node_id, status in outcome.nodes.items():
        print(f"  {node_id:<8} {status.value}")

    print(f"\n每个任务都带 pipeline_run_id：{_run_id_of(controller, run, 'load')}")
    controller.close(timeout=10)
    print("\n要点：依赖由状态机表达（blocked -> queued），调度器只负责放行。")


def _run_id_of(controller: Controller, run: Any, node_id: str) -> str:
    snapshot = controller.get(run.task_id_of(node_id))
    return "?" if snapshot is None else str(snapshot.pipeline_run_id)


if __name__ == "__main__":
    main()