from __future__ import annotations
import argparse
import io
import shutil
import sys
import tempfile
import traceback

from engine.harness import calibrate_cprofile_overhead, collect_environment, load_source
from engine.schema import SCHEMA_VERSION, AnalysisResult, Calibration, FlameNode, WallTime
from engine.structure import profile_and_measure


# Result 文件大小防护（防用户脚本异常把磁盘撑爆）
# - error.message 上限 1KB:1KB 一般足够覆盖正常异常说明
# - error.traceback 上限 8KB:100 帧左右 traceback 的体量,UI 可读范围
# - functions 列表上限 5000:UI 表格默认只看前 200,再多只是给后续解析加负载
_MAX_ERROR_MESSAGE_BYTES = 1024
_MAX_TRACEBACK_BYTES = 8 * 1024
_MAX_FUNCTIONS = 5000


def _truncate_text(text: str, max_bytes: int) -> str:
    """把文本按 UTF-8 字节截断到 max_bytes 以内,补一个明确截断标记。

    直接 [:max_chars] 在多字节字符中间切会破坏 UTF-8,UI 解析挂。encode → 截
    → 走 errors='ignore' 让不完整尾部自然掉落 → decode,保证结果仍是合法 UTF-8。
    """
    if not text:
        return text
    encoded = text.encode("utf-8", errors="replace")
    if len(encoded) <= max_bytes:
        return text
    truncated = encoded[:max_bytes].decode("utf-8", errors="ignore")
    # 留 64 字节给标记(总长仍 ≤ max_bytes)
    suffix = f"\n\n[... 截断:共 {len(encoded)} 字节,仅保留前 {max_bytes} 字节 ...]"
    # 如果 suffix 太长导致总超,再砍一次
    return (truncated + suffix)[:max_bytes]


def _cap_functions(functions: list, flame: FlameNode | None) -> tuple:
    """截断 functions 列表到 _MAX_FUNCTIONS,并把超出部分累加到 flame 的 root。

    为什么不只截 functions 不动 flame?
    - flame 是 UI 看到的"时间分布",丢了 functions 里那些小函数会让 flame
      总和 < 用户报的总耗时,UI 显示出现 100% 但 sum 不到 80% 的诡异情况。
    - 把被丢函数的 tottime 加到 root.value 上,UI 那边仍按 flame 内部占比显示,
      多出来的部分归到 root 节点（最大最宽的 tile）,用户视觉上只是「其他」变
      大了一点 —— 总和不撒谎。
    """
    if len(functions) <= _MAX_FUNCTIONS:
        return functions, flame
    kept = functions[:_MAX_FUNCTIONS]
    dropped_tottime = sum(f.tottime for f in functions[_MAX_FUNCTIONS:])
    if flame is None:
        flame = FlameNode(name="root", value=dropped_tottime, children=[])
    else:
        flame.value += dropped_tottime
    return kept, flame


def _compute_wall_time(instrumented_sec: float, calibration: Calibration) -> float:
    """用校准 ratio 折算 instrumented 时间为估计的干净耗时。

    思路:cProfile 是插桩式的,wall-time 会被膨胀 ratio 倍(instrumented / clean);
    用同一段 tight-loop 测出 ratio 后,把用户的 instrumented 时间除回去,就是
    「去掉 cProfile 开销后的估计耗时」。

    边界:calibration.ratio 是 0 / 负数 / NaN 时直接信 instrumented 时间 —— 跟
    calibrate_cprofile_overhead 在「校准耗时过短」时回退 ratio=1.0 的兜底策略
    保持一致,UI 显示真实测得的 instrumented 时间而不是除以无效值。
    """
    ratio = calibration.ratio
    # `ratio == ratio` 排除 NaN;`ratio > 0` 排除 0 / 负数
    if ratio == ratio and ratio > 0:
        return instrumented_sec / ratio
    return instrumented_sec


def _cap_result(result: AnalysisResult) -> AnalysisResult:
    """给 result.json 上一道防线,防止异常情况把磁盘/IPC 通道撑爆。

    触发场景:
    - 用户代码 raise 一个带 50MB 文件内容的 ValueError → error.message 撑爆
    - 用户代码触发深递归异常,Python 自动 dump 几千层 traceback → 撑爆
    - 用户脚本 import 一个庞大的库(numpy/pandas)scope=all 把整棵树都归因 →
      38MB+ 的 functions 列表

    截断而非拒绝:UI 仍能看到「有错」「发生了什么」,只是具体细节截断;functions
    截断后仍按 tottime 排序,前 5000 一定是最值得用户关注的热点。
    """
    if result.error:
        msg = result.error.get("message")
        if isinstance(msg, str) and msg:
            result.error["message"] = _truncate_text(msg, _MAX_ERROR_MESSAGE_BYTES)
        tb = result.error.get("traceback")
        if isinstance(tb, str) and tb:
            result.error["traceback"] = _truncate_text(tb, _MAX_TRACEBACK_BYTES)
    if result.functions:
        result.functions, result.flame = _cap_functions(result.functions, result.flame)
    return result


def _progress(phase, pct):
    # 前导 \n：tqdm / 进度条类库会写"working..."到 stderr 但不换行,
    # 直接接 PROGRESS 会拼成 "working...PROGRESS running 10",主进程正则匹配挂、
    # 进度事件掉一档且把工程协议行混进 stderrTail 让用户看到 \r 进度条残片。
    print(f"\nPROGRESS {phase} {pct}", file=sys.stderr, flush=True)


def _emit(result: AnalysisResult, out_path=None):
    """发出结果。

    out_path 给定时结果写入独立文件，stdout/stderr 就完整留给用户代码 ——
    否则用户代码里一句 sys.stdout.write("done") 或 print(..., end='') 就会和
    JSON 挤在同一行，主进程按"最后一行"取 JSON 时解析失败，
    用户看到的却是"结果解析失败"。不给 out_path 时退回 stdout（直接跑 CLI / 老测试）。
    """
    payload = result.to_json()
    if out_path:
        with io.open(out_path, "w", encoding="utf-8", newline="\n") as f:
            f.write(payload)
    else:
        print(payload, flush=True)


def _base(env, config, status, error=None, wall_time=None, functions=None, flame=None, calibration=None):
    return AnalysisResult(
        schemaVersion=SCHEMA_VERSION,
        environment=env,
        config=config,
        status=status,
        error=error,
        wallTime=WallTime(seconds=wall_time, unit="s") if wall_time is not None else None,
        functions=functions or [],
        flame=flame,
        calibration=calibration,
    )


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--script", required=True)
    ap.add_argument("--out", default=None,
                    help="结果 JSON 的输出文件；不给则打到 stdout")
    ap.add_argument("--scope", choices=["user", "all"], default="user",
                    help="剖析范围：user（默认）只归因用户脚本里的函数；"
                         "all 包含所有非 cProfile 内部帧（标准库 + 第三方包 + 用户代码），"
                         "让耗时可以下钻到 import 的包里。")
    ap.add_argument("--workdir", default=None,
                    help="校准文件落盘的临时目录；不给则 runner 自己 mkdtemp。"
                         "profiler-service 一般传过来(它创建的 pyprof-xxx),"
                         "直跑 CLI 不传也行。")
    ap.add_argument("--hide-internal", dest="hide_internal",
                    action=argparse.BooleanOptionalAction, default=True,
                    help="过滤掉测试代码(tests/ / test_*.py / _pyrof_calib 等)+ "
                         "tottime=0 的量化噪声帧。默认 on,用户原话:「软件内部的测试部分"
                         "默认百分百过滤掉」。--no-hide-internal 关掉这条过滤,看完整数据。")
    args = ap.parse_args(argv)

    # 统一走这个 emit：结果去 --out 指定的文件，stdout 留给用户代码。
    # 每条 emit 之前过 _cap_result 给 result.json 上一道防线 —— error.traceback /
    # error.message / functions 列表都设硬上限，防止用户异常把磁盘 / IPC 撑爆。
    def emit(result):
        _emit(_cap_result(result), args.out)

    env = collect_environment()
    config: dict = {}

    # 版本预检：主进程的 validateInterpreter 已经挡过一层，这里是纵深防御——
    # 用户可能绕过 UI 直接调引擎，或 PATH 上的 python 在探测后被换掉。
    # 最低版本 3.8：4 个引擎模块全用了 `from __future__ import annotations`
    # (PEP 563, 3.7+),PEP 604 / PEP 585 注解都字符串化 —— 3.8 也 import 得住。
    # 3.7 故意不锁：dataclass + 字符串化 PEP 585 的边界 case 多,且 3.7 已 EOL 2 年。
    if sys.version_info < (3, 8):
        got = ".".join(str(p) for p in sys.version_info[:3])
        emit(_base(env, config, "runtime_error",
                    error={"type": "PythonTooOld",
                           "message": f"引擎需要 Python 3.8 或更高版本，当前为 {got}"}))
        return

    # 一次 IO 拿到 src：语法预检和后续 profile 都用这份文本
    # 之前没 try/except：临时目录被删 / 路径权限不够 → 进程崩，无 JSON 输出 → UI 显示 "引擎无输出"
    # UnicodeDecodeError 不是 OSError 的子类：二进制文件 / 非 UTF-8 脚本会绕过这层保护再次崩溃。
    try:
        src = load_source(args.script)
    except (OSError, UnicodeDecodeError) as e:
        emit(_base(env, config, "runtime_error",
                    error={"type": type(e).__name__, "message": f"无法读取脚本: {e}"}))
        return

    # Syntax precheck (用文本而非再读一次)
    # compile() 还会在源码含 null 字节时抛 ValueError，也一并按 syntax_error 归类
    # 把 code object 留起来 —— profile_and_measure 还需要再 exec 一次,
    # 直接复用这份省一次 read + compile。
    try:
        code = compile(src, args.script, "exec")
    except (SyntaxError, ValueError) as e:
        emit(_base(env, config, "syntax_error", error={"type": type(e).__name__, "message": str(e)}))
        return

    # v4 单跑架构：校准 + 用户代码 under cProfile 一次拿到归因 + instrumented 时间。
    # 跟 v3 的「先裸跑测 wall_time 再 cProfile 跑」相比，
    # 用户脚本的副作用(plot / file write / etc)只发生一次,语义更直观。
    # workdir 优先用 caller 传的(profiler-service 已经 mkdir pyprof-xxx 落 _pyrof_calib.py
    # 和用户脚本);caller 不传则自己 mkdtemp —— 这种情况跑完自己清掉。
    workdir = args.workdir
    workdir_owned = False
    if not workdir:
        workdir = tempfile.mkdtemp(prefix="pyprof-")
        workdir_owned = True

    try:
        try:
            _progress("calibrating", 5)
            # 校准脚本落在 caller 提供的 workdir 里 —— profiler-service 已经 mkdir
            # pyrof-xxx 装用户脚本,把 _pyrof_calib.py 放一起能让 rmtree(workdir) 一次
            # 清干净。runner 自己 mkdtemp 时(直跑 CLI 场景)workdir_owned=True,
            # calibration 拿到的是同一个目录,临时文件一起被 finally 兜底删。
            calibration = calibrate_cprofile_overhead(workdir)

            _progress("running", 30)
            st, instrumented = profile_and_measure(
                args.script, scope=args.scope, hide_internal=args.hide_internal,
                code=code, src=src,
            )
            wall_time = _compute_wall_time(instrumented, calibration)

            _progress("done", 100)
            emit(_base(env, config, "ok",
                        wall_time=wall_time,
                        functions=st.functions,
                        flame=st.flame,
                        calibration=calibration))
            return
        except SystemExit as e:
            # sys.exit() 抛 SystemExit（BaseException 的子类）—— 不被下面的 except BaseException 捕获，
            # 子进程直接退出无 JSON，UI 就会显示 "引擎无输出 (exit N)"。转成 runtime_error。
            # e.code 默认 None（无参 sys.exit()），文案要单独处理，否则会出现"sys.exit(None)"
            code = e.code if isinstance(e.code, int) else (0 if e.code is None else f"非整数 {e.code!r}")
            emit(_base(env, config, "runtime_error",
                        error={"type": "SystemExit", "message": f"用户代码调用 sys.exit({code})"}))
        except BaseException as e:
            # KeyboardInterrupt / GeneratorExit / 自定义 BaseException 子类 以及
            # 普通 Exception 一并兜底 —— 之前再写一条 `except Exception` 是 dead code
            # （BaseException 是 Exception 的父类,先匹配的赢）。父进程端有 cancel +
            # 进程树 kill 在管 Ctrl+C,这里只负责兜底写 JSON。
            emit(_base(env, config, "runtime_error",
                        error={"type": type(e).__name__, "message": str(e),
                               "traceback": traceback.format_exc()}))
    finally:
        # runner 自己 mkdtemp 的目录要自己清 —— profiler-service 已经传 --workdir 时
        # 整个 workdir 由 caller 管(主进程清理临时目录的逻辑在 before-quit),这里只兜
        # 「直接 python -m engine.runner」的直跑场景。
        if workdir_owned:
            shutil.rmtree(workdir, ignore_errors=True)


if __name__ == "__main__":
    main()