import asyncio

import math

import os

import threading

import time

from concurrent.futures import ProcessPoolExecutor

def count_primes(limit: int) -> dict:

    """

    CPU 密集型任务：统计 2..limit 之间的质数数量

    这个函数会运行在子进程中

    """

    start = time.perf_counter()

    count = 0

    for num in range(2, limit + 1):

        is_prime = True

        sqrt_n = int(math.sqrt(num))

        for i in range(2, sqrt_n + 1):

            if num % i == 0:

                is_prime = False

                break

        if is_prime:

            count += 1

    duration = time.perf_counter() - start

    return {

        "pid": os.getpid(),

        "thread_id": threading.get_native_id(),

        "limit": limit,

        "prime_count": count,

        "duration": round(duration, 3),

    }

async def run_cpu_task(pool: ProcessPoolExecutor, limit: int) -> dict:

    """

    在 asyncio 中把 CPU 密集型任务提交到进程池

    """

    loop = asyncio.get_running_loop()

    print(f"[main] 提交任务 limit={limit}")

    result = await loop.run_in_executor(pool, count_primes, limit)

    print(

        f"[main] 完成任务 limit={limit}, "

        f"pid={result['pid']}, "

        f"thread_id={result['thread_id']}, "

        f"prime_count={result['prime_count']}, "

        f"耗时={result['duration']}s"

    )

    return result

async def heartbeat():

    """

    模拟事件循环中同时还在处理别的异步任务，并输出线程监视信息

    """

    for i in range(10000000):

        threads = threading.enumerate()

        current_thread = threading.current_thread()

        process_id = os.getpid()

        thread_names = ", ".join(

            f"{thread.name}:{thread.native_id}" for thread in threads

        )

        print(

            f"[heartbeat] event loop 仍然活着... {i} | "

            f"主进程PID={process_id} | "

            f"当前线程={current_thread.name} | "

            f"当前线程ID={current_thread.native_id} | "

            f"活跃线程数={len(threads)} | "

            f"线程列表=[{thread_names}]"

        )

        await asyncio.sleep(0.5)

async def main():

    limits = [10000000,

              210000,

              310000,

              410000,

              510000,

              610000,

              710000,

              810000,

              910000,

              910000,

              10000,

              100,

              2210000,

              210000,

              2100003,

              230100,

              1100001,

              1100002,

              1100003,

              1100004,

              1100005,

              1100006,

              10000,

              12000,

              130000,

              140000,

              150000

              ]

    start = time.perf_counter()

    # Windows 下一定要放在 if __name__ == "__main__" 保护下运行

    with ProcessPoolExecutor(60) as pool:

        cpu_tasks = [run_cpu_task(pool, limit) for limit in limits]

        # 同时跑：CPU 任务 + 普通异步任务

        results, _ = await asyncio.gather(

            asyncio.gather(*cpu_tasks),

            heartbeat(),

        )

    total = time.perf_counter() - start

    print("\n=== 汇总结果 ===")

    for item in results:

        print(item)

    print(f"\n总耗时: {total:.3f}s")

if __name__ == "__main__":

    asyncio.run(main())