use scarlett::{builder::IrModuleBuilder, ir::*, validate::validate_module};

#[test]
fn test_float_ir() {
    let mut module_builder = IrModuleBuilder::new();
    let f64_ty = Type::F64;

    let func_id = module_builder.new_function_id();
    {
        let f_builder =
            module_builder.new_function(func_id, "float_math", vec![&f64_ty, &f64_ty], f64_ty);

        let a = f_builder.get_param(0).unwrap();
        let b = f_builder.get_param(1).unwrap();

        let sum = f_builder.build_fadd(f64_ty, a, b);
        let diff = f_builder.build_fsub(f64_ty, sum, a);
        let prod = f_builder.build_fmul(f64_ty, diff, b);
        let quot = f_builder.build_fdiv(f64_ty, prod, a);
        let neg = f_builder.build_fneg(f64_ty, quot);

        let is_gt = f_builder.build_fcmp(FCmpOp::Ogt, neg, Operand::Float(0.0));

        let then_block = f_builder.create_block();
        let else_block = f_builder.create_block();

        f_builder.build_branch(is_gt, then_block, else_block);

        f_builder.switch_to_block(then_block);
        f_builder.build_return(f64_ty, neg);

        f_builder.switch_to_block(else_block);
        f_builder.build_return(f64_ty, Operand::Float(42.0));
    }
    module_builder.complete_function();

    let module = module_builder.finish();

    validate_module(&module).expect("Module validation failed");

    let printed = format!("{}", module);
    println!("{}", printed);

    assert!(printed.contains("f64"));
    assert!(printed.contains("fadd"));
    assert!(printed.contains("fsub"));
    assert!(printed.contains("fmul"));
    assert!(printed.contains("fdiv"));
    assert!(printed.contains("fneg"));
    assert!(printed.contains("fcmp ogt"));
}

#[test]
fn test_float_constant_folding() {
    let mut module_builder = IrModuleBuilder::new();
    let f64_ty = Type::F64;

    let func_id = module_builder.new_function_id();
    {
        let f_builder = module_builder.new_function(func_id, "fold_me", vec![], f64_ty);

        let a = Operand::Float(10.0);
        let b = Operand::Float(2.0);

        let sum = f_builder.build_fadd(f64_ty, a, b); // 12.0
        let prod = f_builder.build_fmul(f64_ty, sum, b); // 24.0

        f_builder.build_return(f64_ty, prod);
    }
    module_builder.complete_function();

    let mut module = module_builder.finish();

    validate_module(&module).expect("Module validation failed");

    scarlett::passes::cfp::fold_constants(&mut module);

    let printed = format!("{}", module);
    println!("{}", printed);

    assert!(printed.contains("$24"));
    assert!(!printed.contains("fadd"));
    assert!(!printed.contains("fmul"));
}

#[test]
fn test_float_codegen() {
    let mut module_builder = IrModuleBuilder::new();
    let f64_ty = Type::F64;

    let func_id = module_builder.new_function_id();
    {
        let f_builder =
            module_builder.new_function(func_id, "add_floats", vec![&f64_ty, &f64_ty], f64_ty);

        let a = f_builder.get_param(0).unwrap();
        let b = f_builder.get_param(1).unwrap();

        let res = f_builder.build_fadd(f64_ty, a, b);
        f_builder.build_return(f64_ty, res);
    }
    module_builder.complete_function();

    let module = module_builder.finish();

    let assembly = scarlett::backend::x86_64::X86Backend::new(&module).compile_module();
    println!("{}", assembly);

    assert!(assembly.contains("addsd"));
    assert!(assembly.contains("movsd"));
    assert!(assembly.contains("%xmm0"));
    assert!(assembly.contains("%xmm1"));
}