"""示例 8：``@task`` 装饰器与 ``TaskContext``。

演示四件事：

* ``@task`` 把元数据写在函数旁边，``register_module`` 一次登记整个模块；
* ``ctx.progress(...)`` 上报进度——它会进快照、进存储，并出现在 ``ps`` 的 PROGRESS 列；
* ``ctx.log(...)`` 走标准库 ``logging``，logger 名字是 ``chongming_task.task.<task_id>``；
* ``ctx.input(节点id)`` 取上游节点的返回值——管道之间真正能传值。

运行::

    uv run python examples/08_task_context.py
"""

from __future__ import annotations

import logging
import sys
import time
from typing import Any

from chongming_task import Controller, Pipeline, RetryPolicy, render_table, task


@task(tags={"io"}, retry=RetryPolicy(max_attempts=2))
def extract(ctx: Any, day: str) -> list[str]:
    """按天抽取分区列表。"""
    ctx.log(f"开始抽取 {day}")
    partitions = [f"{day}/part{i}" for i in range(6)]
    for index, part in enumerate(partitions, start=1):
        ctx.checkpoint()
        ctx.progress(index, total=len(partitions), unit="partitions", message=part)
        time.sleep(0.05)
    return partitions


@task
def transform(ctx: Any, day: str) -> int:
    """把上游抽出的分区逐个转换，返回处理条数。"""
    partitions = ctx.input("extract", [])
    ctx.log(f"待转换分区：{len(partitions)} 个")
    rows = 0
    for index, part in enumerate(partitions, start=1):
        ctx.checkpoint()
        rows += 100
        ctx.progress(index, total=len(partitions), unit="partitions", message=part)
        time.sleep(0.05)
    return rows


@task
def load(ctx: Any, day: str) -> str:
    """入库：条数来自上游的返回值，而不是自己再算一遍。"""
    rows = ctx.input("transform", 0)
    ctx.log(f"{day} 共 {rows} 行待入库")
    for index in range(1, 5):
        ctx.checkpoint()
        ctx.progress(index, total=4, unit="batch")
        time.sleep(0.04)
    return f"{day}: {rows} 行已入库"


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(levelname)-7s %(name)s %(message)s",
        stream=sys.stdout,
    )

    controller = Controller(max_parallel=2)

    # ---- 一次登记整个模块 ----
    registered = controller.register_module(sys.modules[__name__])
    print("从 @task 元数据登记的：")
    for item in registered:
        print(f"  {item.name:<10} tags={sorted(item.tags) or '-'}  {item.description}")

    # ---- 串成管道：数据沿依赖流动 ----
    pipeline = Pipeline("daily-etl")
    extract_node = pipeline.node("extract", "2026-10-02", node_id="extract")
    transform_node = pipeline.node(
        "transform", "2026-10-02", node_id="transform", depends_on=[extract_node]
    )
    pipeline.node("load", "2026-10-02", node_id="load", depends_on=[transform_node])

    run = controller.submit_pipeline(pipeline)
    print(f"\n已提交管道 {run.run_id}\n")

    # ---- 跑的过程中看进度 ----
    deadline = time.monotonic() + 30
    while time.monotonic() < deadline and not run.settled:
        snapshots = controller.list_tasks(pipeline_run_id=run.run_id)
        cells: list[str] = []
        for snapshot in snapshots:
            node = snapshot.name.split("/", 1)[-1]
            progress = "-"
            if snapshot.progress is not None:
                progress = f"{snapshot.progress.done}/{snapshot.progress.total}"
            cells.append(f"{node:<10} {snapshot.status.value:<10} {progress:<8}")
        sys.stdout.write("\r" + " | ".join(cells).ljust(120))
        sys.stdout.flush()
        time.sleep(0.1)

    outcome = run.wait(timeout=30)
    print("\n\n--- 结束 ---")
    print(render_table(controller.list_tasks(pipeline_run_id=run.run_id)))
    print(f"\n管道结论：{outcome.status.value}")

    final = controller.get(run.task_id_of("load"))
    print(f"load 的返回值：{final.outcome if final else None}")  # type: ignore[union-attr]

    controller.close(timeout=10)
    print(
        "\n要点：ctx 是任务与运行时之间唯一的协作入口——停止、暂停、进度、日志、上游数据都在这里。"
    )


if __name__ == "__main__":
    main()