use prometeu_compiler::frontends::pbs::parser::Parser;
use prometeu_compiler::frontends::pbs::collector::SymbolCollector;
use prometeu_compiler::frontends::pbs::symbols::ModuleSymbols;
use prometeu_compiler::frontends::pbs::lowering::Lowerer;
use prometeu_compiler::ir_core;

#[test]
fn test_basic_lowering() {
    let code = "
        fn add(a: int, b: int): int {
            return a + b;
        }
        fn main() {
            let x = add(10, 20);
        }
    ";
    let mut parser = Parser::new(code, 0);
    let ast = parser.parse_file().expect("Failed to parse");

    let mut collector = SymbolCollector::new();
    let (type_symbols, value_symbols) = collector.collect(&ast).expect("Failed to collect symbols");
    let module_symbols = ModuleSymbols { type_symbols, value_symbols };

    let lowerer = Lowerer::new(&module_symbols);
    let program = lowerer.lower_file(&ast, "test");

    // Verify program structure
    assert_eq!(program.modules.len(), 1);
    let module = &program.modules[0];
    assert_eq!(module.functions.len(), 2);

    let add_func = module.functions.iter().find(|f| f.name == "add").unwrap();
    assert_eq!(add_func.params.len(), 2);
    assert_eq!(add_func.return_type, ir_core::Type::Int);

    // Verify blocks
    assert!(add_func.blocks.len() >= 1);
    let first_block = &add_func.blocks[0];
    // Check for Add instruction
    assert!(first_block.instrs.iter().any(|i| matches!(i, ir_core::Instr::Add)));
}

#[test]
fn test_control_flow_lowering() {
    let code = "
        fn max(a: int, b: int): int {
            if (a > b) {
                return a;
            } else {
                return b;
            }
        }
    ";
    let mut parser = Parser::new(code, 0);
    let ast = parser.parse_file().expect("Failed to parse");

    let mut collector = SymbolCollector::new();
    let (type_symbols, value_symbols) = collector.collect(&ast).expect("Failed to collect symbols");
    let module_symbols = ModuleSymbols { type_symbols, value_symbols };

    let lowerer = Lowerer::new(&module_symbols);
    let program = lowerer.lower_file(&ast, "test");

    let max_func = &program.modules[0].functions[0];
    // Should have multiple blocks for if-else
    assert!(max_func.blocks.len() >= 3); 
}

#[test]
fn test_hip_lowering() {
    let code = "
        fn test_hip() {
            let g = alloc int;
            mutate g as x {
                let y = x + 1;
            }
        }
    ";
    let mut parser = Parser::new(code, 0);
    let ast = parser.parse_file().expect("Failed to parse");

    let mut collector = SymbolCollector::new();
    let (type_symbols, value_symbols) = collector.collect(&ast).expect("Failed to collect symbols");
    let module_symbols = ModuleSymbols { type_symbols, value_symbols };

    let lowerer = Lowerer::new(&module_symbols);
    let program = lowerer.lower_file(&ast, "test");

    let func = &program.modules[0].functions[0];
    let instrs: Vec<_> = func.blocks.iter().flat_map(|b| b.instrs.iter()).collect();
    
    assert!(instrs.iter().any(|i| matches!(i, ir_core::Instr::Alloc)));
    assert!(instrs.iter().any(|i| matches!(i, ir_core::Instr::ReadGate)));
    assert!(instrs.iter().any(|i| matches!(i, ir_core::Instr::WriteGate)));
}