#!/usr/bin/env python3
import argparse
import json
import os
import re
import resource
import subprocess
import sys
import tempfile
import time
from pathlib import Path

# Regex to parse reachability stats from stderr
REACH_ROW_RE = re.compile(
    r"^\s*(?P<node>\w+):\s+encoded=(?P<encoded>\d+)\s+bailed=(?P<bailed>\d+)"
    r"\s+bail_rate=(?P<rate>[0-9.]+)%"
    r"(?:\s+top_bails=(?P<top>.*))?$"
)

def find_fold_binary() -> str:
    # Try common locations
    candidates = [
        Path("build/release/src/cli/fold"),
        Path("build/src/cli/fold"),
        Path("bin/fold"),
    ]
    for c in candidates:
        if c.exists() and os.access(c, os.X_OK):
            return str(c.resolve())
    return "fold"

def run_benchmark(args):
    fold_bin = args.fold_path or find_fold_binary()
    print(f"Using fold binary: {fold_bin}")

    # Prepare temp files for output and profile
    with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f_out, \
         tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f_prof:
        findings_json_path = f_out.name
        profile_json_path = f_prof.name

    # Build command
    cmd = [
        fold_bin, "scan",
        "--mode", args.mode,
        "-o", findings_json_path,
        "-f", "json",
        "--profile",
        "--profile-output", profile_json_path,
        "--profile-reach"
    ]

    if args.compile_db:
        cmd.extend(["--compile-db", args.compile_db])
    if args.binary:
        cmd.extend(["--binary", args.binary])
    if args.processes is not None:
        cmd.extend(["--processes", str(args.processes)])
    if args.detectors:
        cmd.extend(["--detectors", args.detectors])

    # Append any extra files/folders
    if args.paths:
        cmd.extend(args.paths)

    print(f"Running command: {' '.join(cmd)}")

    # Record resource usage before (child usage accumulates, so we measure the diff)
    usage_before = resource.getrusage(resource.RUSAGE_CHILDREN)
    start_time = time.perf_counter()

    # Run fold scan
    proc = subprocess.run(cmd, capture_output=True, text=True)

    end_time = time.perf_counter()
    usage_after = resource.getrusage(resource.RUSAGE_CHILDREN)

    # Compute resource diffs
    wall_time = end_time - start_time
    user_cpu = usage_after.ru_utime - usage_before.ru_utime
    sys_cpu = usage_after.ru_stime - usage_before.ru_stime
    total_cpu = user_cpu + sys_cpu
    
    # Max RSS: on macOS ru_maxrss is in bytes, on Linux in kilobytes
    max_rss = usage_after.ru_maxrss
    if sys.platform != 'darwin':
        max_rss *= 1024  # convert to bytes

    print(f"Scan finished with exit code {proc.returncode}")

    # Read findings
    findings = []
    if os.path.exists(findings_json_path) and os.path.getsize(findings_json_path) > 0:
        try:
            with open(findings_json_path, "r") as f:
                data = json.load(f)
                if isinstance(data, dict) and "runs" in data:
                    findings = data["runs"][0].get("results", [])
                else:
                    findings = data
        except Exception as e:
            print(f"Failed to read findings JSON: {e}", file=sys.stderr)
    
    # Clean up findings temp file
    if os.path.exists(findings_json_path):
        os.unlink(findings_json_path)

    # Read profile
    profile_data = {}
    if os.path.exists(profile_json_path) and os.path.getsize(profile_json_path) > 0:
        try:
            with open(profile_json_path, "r") as f:
                profile_data = json.load(f)
        except Exception as e:
            print(f"Failed to read profile JSON: {e}", file=sys.stderr)

    # Clean up profile temp file
    if os.path.exists(profile_json_path):
        os.unlink(profile_json_path)

    # Parse reachability stats from stderr
    reach_stats = {}
    in_reach = False
    for line in proc.stderr.splitlines():
        if "[reach-profile] reachability" in line:
            in_reach = True
            continue
        if in_reach:
            if "[reach-profile]" not in line and not line.strip().startswith("[reach-profile]"):
                if line.strip() and not line.startswith(" ") and not line.startswith("\t"):
                    in_reach = False
                    continue
            content = line.replace("[reach-profile]", "", 1).strip()
            if not content:
                continue
            match = REACH_ROW_RE.match(content)
            if match:
                node = match.group("node")
                reach_stats[node] = {
                    "encoded": int(match.group("encoded")),
                    "bailed": int(match.group("bailed")),
                    "rate": float(match.group("rate")),
                    "top": match.group("top").strip() if match.group("top") else ""
                }

    # Aggregate findings by detector type
    findings_by_detector = {}
    for f in findings:
        if isinstance(f, dict):
            fid = f.get("id")
            if fid is None:
                fid = f.get("properties", {}).get("id")
            if fid is None:
                fid = f.get("ruleId", "unknown")
            det = fid.split(":")[0]
        else:
            det = str(f).split(":")[0]
        findings_by_detector[det] = findings_by_detector.get(det, 0) + 1

    # Print summary
    print("\n" + "="*50)
    print(" FOLD BENCHMARK SUMMARY")
    print("="*50)
    print(f"Scan Mode:        {args.mode}")
    print(f"Exit Code:        {proc.returncode}")
    print(f"Wall Time (s):    {wall_time:.3f}")
    print(f"CPU Time (s):     {total_cpu:.3f} (User: {user_cpu:.3f}, Sys: {sys_cpu:.3f})")
    print(f"Peak Memory (MB): {max_rss / (1024*1024):.2f}")
    print(f"Total Findings:   {len(findings)}")
    if findings_by_detector:
        print("Findings by detector:")
        for det, count in findings_by_detector.items():
            print(f"  - {det}: {count}")
    
    if reach_stats:
        print("\nReachability & Bail Rates:")
        for node, stats in reach_stats.items():
            print(f"  - {node}: encoded={stats['encoded']}, bailed={stats['bailed']} ({stats['rate']}%)")

    # Combine all results for JSON output
    report = {
        "mode": args.mode,
        "exit_code": proc.returncode,
        "metrics": {
            "wall_time_seconds": wall_time,
            "cpu_time_user_seconds": user_cpu,
            "cpu_time_sys_seconds": sys_cpu,
            "cpu_time_total_seconds": total_cpu,
            "peak_memory_bytes": max_rss,
        },
        "findings": {
            "total": len(findings),
            "by_detector": findings_by_detector
        },
        "reachability": reach_stats,
        "profile": profile_data
    }

    if args.output:
        try:
            with open(args.output, "w") as f:
                json.dump(report, f, indent=2)
            print(f"\nReport successfully saved to {args.output}")
        except Exception as e:
            print(f"Failed to write output report: {e}", file=sys.stderr)

    return proc.returncode == 0 or proc.returncode == 1

def main():
    parser = argparse.ArgumentParser(description="Benchmark fold scanner performance.")
    parser.add_argument("--fold-path", help="Path to fold executable")
    parser.add_argument("--mode", choices=["source", "binary", "hybrid"], default="source",
                        help="Scan mode (default: source)")
    parser.add_argument("--compile-db", help="Path to compile_commands.json")
    parser.add_argument("--binary", help="Path to binary/firmware for binary analysis")
    parser.add_argument("-j", "--processes", type=int, help="Number of worker processes")
    parser.add_argument("--detectors", help="Detectors to run (comma separated)")
    parser.add_argument("-o", "--output", help="Write JSON benchmark results to this file")
    parser.add_argument("paths", nargs="*", help="Optional source file list")

    args = parser.parse_args()

    # Validate arguments based on mode
    if args.mode in ["source", "hybrid"] and not args.compile_db:
        # Check if compile_commands.json exists in current folder
        if os.path.exists("build/release/compile_commands.json"):
            args.compile_db = "build/release/compile_commands.json"
        else:
            parser.error("--compile-db is required for source or hybrid mode")
    
    if args.mode in ["binary", "hybrid"] and not args.binary:
        parser.error("--binary is required for binary or hybrid mode")

    success = run_benchmark(args)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()