"""PyFlowX 性能基准套件.

用法::

    python -m benchmarks              # 运行全部基准
    python -m benchmarks graph        # 仅图构建基准
    python -m benchmarks execution    # 仅执行基准
    python -m benchmarks context      # 仅上下文注入基准
    python -m benchmarks storage      # 仅状态后端基准
    python -m benchmarks advanced     # 仅高级基准
"""

from __future__ import annotations

import math
import sys
import tempfile
from collections.abc import Callable
from pathlib import Path
from typing import Any

from rich.console import Console

import pyflowx as px
from benchmarks import print_results, time_it
from benchmarks.bench_advanced import run_advanced
from pyflowx import Graph, GraphDefaults, RetryPolicy, TaskSpec
from pyflowx._json import _HAS_ORJSON
from pyflowx._json import dumps as _dumps
from pyflowx._json import loads as _loads
from pyflowx.context import build_call_args
from pyflowx.storage import JSONBackend, MemoryBackend, SQLiteBackend

# ============================================================================
# 图生成工具
# ============================================================================


def make_chain(n: int) -> list[TaskSpec]:
    """生成 n 个任务的链式 DAG。"""
    specs = [TaskSpec(f"t{i}", cmd=["true"]) for i in range(n)]
    for i in range(1, n):
        specs[i] = TaskSpec(f"t{i}", cmd=["true"], depends_on=(f"t{i - 1}",))
    return specs


def make_diamond(n: int) -> list[TaskSpec]:
    """生成 n 个任务的菱形 DAG（每层宽度约 sqrt(n)）。"""
    width = max(1, int(math.sqrt(n)))
    specs: list[TaskSpec] = []
    prev_layer: list[str] = []
    layer = 0
    count = 0
    while count < n:
        cur_layer: list[str] = []
        for j in range(width):
            if count >= n:
                break
            name = f"l{layer}_t{j}"
            deps = tuple(prev_layer) if prev_layer else ()
            specs.append(TaskSpec(name, cmd=["true"], depends_on=deps))
            cur_layer.append(name)
            count += 1
        prev_layer = cur_layer
        layer += 1
    return specs


def make_wide(n: int) -> list[TaskSpec]:
    """生成 n 个独立任务（无依赖，最大并行度）。"""
    return [TaskSpec(f"t{i}", cmd=["true"]) for i in range(n)]


# ============================================================================
# 基准：图构建
# ============================================================================


def bench_construction() -> None:
    """图构建（from_specs + validate）基准。"""
    results = []
    for n in (10, 100, 500, 1000):
        specs = make_chain(n)
        ms, _ = time_it(lambda s=specs: Graph.from_specs(s), iterations=20)
        results.append((f"chain({n})", 20, ms, 1.0 / (ms / 1000) if ms > 0 else 0))

    for n in (10, 100, 500, 1000):
        specs = make_diamond(n)
        ms, _ = time_it(lambda s=specs: Graph.from_specs(s), iterations=20)
        results.append((f"diamond({n})", 20, ms, 1.0 / (ms / 1000) if ms > 0 else 0))

    print_results("图构建 (from_specs + validate)", results)


def bench_layers() -> None:
    """拓扑分层基准（冷启动 vs 缓存命中）。"""
    results = []
    for n in (100, 500, 1000):
        specs = make_diamond(n)
        graph = Graph.from_specs(specs)

        def _cold(g: Graph = graph) -> None:
            g._layers_cache = None  # type: ignore[attr-defined]
            g.layers()

        ms_cold, ops_cold = time_it(_cold, iterations=50, warmup=5)
        results.append((f"layers(cold,{n})", 50, ms_cold, ops_cold))

        ms_hot, ops_hot = time_it(lambda g=graph: g.layers(), iterations=200, warmup=10)
        results.append((f"layers(cached,{n})", 200, ms_hot, ops_hot))

    print_results("拓扑分层 (layers)", results)


def bench_resolved_spec() -> None:
    """resolved_spec 缓存命中基准。"""
    results = []
    for n in (100, 500, 1000):
        specs = make_chain(n)
        defaults = GraphDefaults(retry=RetryPolicy(max_attempts=2))
        graph = Graph.from_specs(specs, defaults=defaults)
        name = f"t{n // 2}"
        ms, ops = time_it(lambda g=graph, nm=name: g.resolved_spec(nm), iterations=500, warmup=20)
        results.append((f"resolved_spec(cached,{n})", 500, ms, ops))

    print_results("resolved_spec (缓存命中)", results)


def run_graph() -> None:
    """运行全部图基准。"""
    bench_construction()
    bench_layers()
    bench_resolved_spec()


# ============================================================================
# 基准：任务执行
# ============================================================================


def bench_sequential() -> None:
    """sequential 策略执行基准。"""
    results = []

    def noop() -> None:
        pass

    for n in (50, 200, 500):
        specs = [TaskSpec(f"t{i}", fn=noop) for i in range(n)]
        graph = Graph.from_specs(specs)
        ms, ops = time_it(lambda g=graph: px.run(g, strategy="sequential"), iterations=10, warmup=2)
        results.append((f"sequential({n})", 10, ms, ops))

    print_results("执行策略: sequential", results)


def bench_thread() -> None:
    """thread 策略执行基准。"""
    results = []

    def noop() -> None:
        pass

    for n in (50, 200, 500):
        specs = [TaskSpec(f"t{i}", fn=noop) for i in range(n)]
        graph = Graph.from_specs(specs)
        ms, ops = time_it(lambda g=graph: px.run(g, strategy="thread"), iterations=10, warmup=2)
        results.append((f"thread({n})", 10, ms, ops))

    print_results("执行策略: thread", results)


def bench_async() -> None:
    """async 策略执行基准。"""
    results = []

    def noop() -> None:
        pass

    for n in (50, 200, 500):
        specs = [TaskSpec(f"t{i}", fn=noop) for i in range(n)]
        graph = Graph.from_specs(specs)
        ms, ops = time_it(lambda g=graph: px.run(g, strategy="async"), iterations=10, warmup=2)
        results.append((f"async({n})", 10, ms, ops))

    print_results("执行策略: async", results)


def bench_dependency() -> None:
    """dependency 策略执行基准。"""
    results = []

    def noop() -> None:
        pass

    for n in (50, 200, 500):
        specs = [TaskSpec(f"t{i}", fn=noop) for i in range(n)]
        graph = Graph.from_specs(specs)
        ms, ops = time_it(lambda g=graph: px.run(g, strategy="dependency"), iterations=10, warmup=2)
        results.append((f"dependency({n})", 10, ms, ops))

    print_results("执行策略: dependency", results)


def bench_large_graph() -> None:
    """大图调度基准：验证增量就绪集优化效果（1k–10k 任务）。

    覆盖三种图形状：
    * chain  —— 深链（每轮仅 1 个就绪，调度轮数 = N）
    * wide   —— 完全并行（首轮全部就绪，仅 1 轮调度）
    * diamond —— 菱形（多层，每层多任务）
    """
    results = []

    def noop() -> None:
        pass

    for shape, maker in (("chain", make_chain), ("diamond", make_diamond), ("wide", make_wide)):
        for n in (1000, 5000, 10000):
            specs = maker(n)
            # 替换为 noop fn（避免子进程开销，纯测调度性能）
            specs = [TaskSpec(s.name, fn=noop, depends_on=s.depends_on) for s in specs]
            graph = Graph.from_specs(specs)
            ms, ops = time_it(lambda g=graph: px.run(g, strategy="dependency"), iterations=3, warmup=1)
            results.append((f"dependency-{shape}({n})", 3, ms, ops))

    print_results("大图调度 (dependency, 增量就绪集)", results)


def bench_cmd_execution() -> None:
    """cmd 任务执行基准（真实子进程）。"""
    results = []
    for n in (10, 50, 100):
        specs = [TaskSpec(f"t{i}", cmd=["true"]) for i in range(n)]
        graph = Graph.from_specs(specs)
        ms, ops = time_it(lambda g=graph: px.run(g, strategy="sequential"), iterations=5, warmup=1)
        results.append((f"cmd-sequential({n})", 5, ms, ops))

    for n in (10, 50, 100):
        specs = [TaskSpec(f"t{i}", cmd=["true"]) for i in range(n)]
        graph = Graph.from_specs(specs)
        ms, ops = time_it(lambda g=graph: px.run(g, strategy="thread", max_workers=8), iterations=5, warmup=1)
        results.append((f"cmd-thread({n})", 5, ms, ops))

    print_results("cmd 任务执行 (['true'])", results)


def run_execution() -> None:
    """运行全部执行基准。"""
    bench_sequential()
    bench_thread()
    bench_async()
    bench_dependency()
    bench_large_graph()
    bench_cmd_execution()


# ============================================================================
# 基准：上下文注入
# ============================================================================


def bench_context_no_deps() -> None:
    """无依赖 fn 任务的上下文注入基准。"""
    results = []

    def noop() -> None:
        pass

    spec = TaskSpec("noop", fn=noop)
    context: dict[str, Any] = {}
    ms, ops = time_it(lambda s=spec, c=context: build_call_args(s, c), iterations=2000, warmup=100)
    results.append(("fn(no-deps)", 2000, ms, ops))

    # cmd 任务快速路径
    spec_cmd = TaskSpec("cmd", cmd=["true"])
    ms, ops = time_it(lambda s=spec_cmd, c=context: build_call_args(s, c), iterations=2000, warmup=100)
    results.append(("cmd(fast-path)", 2000, ms, ops))

    print_results("上下文注入 (build_call_args)", results)


def bench_context_with_deps() -> None:
    """有依赖 fn 任务的上下文注入基准。"""
    results = []

    def consumer(a: int, b: int) -> int:
        return a + b

    spec = TaskSpec("consumer", fn=consumer, depends_on=("a", "b"))
    context = {"a": 1, "b": 2, "c": 3}
    ms, ops = time_it(lambda s=spec, c=context: build_call_args(s, c), iterations=2000, warmup=100)
    results.append(("fn(2-deps)", 2000, ms, ops))

    # Context 标注
    from pyflowx.task import Context

    def ctx_fn(ctx: Context) -> int:
        return sum(ctx.values())

    spec_ctx = TaskSpec("ctx", fn=ctx_fn, depends_on=("a", "b"))
    ms, ops = time_it(lambda s=spec_ctx, c=context: build_call_args(s, c), iterations=2000, warmup=100)
    results.append(("fn(Context-annotated)", 2000, ms, ops))

    # **kwargs
    def kwargs_fn(**kwargs: int) -> int:
        return sum(kwargs.values())

    spec_kw = TaskSpec("kw", fn=kwargs_fn, depends_on=("a", "b"))
    ms, ops = time_it(lambda s=spec_kw, c=context: build_call_args(s, c), iterations=2000, warmup=100)
    results.append(("fn(**kwargs)", 2000, ms, ops))

    print_results("上下文注入 (有依赖)", results)


def run_context() -> None:
    """运行全部上下文注入基准。"""
    bench_context_no_deps()
    bench_context_with_deps()


# ============================================================================
# 基准：状态后端
# ============================================================================


def bench_storage() -> None:
    """状态后端 save/load 基准。"""
    results = []

    # MemoryBackend
    mem = MemoryBackend()
    ms, ops = time_it(lambda: mem.save("key", "value"), iterations=1000, warmup=50)
    results.append(("MemoryBackend.save", 1000, ms, ops))

    ms, ops = time_it(mem.load, iterations=1000, warmup=50)
    results.append(("MemoryBackend.load", 1000, ms, ops))

    # JSONBackend（batch 模式）
    tmp_dir = tempfile.mkdtemp()
    json_path = str(Path(tmp_dir) / "state.json")

    json_backend = JSONBackend(json_path)
    with json_backend.batch():
        for i in range(100):
            json_backend.save(f"task_{i}", f"result_{i}")

    def _json_save() -> None:
        jb = JSONBackend(json_path)
        with jb.batch():
            for i in range(10):
                jb.save(f"bench_{i}", f"val_{i}")

    ms, ops = time_it(_json_save, iterations=50, warmup=5)
    results.append(("JSONBackend.save(batch=10)", 50, ms, ops))

    # 复杂 value（嵌套 dict）—— 展示 batch 模式延迟验证的优化效果
    complex_value = {
        "output": {"path": "/data/result.json", "size": 1024, "checksum": "abc123"},
        "metrics": {"accuracy": 0.95, "latency_ms": 120, "samples": 1000},
        "artifacts": [f"artifact_{j}.bin" for j in range(10)],
        "metadata": {"version": "1.0", "tags": ["trained", "validated"], "created_at": "2026-07-07"},
    }

    def _json_save_complex() -> None:
        jb = JSONBackend(json_path)
        with jb.batch():
            for i in range(10):
                jb.save(f"bench_complex_{i}", complex_value)

    ms, ops = time_it(_json_save_complex, iterations=50, warmup=5)
    results.append(("JSONBackend.save(batch=10,complex)", 50, ms, ops))

    ms, ops = time_it(json_backend.load, iterations=200, warmup=10)
    results.append(("JSONBackend.load", 200, ms, ops))

    # SQLiteBackend
    db_path = str(Path(tmp_dir) / "state.db")

    sqlite_backend = SQLiteBackend(db_path)
    with sqlite_backend.batch():
        for i in range(100):
            sqlite_backend.save(f"task_{i}", f"result_{i}")

    def _sqlite_save() -> None:
        sb = SQLiteBackend(db_path)
        with sb.batch():
            for i in range(10):
                sb.save(f"bench_{i}", f"val_{i}")

    ms, ops = time_it(_sqlite_save, iterations=50, warmup=5)
    results.append(("SQLiteBackend.save(batch=10)", 50, ms, ops))

    ms, ops = time_it(sqlite_backend.load, iterations=200, warmup=10)
    results.append(("SQLiteBackend.load", 200, ms, ops))

    print_results("状态后端 (save/load)", results)

    # 序列化基准：_json 抽象层 dumps/loads
    ser_results = []
    report_like = {
        "run_id": "abc12345",
        "success": True,
        "results": [
            {
                "name": f"task_{i}",
                "status": "success",
                "attempts": 1,
                "duration_seconds": 0.123 + i * 0.001,
                "value": {"output": [i, i + 1, i + 2], "meta": {"tag": "api"}},
            }
            for i in range(100)
        ],
    }
    serialized = _dumps(report_like)

    ms, ops = time_it(lambda: _dumps(report_like), iterations=200, warmup=10)
    ser_results.append(("dumps(report-100)", 200, ms, ops))

    ms, ops = time_it(lambda: _loads(serialized), iterations=200, warmup=10)
    ser_results.append(("loads(report-100)", 200, ms, ops))

    backend_name = "orjson" if _HAS_ORJSON else "stdlib"
    print_results(f"JSON 序列化 (后端={backend_name})", ser_results)

    # 清理临时目录
    import shutil

    shutil.rmtree(tmp_dir, ignore_errors=True)


def run_storage() -> None:
    """运行全部状态后端基准。"""
    bench_storage()


# ============================================================================
# 主入口
# ============================================================================


BENCH_MODULES: dict[str, Callable[[], None]] = {
    "graph": run_graph,
    "execution": run_execution,
    "context": run_context,
    "storage": run_storage,
    "advanced": run_advanced,
}


def main(argv: list[str] | None = None) -> int:
    """CLI 入口。"""
    args = argv if argv is not None else sys.argv[1:]
    console = Console()
    console.print("[bold cyan]PyFlowX 性能基准套件[/bold cyan]\n")

    if not args or args[0] in ("--all", "-a"):
        for name, fn in BENCH_MODULES.items():
            console.print(f"[bold]运行: {name}[/bold]")
            fn()
    elif args[0] in BENCH_MODULES:
        BENCH_MODULES[args[0]]()
    elif args[0] in ("--help", "-h"):
        console.print("用法: python -m benchmarks [graph|execution|context|storage]")
        console.print("  无参数 = 运行全部基准")
    else:
        console.print(f"[red]未知基准模块: {args[0]}[/red]")
        console.print(f"可用: {', '.join(BENCH_MODULES)}")
        return 1

    return 0


if __name__ == "__main__":
    sys.exit(main())