RustAst/src/utils/tester.rs

use crate::ast::environment::Environment;
use regex::Regex;
use std::fs;
use std::time::{Duration, Instant};

pub struct TestResult {
    pub name: String,
    pub success: bool,
    pub message: String,
}

pub struct BenchmarkResult {
    pub name: String,
    pub median: Duration,
    pub baseline: Option<Duration>,
    pub diff_pct: Option<f64>,
    pub status: String,
}

pub fn run_functional_tests() -> Vec<TestResult> {
    run_functional_tests_with_optimization(false)
}

pub fn run_functional_tests_with_optimization(enabled: bool) -> Vec<TestResult> {
    let mut results = Vec::new();
    let entries = fs::read_dir("examples").unwrap();
    let output_re = Regex::new(r";; Output: (.*)").unwrap();

    for entry in entries.filter_map(|e| e.ok()) {
        let mut env = Environment::new(); // Fresh environment per test file
        env.add_search_path(".");
        env.optimization = enabled;
        let path = entry.path();
        if path.extension().is_some_and(|ext| ext == "myc") {
            let content = fs::read_to_string(&path).unwrap();
            let name = path.file_name().unwrap().to_string_lossy().to_string();

            let expected_output = output_re
                .captures(&content)
                .map(|m| m.get(1).unwrap().as_str().trim().to_string());

            if let Some(expected) = expected_output {
                if let Err(e) = env.preload_dependencies(&content, Some(&path)) {
                    results.push(TestResult {
                        name,
                        success: false,
                        message: format!("Dependency Error: {}", e),
                    });
                    continue;
                }

                match env.run_script(&content) {
                    Ok(val) => {
                        let val_str = format!("{}", val);
                        if val_str == expected {
                            results.push(TestResult {
                                name,
                                success: true,
                                message: format!("OK: {}", val_str),
                            });
                        } else {
                            results.push(TestResult {
                                name,
                                success: false,
                                message: format!(
                                    "Opt {}: Expected {}, got {}",
                                    if enabled { "ON" } else { "OFF" },
                                    expected,
                                    val_str
                                ),
                            });
                        }
                    }
                    Err(e) => results.push(TestResult {
                        name,
                        success: false,
                        message: format!(
                            "Opt {}: Error: {}",
                            if enabled { "ON" } else { "OFF" },
                            e
                        ),
                    }),
                }
            }
        }
    }
    results
}

pub fn run_benchmarks(update: bool, filter: Option<&str>) -> Vec<BenchmarkResult> {
    let mut results = Vec::new();
    let entries = fs::read_dir("examples").unwrap();
    let is_release = !cfg!(debug_assertions);
    let baseline_re = Regex::new(r";; Benchmark: ([\d\.]+\w+)").unwrap();
    let repeat_re = Regex::new(r";; Benchmark-Repeat: (\d+)").unwrap();

    for entry in entries.filter_map(|e| e.ok()) {
        let path = entry.path();
        if path.extension().is_none_or(|ext| ext != "myc") {
            continue;
        }

        let name = path.file_name().unwrap().to_string_lossy().to_string();
        if let Some(f) = filter
            && name != f
        {
            continue;
        }

        let content = fs::read_to_string(&path).unwrap();

        let baseline_match = baseline_re.captures(&content);
        if !update && baseline_match.is_none() {
            continue;
        }
        let repeat_match = repeat_re.captures(&content);

        let mut repeats = repeat_match
            .and_then(|m| m.get(1))
            .and_then(|m| m.as_str().parse::<u32>().ok())
            .unwrap_or(1);

        // Compile once for this file (symbols/types are compatible with all fresh environments)
        let initial_env = Environment::new();
        initial_env.add_search_path(".");
        
        if let Err(e) = initial_env.preload_dependencies(&content, Some(&path)) {
            results.push(BenchmarkResult {
                name,
                median: Duration::ZERO,
                baseline: None,
                diff_pct: None,
                status: format!("DEPENDENCY ERROR: {}", e),
            });
            continue;
        }

        let compiled_once = match initial_env.compile(&content).into_result() {
            Ok(c) => c,
            Err(e) => {
                results.push(BenchmarkResult {
                    name,
                    median: Duration::ZERO,
                    baseline: None,
                    diff_pct: None,
                    status: format!("COMPILE ERROR: {}", e),
                });
                continue;
            }
        };

        // Helper to measure sum of VM execution times over N executions in one environment
        let measure_sum =
            |n: u32, node: &crate::ast::compiler::TypedNode| -> Result<Duration, String> {
                let env = Environment::new();
                env.add_search_path(".");
                if let Err(e) = env.preload_dependencies(&content, Some(&path)) {
                    return Err(format!("Dependency Error in measurement: {}", e));
                }
                
                // Link once per sample
                let linked = env.link(node.clone());
                let func = env.instantiate(linked);
                let mut total = Duration::ZERO;
                for _ in 0..n {
                    let start = Instant::now();
                    let _ = (func.func)(&[]);
                    total += start.elapsed();
                }
                Ok(total)
            };

        if update {
            repeats = 1;
            loop {
                match measure_sum(repeats, &compiled_once) {
                    Ok(total) => {
                        if total >= Duration::from_millis(2) || repeats >= 100_000 {
                            break;
                        }
                        let nanos = total.as_nanos().max(1) as f64;
                        let factor = 2_000_000.0 / nanos;
                        repeats = (repeats as f64 * factor).ceil() as u32;
                        repeats = repeats.max(repeats + 1);
                    }
                    Err(e) => {
                        results.push(BenchmarkResult {
                            name: name.clone(),
                            median: Duration::ZERO,
                            baseline: None,
                            diff_pct: None,
                            status: format!("ERROR: {}", e),
                        });
                        break;
                    }
                }
            }
            if results.last().is_some_and(|r| r.name == name) {
                continue;
            }
        }

        let mut runs = Vec::new();
        let mut error = None;
        // Adaptive samples: High repeats need fewer samples for stable median
        let num_samples = if repeats > 1000 {
            10
        } else if repeats > 100 {
            30
        } else {
            100
        };

        for _ in 0..num_samples {
            match measure_sum(repeats, &compiled_once) {
                Ok(d) => runs.push(d),
                Err(e) => {
                    error = Some(e);
                    break;
                }
            }
        }

        if let Some(e) = error {
            results.push(BenchmarkResult {
                name,
                median: Duration::ZERO,
                baseline: None,
                diff_pct: None,
                status: format!("ERROR: {}", e),
            });
            continue;
        }

        runs.sort();
        let median_total = runs[runs.len() / 2];
        let median_single = median_total / repeats;

        if update {
            let new_val = format_duration(median_single);
            let mut updated_content = content.clone();

            // 1. Update/Insert Benchmark
            let bench_line = format!(";; Benchmark: {}", new_val);
            if let Some(m) = baseline_match {
                updated_content = updated_content.replace(m.get(0).unwrap().as_str(), &bench_line);
            } else {
                updated_content = format!("{}\n{}", bench_line, updated_content);
            }

            // 2. Update/Insert/Remove Benchmark-Repeat
            let repeat_line = if repeats > 1 {
                Some(format!(";; Benchmark-Repeat: {}", repeats))
            } else {
                None
            };
            let current_repeat_str = repeat_re
                .captures(&updated_content)
                .map(|m| m.get(0).unwrap().as_str().to_string());

            if let Some(line) = repeat_line {
                if let Some(old_line) = current_repeat_str {
                    updated_content = updated_content.replace(&old_line, &line);
                } else if let Some(m) = baseline_re.captures(&updated_content) {
                    let b_str = m.get(0).unwrap().as_str().to_string();
                    if let Some(pos) = updated_content.find(&b_str) {
                        updated_content.insert_str(pos + b_str.len(), &format!("\n{}", line));
                    }
                }
            } else if let Some(old_line) = current_repeat_str {
                updated_content = updated_content.replace(&format!("{}\n", old_line), "");
                updated_content = updated_content.replace(&old_line, "");
            }

            fs::write(&path, updated_content).unwrap();
            results.push(BenchmarkResult {
                name,
                median: median_single,
                baseline: None,
                diff_pct: None,
                status: format!("UPDATED: {}", new_val),
            });
        } else if let Some(m) = baseline_match {
            let baseline_str = m.get(1).unwrap().as_str();
            let baseline = parse_duration(baseline_str).unwrap();

            let diff = (median_single.as_nanos() as f64 / baseline.as_nanos() as f64) - 1.0;

            let threshold = if is_release { 0.15 } else { 0.30 };
            let status = if median_single > baseline && diff > threshold {
                "FAILED"
            } else {
                "OK"
            };

            results.push(BenchmarkResult {
                name,
                median: median_single,
                baseline: Some(baseline),
                diff_pct: Some(diff * 100.0),
                status: status.to_string(),
            });
        }
    }
    results
}

fn format_duration(d: Duration) -> String {
    if d.as_nanos() < 1000 {
        format!("{}ns", d.as_nanos())
    } else if d.as_micros() < 1000 {
        format!("{:.1}us", d.as_nanos() as f64 / 1000.0)
    } else if d.as_millis() < 1000 {
        format!("{:.1}ms", d.as_micros() as f64 / 1000.0)
    } else {
        format!("{:.1}s", d.as_millis() as f64 / 1000.0)
    }
}

fn parse_duration(s: &str) -> Option<Duration> {
    use std::sync::OnceLock;
    static RE: OnceLock<Regex> = OnceLock::new();
    let re = RE.get_or_init(|| Regex::new(r"([\d\.]+)(\w+)").unwrap());

    let caps = re.captures(s)?;
    let val: f64 = caps[1].parse().ok()?;
    let unit = &caps[2];
    match unit {
        "ns" => Some(Duration::from_nanos(val as u64)),
        "us" => Some(Duration::from_nanos((val * 1000.0) as u64)),
        "ms" => Some(Duration::from_nanos((val * 1_000_000.0) as u64)),
        "s" => Some(Duration::from_nanos((val * 1_000_000_000.0) as u64)),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    #[test]
    #[cfg(not(debug_assertions))]
    fn benchmark_regression_test() {
        use super::*;
        let results = run_benchmarks(false, None);
        let failures: Vec<_> = results.iter().filter(|r| r.status == "FAILED").collect();

        if !failures.is_empty() {
            let error_msg = failures
                .iter()
                .map(|r| {
                    format!(
                        "{}: Median {:?}, Baseline {:?}, Diff {:.2}%",
                        r.name,
                        r.median,
                        r.baseline.unwrap_or_default(),
                        r.diff_pct.unwrap_or(0.0)
                    )
                })
                .collect::<Vec<_>>()
                .join("\n");

            panic!("Performance regression detected:\n{}", error_msg);
        }
    }
}