"""Benchmarks for ads_foundation: proof size, verification time, pipeline overhead.

Generates CSV files and prints tables suitable for inclusion in the ENASE paper.
"""

import csv
import math
import statistics
import sys
import time

sys.setrecursionlimit(50000)  # build_merkle is recursive; needed for 100K

from ads_foundation import (
    ProofTrailMonitor,
    build_merkle,
    generate_proof,
    merkle_verify_proof as verify_proof,
)


# ---------------------------------------------------------------------------
# Benchmark 1: Proof size vs chain length
# ---------------------------------------------------------------------------

def bench_proof_size():
    sizes = [1, 10, 100, 1_000, 10_000, 100_000]
    rows = []
    for n in sizes:
        data = list(range(n))
        tree = build_merkle(data)
        proof = generate_proof(n // 2, tree, n)
        path_len = len(proof.path)
        proof_bytes = path_len * 32  # each sibling hash is 32 bytes
        theoretical = math.ceil(math.log2(n)) if n > 1 else 0
        rows.append((n, path_len, theoretical, proof_bytes))

    print("\n=== Benchmark 1: Proof Size vs Chain Length ===")
    print(f"{'n':>10}  {'path_len':>10}  {'log2(n)':>10}  {'bytes':>10}")
    for n, pl, th, byt in rows:
        print(f"{n:>10,}  {pl:>10}  {th:>10}  {byt:>10}")

    with open("benchmarks/proof_size.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["n", "path_length", "theoretical_log2", "proof_bytes"])
        w.writerows(rows)
    return rows


# ---------------------------------------------------------------------------
# Benchmark 2: Verification time vs chain length
# ---------------------------------------------------------------------------

def bench_verify_time():
    sizes = [10, 100, 1_000, 10_000, 100_000]
    iterations = 1_000
    rows = []
    for n in sizes:
        data = list(range(n))
        tree = build_merkle(data)
        idx = n // 2
        value = data[idx]
        proof = generate_proof(idx, tree, n)

        times_ns = []
        for _ in range(iterations):
            t0 = time.perf_counter_ns()
            verify_proof(value, proof)
            t1 = time.perf_counter_ns()
            times_ns.append(t1 - t0)

        mean_us = statistics.mean(times_ns) / 1_000
        std_us = statistics.stdev(times_ns) / 1_000
        rows.append((n, round(mean_us, 2), round(std_us, 2)))

    print("\n=== Benchmark 2: Verification Time vs Chain Length ===")
    print(f"{'n':>10}  {'mean_us':>10}  {'std_us':>10}")
    for n, m, s in rows:
        print(f"{n:>10,}  {m:>10.2f}  {s:>10.2f}")

    with open("benchmarks/verify_time.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["n", "mean_us", "std_us"])
        w.writerows(rows)
    return rows


# ---------------------------------------------------------------------------
# Benchmark 3: ProofTrail pipeline overhead
# ---------------------------------------------------------------------------

def bench_pipeline_overhead():
    def process_payload_raw(payload: dict) -> dict:
        return {"status": "ok", "size": len(payload)}

    payloads = {
        "small": {"key": "value"},
        "medium": {f"k{i}": f"v{i}" for i in range(100)},
        "large": {f"k{i}": f"v{i}" for i in range(1000)},
    }

    iterations = 1_000
    rows = []

    for label, payload in payloads.items():
        # Raw baseline
        raw_ns = []
        for _ in range(iterations):
            t0 = time.perf_counter_ns()
            process_payload_raw(payload)
            t1 = time.perf_counter_ns()
            raw_ns.append(t1 - t0)

        # With tracking -- fresh monitor per payload to keep chain short
        tracked_ns = []
        for _ in range(iterations):
            mon = ProofTrailMonitor()

            @mon.track_tool_call
            def _tracked(p: dict) -> dict:
                return {"status": "ok", "size": len(p)}

            with mon.session(agent_id="bench"):
                t0 = time.perf_counter_ns()
                _tracked(payload)
                t1 = time.perf_counter_ns()
            tracked_ns.append(t1 - t0)

        raw_mean = statistics.mean(raw_ns) / 1_000
        tracked_mean = statistics.mean(tracked_ns) / 1_000
        overhead = tracked_mean - raw_mean
        pct_of_api = (overhead / 200_000) * 100  # 200ms typical API latency
        rows.append((
            label,
            round(raw_mean, 1),
            round(tracked_mean, 1),
            round(overhead, 1),
            round(pct_of_api, 4),
        ))

    print("\n=== Benchmark 3: ProofTrail Pipeline Overhead ===")
    print(f"{'payload':>10}  {'raw_us':>10}  {'tracked_us':>12}  {'overhead_us':>12}  {'% of 200ms':>12}")
    for label, raw, tracked, oh, pct in rows:
        print(f"{label:>10}  {raw:>10.1f}  {tracked:>12.1f}  {oh:>12.1f}  {pct:>11.4f}%")

    with open("benchmarks/pipeline_overhead.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["payload", "raw_us", "tracked_us", "overhead_us", "pct_of_200ms"])
        w.writerows(rows)
    return rows


# ---------------------------------------------------------------------------

if __name__ == "__main__":
    bench_proof_size()
    bench_verify_time()
    bench_pipeline_overhead()
    print("\nCSV files written to benchmarks/")
