"""示例 5：任务分桶 —— 本地并发隔离，以及为分布式预留的分配单位。

运行::

    uv run python examples/05_buckets.py
"""

from __future__ import annotations

import time

from chongming_task.controller import Controller, render_table


def sync_tenant(ctx, tenant: str) -> str:
    """同步一个租户的数据。"""
    for _ in range(20):
        ctx.checkpoint()
        time.sleep(0.02)
    return f"{tenant} 同步完成"


def main() -> None:
    controller = Controller(max_parallel=10, max_parallel_per_bucket=1)
    controller.register("sync", sync_tenant)
    policy = controller.bucket_policy

    print(f"分桶策略：{policy.num_buckets} 个桶")
    print("并发限制：全局 10，但每个桶最多 1 个\n")

    for index in range(4):
        controller.submit(
            "sync", "租户A", task_id=f"a{index}", label="A-同步", bucket_key="tenant-A"
        )
    for index in range(2):
        controller.submit(
            "sync", "租户B", task_id=f"b{index}", label="B-同步", bucket_key="tenant-B"
        )

    print("业务键 -> 桶号（同一个键在任何进程、任何机器上都会落到同一个桶）：")
    for key in ("tenant-A", "tenant-B"):
        print(f"  {key:<10} -> bucket {policy.bucket_for(key)}")

    time.sleep(0.5)
    print("\n--- 查询：A 桶积压，但不影响 B 桶 ---")
    print(render_table(controller.list_tasks()))
    print(
        f"    占用额度 {controller.active_count}，排队 {controller.queued_count}，"
        f"A 桶在用 {controller.active_in_bucket(policy.bucket_for('tenant-A'))}"
    )

    controller.wait_all(timeout=30)

    print("\n--- 查询：全部结束 ---")
    print(render_table(controller.list_tasks()))

    print("\n各桶任务数（分布式调度要靠它决定把桶分给谁）：")
    for bucket, count in sorted(
        controller.bucket_counts().items(), key=lambda kv: (kv[0] is None, kv[0])
    ):
        print(f"  bucket {bucket}: {count}")

    bucket_a = policy.bucket_for("tenant-A")
    print(f"\n按桶查询（只取 A 桶）：{[s.task_id for s in controller.list_tasks(bucket=bucket_a)]}")

    print("\n分布式预留：桶是分配单位——一台机器拥有一组桶，只取自己那组桶的任务即可。")


if __name__ == "__main__":
    main()