use std::collections::BTreeSet;
use std::path::PathBuf;

use armfortas::driver::OptLevel;
use armfortas::testing::{capture_from_path, CaptureRequest, CapturedStage, Stage};

fn fixture(name: &str) -> PathBuf {
    let path = PathBuf::from("test_programs").join(name);
    assert!(path.exists(), "missing test fixture {}", path.display());
    path
}

fn capture_text(request: CaptureRequest, stage: Stage) -> String {
    let result = capture_from_path(&request).expect("capture should succeed");
    match result.get(stage) {
        Some(CapturedStage::Text(text)) => text.clone(),
        Some(CapturedStage::Run(_)) => panic!("expected text stage for {}", stage.as_str()),
        None => panic!("missing requested stage {}", stage.as_str()),
    }
}

#[test]
fn o2_fully_unrolls_small_do_concurrent_loop() {
    let source = fixture("audit6_do_concurrent.f90");

    let raw_ir = capture_text(
        CaptureRequest {
            input: source.clone(),
            requested: BTreeSet::from([Stage::Ir]),
            opt_level: OptLevel::O2,
        },
        Stage::Ir,
    );
    let opt_ir = capture_text(
        CaptureRequest {
            input: source,
            requested: BTreeSet::from([Stage::OptIr]),
            opt_level: OptLevel::O2,
        },
        Stage::OptIr,
    );

    assert!(
        raw_ir.contains("doconc_check_")
            && raw_ir.contains("doconc_body_")
            && raw_ir.contains("doconc_incr_")
            && raw_ir.contains("doconc_exit_"),
        "lowered IR should preserve DO CONCURRENT loop identity:\n{}",
        raw_ir
    );
    assert!(
        !opt_ir.contains("doconc_check_")
            && !opt_ir.contains("doconc_body_")
            && !opt_ir.contains("doconc_incr_"),
        "O2 optimized IR should fully unroll the small DO CONCURRENT loop:\n{}",
        opt_ir
    );
}