#!/usr/bin/env python3
"""#2 决策:流式 get_edges(现状)vs 批量 afl-showmap -i dir。同一批输入,比总耗时 + 覆盖一致性。"""
import sys, os, time, subprocess, tempfile, glob
sys.path.insert(0, "util")
import mpi_fuzzing_helper as m

SHOWMAP = "/usr/local/bin/afl-showmap"
TARGET = "/tmp/synth"
INDIR = "/tmp/synth_in"
files = sorted(glob.glob(os.path.join(INDIR, "*")))[:600]
# 只保留子集目录给批量模式用
SUBDIR = "/tmp/synth_in600"; os.makedirs(SUBDIR, exist_ok=True)
import shutil
for f in files:
    shutil.copy2(f, os.path.join(SUBDIR, os.path.basename(f)))
INDIR = SUBDIR
contents = [open(f, "rb").read() for f in files]
n = len(contents)
print(f"输入 = {n}", flush=True)

def edgeset(e):
    return set(eid for eid, _ in e) if e else set()

# --- 流式(现状):StreamingShowmap.get_edges 逐个 ---
def run_stream():
    sm = m.StreamingShowmap(SHOWMAP, [TARGET])
    cov = set(); t0 = time.monotonic()
    for c in contents:
        e = sm.get_edges(c)
        if e:
            cov |= edgeset(e)
    dt = time.monotonic() - t0
    sm.close()
    return dt, len(cov)

# --- 批量:afl-showmap -i dir -o mapdir 一次,再读所有 map ---
def run_batch():
    mapdir = tempfile.mkdtemp(prefix="bm_")
    t0 = time.monotonic()
    cmd = [SHOWMAP, "-i", INDIR, "-o", mapdir, "-t", "5000", "-m", "none", "-q", "--", TARGET]
    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    cov = set()
    for mf in os.listdir(mapdir):
        p = os.path.join(mapdir, mf)
        try:
            for line in open(p):                 # afl-showmap -o 文件:每行 "edge_id:count"
                line = line.strip()
                if line:
                    cov.add(int(line.split(":")[0]))
        except (IOError, ValueError):
            pass
    dt = time.monotonic() - t0
    import shutil; shutil.rmtree(mapdir, ignore_errors=True)
    return dt, len(cov), len(os.listdir(mapdir)) if os.path.isdir(mapdir) else 0

print("跑流式...", flush=True); s_t, s_cov = run_stream()
print(f"  流式 {s_t*1000:.0f}ms cov={s_cov}", flush=True)
print("跑批量...", flush=True); b_t, b_cov, nmaps = run_batch()
print(f"  批量 {b_t*1000:.0f}ms cov={b_cov} maps={nmaps}", flush=True)
print(f"\n{'方式':8} {'总耗时':>9} {'每输入':>9} {'覆盖边':>8}")
print(f"{'流式':8} {s_t*1000:>7.0f}ms {s_t/n*1e6:>7.0f}us {s_cov:>8}")
print(f"{'批量-i':8} {b_t*1000:>7.0f}ms {b_t/n*1e6:>7.0f}us {b_cov:>8}")
faster = "批量" if b_t < s_t else "流式"
print(f"\n更快 = {faster} (差 {max(s_t,b_t)/max(1e-9,min(s_t,b_t)):.2f}×);覆盖接近 = {abs(s_cov-b_cov)<=2} ({s_cov} vs {b_cov})")