"""
Benchmark: Redline analysis speed vs function complexity.

Generates Rust source files of varying complexity (simple, medium, complex),
compiles them with Redline's proc macro, and measures wall-clock compile time.
Compares against a baseline (no #[redline] annotation).
"""

import subprocess
import tempfile
import time
import os
import json
import shutil

REDLINE_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# --- Code generators ---

def gen_simple(use_redline: bool) -> str:
    """~5 lines, pure arithmetic."""
    attr = '#[redline(latency = "< 1ms", allocs = 0)]' if use_redline else ""
    return f"""
use redline::redline;

{attr}
fn simple_add(a: u64, b: u64) -> u64 {{
    let c = a + b;
    let d = c * 2;
    let e = d - a;
    let f = e / 3;
    f + 1
}}

fn main() {{}}
"""


def gen_medium(use_redline: bool) -> str:
    """~50 lines, loops + branches + method calls."""
    attr = '#[redline(latency = "< 10ms", allocs = "< 20")]' if use_redline else ""
    lines = []
    lines.append("use redline::redline;")
    lines.append("")
    lines.append(attr)
    lines.append("fn medium_work(data: &[u8]) -> u64 {")
    lines.append("    let mut sum: u64 = 0;")
    lines.append("    let mut count: u64 = 0;")
    # 10 sequential for-loops, each with branches
    for i in range(10):
        lines.append(f"    for _i{i} in 0..50 {{")
        lines.append(f"        if sum > {i * 100} {{")
        lines.append(f"            sum += {i + 1};")
        lines.append(f"        }} else {{")
        lines.append(f"            count += 1;")
        lines.append(f"        }}")
        lines.append(f"    }}")
    lines.append("    sum + count")
    lines.append("}")
    lines.append("")
    lines.append("fn main() {}")
    return "\n".join(lines)


def gen_complex(use_redline: bool) -> str:
    """~200+ lines, nested loops, allocations, match, closures."""
    attr = '#[redline(latency = "< 100ms", allocs = "< 100")]' if use_redline else ""
    lines = []
    lines.append("use redline::redline;")
    lines.append("")
    lines.append(attr)
    lines.append("fn complex_work(input: &[u8]) -> u64 {")
    lines.append("    let mut result: u64 = 0;")
    lines.append("    let mut scratch = [0u8; 256];")

    # Nested loops with arithmetic
    for outer in range(5):
        lines.append(f"    for _o{outer} in 0..20 {{")
        for inner in range(3):
            lines.append(f"        for _i{outer}_{inner} in 0..10 {{")
            lines.append(f"            let x = result + {outer * 10 + inner};")
            lines.append(f"            let y = x * 3 + {inner + 1};")
            lines.append(f"            if y > {outer * 100 + 50} {{")
            lines.append(f"                result = y / 2;")
            lines.append(f"            }} else {{")
            lines.append(f"                result += y;")
            lines.append(f"            }}")
            lines.append(f"        }}")
        lines.append(f"    }}")

    # Match expressions
    for m in range(5):
        lines.append(f"    match result % {m + 3} {{")
        for arm in range(m + 3):
            lines.append(f"        {arm} => {{ result += {arm * 7 + 1}; }}")
        lines.append(f"        _ => {{ result += 1; }}")
        lines.append(f"    }}")

    # Index operations
    for i in range(10):
        lines.append(f"    scratch[{i}] = (result % 256) as u8;")
        lines.append(f"    result += scratch[{i}] as u64;")

    # More arithmetic to pad out
    for i in range(30):
        lines.append(f"    let t{i} = result.wrapping_mul({i * 13 + 7});")
        lines.append(f"    result = result.wrapping_add(t{i});")

    lines.append("    result")
    lines.append("}")
    lines.append("")
    lines.append("fn main() {}")
    return "\n".join(lines)


GENERATORS = {
    "simple_5_lines": gen_simple,
    "medium_50_lines": gen_medium,
    "complex_200_lines": gen_complex,
}


def compile_and_time(source: str, label: str) -> tuple:
    """Write source to a temp crate that depends on redline, compile it, return wall-clock seconds."""
    tmpdir = tempfile.mkdtemp(prefix="redline_bench_")
    try:
        src_dir = os.path.join(tmpdir, "src")
        os.makedirs(src_dir)

        # Write Cargo.toml pointing to local redline
        cargo_toml = f"""[package]
name = "redline-bench-{label}"
version = "0.0.1"
edition = "2021"

[dependencies]
redline = {{ package = "alia-redline", path = "{REDLINE_ROOT.replace(os.sep, '/')}" }}
"""
        with open(os.path.join(tmpdir, "Cargo.toml"), "w") as f:
            f.write(cargo_toml)

        # Write source
        with open(os.path.join(src_dir, "main.rs"), "w") as f:
            f.write(source)

        # Warm up: build dependencies once (don't count dep compile time)
        env = os.environ.copy()
        env["PATH"] = os.path.expanduser("~/.cargo/bin") + os.pathsep + env.get("PATH", "")
        subprocess.run(
            ["cargo", "check"],
            cwd=tmpdir,
            capture_output=True,
            env=env,
        )

        # Touch the source to force re-analysis on next compile
        os.utime(os.path.join(src_dir, "main.rs"))

        # Timed run -- only measures analysis of our function, not dep compilation
        t0 = time.perf_counter()
        result = subprocess.run(
            ["cargo", "check"],
            cwd=tmpdir,
            capture_output=True,
            env=env,
        )
        t1 = time.perf_counter()

        elapsed = t1 - t0
        status = "pass" if result.returncode == 0 else "fail"
        return elapsed, status
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


def run_benchmarks():
    results = {}
    print("=== Redline Analysis Speed Benchmark ===\n")

    for name, gen_fn in GENERATORS.items():
        # Baseline: no redline attribute
        baseline_src = gen_fn(use_redline=False)
        baseline_time, baseline_status = compile_and_time(baseline_src, f"{name}_baseline")

        # With redline
        redline_src = gen_fn(use_redline=True)
        redline_time, redline_status = compile_and_time(redline_src, f"{name}_redline")

        overhead_ms = (redline_time - baseline_time) * 1000
        overhead_pct = ((redline_time / baseline_time) - 1) * 100 if baseline_time > 0 else 0

        results[name] = {
            "baseline_s": round(baseline_time, 4),
            "redline_s": round(redline_time, 4),
            "overhead_ms": round(overhead_ms, 2),
            "overhead_pct": round(overhead_pct, 2),
            "baseline_status": baseline_status,
            "redline_status": redline_status,
        }

        print(f"{name}:")
        print(f"  baseline: {baseline_time:.4f}s  redline: {redline_time:.4f}s  overhead: {overhead_ms:.2f}ms ({overhead_pct:.1f}%)")

    return results


if __name__ == "__main__":
    results = run_benchmarks()
    out_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "speed_results.json")
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {out_path}")
