use std::{env, fs};

use rasm_core::codegen::aot::{ElfWriter, Formats, MachOWriter, PeWriter, write_native_executable};

use super::helper::compile;

fn read_u16(bytes: &[u8], offset: usize) -> u16 {
    u16::from_le_bytes(bytes[offset..offset + 2].try_into().unwrap())
}

fn read_u32(bytes: &[u8], offset: usize) -> u32 {
    u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
}

fn read_u64(bytes: &[u8], offset: usize) -> u64 {
    u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap())
}

fn contains_bytes(bytes: &[u8], needle: &[u8]) -> bool {
    bytes.windows(needle.len()).any(|window| window == needle)
}

#[test]
fn pe_writer_emits_headers_and_dynamic_imports() {
    let output = compile(
        r#"
            .import kernel32.dll GetCurrentProcessId
            _start:
                sub rsp, 40
                call GetCurrentProcessId
                add rsp, 40
                ret
        "#,
    )
    .expect("program with Windows import must compile");

    let mut bytes = Vec::new();
    PeWriter::new(&output.module)
        .write(&mut bytes)
        .expect("PE writer must emit bytes");

    assert_eq!(&bytes[0..2], b"MZ");

    let pe_offset = read_u32(&bytes, 0x3c) as usize;
    assert_eq!(&bytes[pe_offset..pe_offset + 4], b"PE\0\0");
    assert_eq!(read_u16(&bytes, pe_offset + 4), 0x8664);

    let optional_header = pe_offset + 4 + 20;
    assert_eq!(read_u16(&bytes, optional_header), 0x020b);
    assert_ne!(read_u32(&bytes, optional_header + 16), 0);

    let import_dir_rva = read_u32(&bytes, optional_header + 112 + 8);
    let iat_rva = read_u32(&bytes, optional_header + 112 + 12 * 8);
    assert_ne!(import_dir_rva, 0);
    assert_ne!(iat_rva, 0);

    assert!(contains_bytes(&bytes, b".text"));
    assert!(contains_bytes(&bytes, b".idata"));
    assert!(contains_bytes(&bytes, b"kernel32.dll"));
    assert!(contains_bytes(&bytes, b"GetCurrentProcessId"));
    assert!(contains_bytes(&bytes, b"ExitProcess"));
}

#[test]
fn elf_writer_emits_elf64_executable_with_entry_segment() {
    let output = compile(
        r#"
            _start:
                mov rax, 7
                ret
        "#,
    )
    .expect("program must compile");

    let mut bytes = Vec::new();
    ElfWriter::new(&output.module)
        .write(&mut bytes)
        .expect("ELF writer must emit bytes");

    assert_eq!(&bytes[0..4], b"\x7fELF");
    assert_eq!(bytes[4], 2);
    assert_eq!(bytes[5], 1);
    assert_eq!(read_u16(&bytes, 16), 2);
    assert_eq!(read_u16(&bytes, 18), 62);

    let entry = read_u64(&bytes, 24);
    assert_eq!(entry, 0x401000);
    assert_eq!(read_u64(&bytes, 32), 64);
    assert_eq!(read_u16(&bytes, 54), 56);
    assert_eq!(read_u16(&bytes, 56), 1);

    let first_phdr = 64;
    assert_eq!(read_u32(&bytes, first_phdr), 1);
    assert_eq!(read_u32(&bytes, first_phdr + 4), 5);
    assert_eq!(read_u64(&bytes, first_phdr + 8), 0x1000);
    assert_eq!(read_u64(&bytes, first_phdr + 16), 0x401000);
    assert!(bytes.len() > 0x1000);
}

#[test]
fn bin_writer_writes_sections_in_module_order() {
    let output = compile(
        r#"
            _start:
                ret
            .data
                db 1, 2, 3
        "#,
    )
    .expect("program with data section must compile");

    let mut path = env::temp_dir();
    path.push(format!(
        "rasm-aot-bin-test-{}-{}.bin",
        std::process::id(),
        output.module.sections.len()
    ));

    write_native_executable(
        &output.module,
        path.to_str().expect("temp path must be UTF-8"),
        Formats::BIN,
    )
    .expect("BIN writer must write file");

    let bytes = fs::read(&path).expect("BIN output must be readable");
    let _ = fs::remove_file(&path);

    assert_eq!(bytes, [0xc3, 1, 2, 3]);
}

#[test]
fn macho_writer_emits_macho64_executable() {
    let output = compile(
        r#"
            _start:
                mov rax, 7
                ret
        "#,
    )
    .expect("program must compile");

    let mut bytes = Vec::new();
    MachOWriter::new(&output.module)
        .write(&mut bytes)
        .expect("Mach-O writer must emit bytes");

    assert_eq!(read_u32(&bytes, 0), 0xFEED_FACF);
    assert_eq!(read_u32(&bytes, 4), 0x0100_0007);
    assert_eq!(read_u32(&bytes, 8), 3);
    assert_eq!(read_u32(&bytes, 12), 2);

    let ncmds = read_u32(&bytes, 16);
    assert!(ncmds == 4 || ncmds == 5);

    let sizeofcmds = read_u32(&bytes, 20);

    let pagezero_off = 32;
    assert_eq!(read_u32(&bytes, pagezero_off), 0x19);
    assert_eq!(read_u32(&bytes, pagezero_off + 4), 72);
    assert_eq!(&bytes[pagezero_off + 8..pagezero_off + 18], b"__PAGEZERO");

    let text_off = pagezero_off + 72;
    assert_eq!(read_u32(&bytes, text_off), 0x19);
    assert_eq!(&bytes[text_off + 8..text_off + 14], b"__TEXT");

    let text_vmaddr = read_u64(&bytes, text_off + 24);
    assert_eq!(text_vmaddr, 0x1_0000_0000);

    let main_off = 32 + sizeofcmds as usize - 24 - 24;
    assert_eq!(read_u32(&bytes, main_off), 0x28);
    assert_eq!(read_u32(&bytes, main_off + 4), 24);

    let build_off = main_off + 24;
    assert_eq!(read_u32(&bytes, build_off), 0x32);
    assert_eq!(read_u32(&bytes, build_off + 4), 24);
    assert_eq!(read_u32(&bytes, build_off + 8), 1);

    assert!(bytes.len() > 0x1000, "Mach-O simple: {} bytes", bytes.len());
}

#[test]
fn macho_writer_emits_peekable_macho_with_data_section() {
    let output = compile(
        r#"
            .data
            msg:
                .ascii "Hello"
            .text
            _start:
                lea rax, [rip + msg]
                ret
        "#,
    )
    .expect("program with data must compile");

    let mut bytes = Vec::new();
    MachOWriter::new(&output.module)
        .write(&mut bytes)
        .expect("Mach-O writer must emit bytes");

    assert_eq!(read_u32(&bytes, 0), 0xFEED_FACF);
    assert_eq!(read_u32(&bytes, 12), 2);

    let ncmds = read_u32(&bytes, 16);
    assert_eq!(ncmds, 5);

    assert!(contains_bytes(&bytes, b"__text"));
    assert!(contains_bytes(&bytes, b"__data"));
}
