#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""PROF 综合分析脚本：合并 worker + 主进程 cProfile，输出更有分析价值的报告。

用法:
    python PROF/1.py

输入（PROF/ 目录下，每次编译自动清理并重新生成）:
    worker_prof_*.prof   —— worker 子进程每任务 cProfile（Phase2Translator._ProfileWorker）
    main_prof_*.prof     —— 主进程整次编译 cProfile（Projectrans.py 入口）

输出:
    1.txt                        —— 综合报告（总览 + tottime top + cumtime top +
                                     热点函数 callers 定位 + 模块聚合）
    combined_all.raw.prof        —— 全部合并原始 profile（snakeviz/gprof2dot 用）
    combined_worker_cum.txt      —— worker cumulative（全量）
    combined_worker_tottime.txt  —— worker tottime（全量）
"""
import pstats
from pathlib import Path

PROF_DIR = Path(__file__).resolve().parent

# 关键热点关键词：命中即在 1.txt 中输出其 callers 定位（按实际热点增减）
HOT_SPOTS = (
    '_path_isfile', '_LookupStubFuncType', '_Init', '_TryLoadClassMembersFromPyi',
    '_TryLoadStructFromStub', 'vast_c.parse', 'finalize', '_LoadSha1Map',
    '_BuildPyiCacheFile', '_translate_file', 'Convert', 'GenerateCCode',
    '_build_current_ModuleSha1Map', 'iter_child_nodes', 'read of',
    'iter_fields', 'walk', '_io.open', '_pickle.loads',
)


def _ModuleKey(fname: str) -> str:
    """将文件路径归并到顶层模块（lib/xxx.py 或 lib/xxx/yyy.py 的 xxx）。"""
    parts = fname.replace('\\', '/').split('/')
    try:
        i = parts.index('lib')
    except ValueError:
        return fname.rsplit('/', 1)[-1]
    # lib 之后的第一段即模块归属；直接文件则取文件名
    if i + 1 < len(parts):
        return parts[i + 1]
    return parts[-1]


def _AggregateByModule(stats: pstats.Stats) -> list[tuple[str, int, float, float]]:
    """按顶层模块聚合 tottime/cumtime（ncalls 取和）。返回按 tottime 降序。"""
    agg: dict[str, list[float]] = {}
    for key, val in stats.stats.items():
        fname = key[0]
        tt, ct = val[2], val[3]
        m = _ModuleKey(fname)
        if m not in agg:
            agg[m] = [0, 0.0, 0.0]  # ncalls, tottime, cumtime
        agg[m][0] += val[0]
        agg[m][1] += tt
        agg[m][2] += ct
    rows = [(m, *v) for m, v in agg.items()]
    rows.sort(key=lambda r: r[2], reverse=True)
    return rows


def _WriteCallers(stats: pstats.Stats, fp, keyword: str) -> None:
    """输出含 keyword 的所有函数的 callers（谁在调它，调用次数）。"""
    hits = [k for k in stats.stats if keyword in k[2]]
    if not hits:
        fp.write(f"  [无匹配 {keyword!r} 的函数]\n")
        return
    for key in sorted(hits, key=lambda k: stats.stats[k][2], reverse=True)[:5]:
        cc, nc, tt, ct, callers = stats.stats[key]
        fp.write(f"  {key[2]}: tottime={tt:.3f}s cumtime={ct:.3f}s ncalls={nc}\n")
        if not callers:
            fp.write("    callers: (无)\n")
            continue
        for ckey, ncc in sorted(callers.items(), key=lambda x: x[1], reverse=True)[:8]:
            fp.write(f"      <- {ckey[2]} x{ncc}\n")


def main() -> None:
    worker_files = sorted(PROF_DIR.glob('worker_prof_*.prof'))
    main_files = sorted(PROF_DIR.glob('main_prof_*.prof'))
    all_files = main_files + worker_files
    if not all_files:
        raise SystemExit(f'PROF 目录无 .prof 文件: {PROF_DIR}')

    # 合并全部（主进程 + worker）
    stats = pstats.Stats(str(all_files[0]))
    for f in all_files[1:]:
        stats.add(str(f))

    stats.dump_stats(str(PROF_DIR / 'combined_all.raw.prof'))

    # ---- worker 单独的全量 cum / tottime（兼容旧脚本输出） ----
    if worker_files:
        wstats = pstats.Stats(str(worker_files[0]))
        for f in worker_files[1:]:
            wstats.add(str(f))
        for name, sort in (('combined_worker_cum.txt', 'cumulative'),
                           ('combined_worker_tottime.txt', 'tottime')):
            with open(PROF_DIR / name, 'w', encoding='utf-8') as fp:
                wstats.stream = fp
                wstats.sort_stats(sort)
                wstats.print_stats()

    # ---- 1.txt 综合报告 ----
    report = PROF_DIR / '1.txt'
    total_tt = sum(v[2] for v in stats.stats.values())
    total_ncalls = sum(v[0] for v in stats.stats.values())
    with open(report, 'w', encoding='utf-8') as fp:
        fp.write('=' * 72 + '\n')
        # 上次编译总耗时（Projectrans.py 入口记录，成功/失败均落盘）
        time_txt = PROF_DIR / 'main_time.txt'
        if time_txt.exists():
            try:
                _status, _secs = time_txt.read_text(encoding='utf-8').strip().split('\t')
                fp.write(f'上次编译状态: {_status}   总耗时: {_secs}s\n')
            except Exception:
                pass
        fp.write(f'PROF 综合分析报告  生成时间: {Path(__file__).stat().st_mtime}\n')
        fp.write(f'主进程 prof: {len(main_files)} 个, worker prof: {len(worker_files)} 个\n')
        fp.write(f'总函数调用 {total_ncalls} 次, 总 tottime {total_tt:.1f}s\n')
        fp.write('=' * 72 + '\n\n')

        fp.write('【一】Top 40 热点函数 (tottime = 真实 CPU 消耗, 优先优化)\n')
        fp.write('-' * 72 + '\n')
        tmp = pstats.Stats(str(all_files[0]))
        for f in all_files[1:]:
            tmp.add(str(f))
        tmp.stream = fp
        tmp.sort_stats('tottime')
        tmp.print_stats(40)
        fp.write('\n')

        fp.write('【二】Top 20 调用链耗时 (cumtime = 含子调用, 看大流程)\n')
        fp.write('-' * 72 + '\n')
        tmp.sort_stats('cumulative')
        tmp.print_stats(20)
        fp.write('\n')

        fp.write('【三】热点函数 callers 定位（找出"谁在调用"以定点优化）\n')
        fp.write('-' * 72 + '\n')
        for kw in HOT_SPOTS:
            fp.write(f'\n== {kw} ==\n')
            _WriteCallers(stats, fp, kw)
        fp.write('\n')

        fp.write('【四】模块聚合耗时（按顶层模块归类, tottime 降序）\n')
        fp.write('-' * 72 + '\n')
        fp.write(f"{'模块':<40}{'ncalls':>10}{'tottime(s)':>12}{'cumtime(s)':>12}\n")
        for m, nc, tt, ct in _AggregateByModule(stats):
            fp.write(f'{m:<40}{nc:>10}{tt:>12.2f}{ct:>12.2f}\n')
        fp.write('\n')

    print(f'报告已生成: {report}')
    print(f'原始合并: {PROF_DIR / "combined_all.raw.prof"}')


if __name__ == '__main__':
    main()