#include <scc_pe_builder.h>

#ifdef LOG_INFO
#undef LOG_INFO
#endif
#define LOG_INFO(...)

#define reserve_align(offset, size) ((offset) + ((size) - 1)) & ~((size) - 1)

void scc_pe_builder_init(scc_pe_builder_t *builder, bool is_64,
                         u32 section_alignment, u32 file_alignment) {
    scc_memset(builder, 0, sizeof(scc_pe_builder_t));
    builder->is_64 = is_64;
    builder->section_alignment = section_alignment;
    builder->file_alignment = file_alignment;
    builder->file_offset = 0;
}

static u32 reserve_file(scc_pe_builder_t *builder, u32 len, u32 align) {
    if (len == 0) {
        return len;
    }

    builder->file_offset = reserve_align(builder->file_offset, align);
    u32 offset = builder->file_offset;
    builder->file_offset += len;
    LOG_INFO("reserve_file %u bytes at %u [%u]", len, offset,
             builder->file_offset);
    return offset;
}

static u32 reserve_virtual(scc_pe_builder_t *builder, u32 len, u32 align) {
    if (len == 0) {
        return len;
    }

    u32 offset = builder->virtual_offset;
    builder->virtual_offset += len;
    builder->virtual_offset = reserve_align(builder->virtual_offset, align);
    LOG_INFO("reserve_virtual %u bytes at %u [%u]", len, offset,
             builder->virtual_offset);
    return offset;
}

static void padding_util(scc_pe_builder_t *builder, u32 len) {
    if (builder->buffer.size > len) {
        LOG_FATAL("padding_util less than buffer size [%u]:[%u]",
                  builder->buffer.size, len);
        return;
    }

    for (usize i = builder->buffer.size; i < len; i++) {
        scc_vec_push(builder->buffer, 0);
    }

    Assert(builder->buffer.size == len);
}

void scc_pe_reserve_dos_header_and_stub(scc_pe_builder_t *builder) {
    Assert(builder->file_offset == 0);
    reserve_file(builder, sizeof(IMAGE_DOS_HEADER), 1);
    reserve_file(builder, 64, 1);
}

void scc_pe_reserve_nt_header(scc_pe_builder_t *builder,
                              usize data_directory_num) {
    builder->nt_headers_offset =
        reserve_file(builder,
                     builder->is_64 ? sizeof(IMAGE_NT_HEADERS64)
                                    : sizeof(IMAGE_NT_HEADERS32),
                     8);

    LOG_INFO("scc_pe_reserve_nt_header %u", builder->nt_headers_offset);
    scc_vec_init(builder->image_data_directory_vec);
    for (usize i = 0; i < data_directory_num; i += 1) {
        scc_vec_push(builder->image_data_directory_vec,
                     (IMAGE_DATA_DIRECTORY){0});
    }
    reserve_file(builder, sizeof(IMAGE_DATA_DIRECTORY) * data_directory_num, 1);
}

void scc_pe_reserve_section_headers(scc_pe_builder_t *builder,
                                    u32 num_of_section) {
    builder->section_header_num = num_of_section;
    reserve_file(builder, sizeof(IMAGE_SECTION_HEADER) * num_of_section, 1);
    builder->file_offset =
        reserve_align(builder->file_offset, builder->file_alignment);
    builder->sizeof_headers = builder->file_offset;
    reserve_virtual(builder, builder->file_offset, builder->section_alignment);
}

void scc_pe_reserve_header(scc_pe_builder_t *builder, u32 num_of_section) {
    scc_pe_reserve_dos_header_and_stub(builder);
    scc_pe_reserve_nt_header(builder, IMAGE_NUMBEROF_DIRECTORY_ENTRIES);
    scc_pe_reserve_section_headers(builder, num_of_section);
}

void scc_pe_write_dos_header_and_stub(scc_pe_builder_t *builder) {
    IMAGE_DOS_HEADER dos_header;
    dos_header.e_magic = IMAGE_DOS_SIGNATURE;
    dos_header.e_cblp = 0x90;
    dos_header.e_cp = 3;
    dos_header.e_crlc = 0;
    dos_header.e_cparhdr = 4;
    dos_header.e_minalloc = 0;
    dos_header.e_maxalloc = 0xFFFF;
    dos_header.e_ss = 0x00;
    dos_header.e_sp = 0xB8;
    dos_header.e_csum = 0;
    dos_header.e_ip = 0x00;
    dos_header.e_cs = 0x00;
    dos_header.e_lfarlc = 0x40;
    dos_header.e_ovno = 0x00;
    for (int i = 0; i < 4; i++) {
        dos_header.e_res[i] = 0x00;
    }
    dos_header.e_oemid = 0x00;
    dos_header.e_oeminfo = 0x00;
    for (int i = 0; i < 10; i++) {
        dos_header.e_res2[i] = 0x00;
    }
    dos_header.e_lfanew = builder->nt_headers_offset;

    for (usize i = 0; i < sizeof(dos_header); ++i) {
        scc_vec_push(builder->buffer, ((u8 *)(&dos_header))[i]);
    }
    static const u8 dos_header_bytes[] = {
        0x0e, 0x1f, 0xba, 0x0e, 0x00, 0xb4, 0x09, 0xcd, 0x21, 0xb8, 0x01,
        0x4c, 0xcd, 0x21, 0x54, 0x68, 0x69, 0x73, 0x20, 0x70, 0x72, 0x6f,
        0x67, 0x72, 0x61, 0x6d, 0x20, 0x63, 0x61, 0x6e, 0x6e, 0x6f, 0x74,
        0x20, 0x62, 0x65, 0x20, 0x72, 0x75, 0x6e, 0x20, 0x69, 0x6e, 0x20,
        0x44, 0x4f, 0x53, 0x20, 0x6d, 0x6f, 0x64, 0x65, 0x2e, 0x0d, 0x0d,
        0x0a, 0x24, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    };
    for (usize i = 0; i < sizeof(dos_header_bytes); i++) {
        scc_vec_push(builder->buffer, ((u8 *)(&dos_header_bytes))[i]);
    }
}

void scc_pe_write_nt_header_with_config(scc_pe_builder_t *builder,
                                        const scc_pe_config_t *config) {
    IMAGE_NT_HEADERS32 nt_headers32 = {0};
    IMAGE_NT_HEADERS64 nt_headers64 = {0};

    // 写入PE签名
    DWORD *signature =
        builder->is_64 ? &nt_headers64.Signature : &nt_headers32.Signature;
    *signature = SCC_LE32(IMAGE_NT_SIGNATURE);

    // 构建文件头
    IMAGE_FILE_HEADER *file_header =
        builder->is_64 ? &nt_headers64.FileHeader : &nt_headers32.FileHeader;
    file_header->Machine = SCC_LE16(config->machine);
    file_header->NumberOfSections = SCC_LE16(builder->section_header_num);
    file_header->TimeDateStamp = SCC_LE32(config->time_date_stamp);
    file_header->PointerToSymbolTable = SCC_LE32(builder->symbol_offset);
    file_header->NumberOfSymbols = SCC_LE32(builder->symbol_num);
    // 设置可选头大小
    if (builder->is_64) {
        file_header->SizeOfOptionalHeader =
            SCC_LE16(sizeof(IMAGE_OPTIONAL_HEADER64) +
                     scc_vec_size(builder->image_data_directory_vec) *
                         sizeof(IMAGE_DATA_DIRECTORY));
    } else {
        file_header->SizeOfOptionalHeader =
            SCC_LE16(sizeof(IMAGE_OPTIONAL_HEADER32) +
                     scc_vec_size(builder->image_data_directory_vec) *
                         sizeof(IMAGE_DATA_DIRECTORY));
    }
    file_header->Characteristics = SCC_LE16(config->characteristics);

    // 根据是否为64位写入不同的可选头
    if (builder->is_64) {
        IMAGE_OPTIONAL_HEADER64 *optional = &nt_headers64.OptionalHeader;
        optional->Magic = SCC_LE16(IMAGE_NT_OPTIONAL_HDR64_MAGIC);
        optional->MajorLinkerVersion = config->major_linker_version;
        optional->MinorLinkerVersion = config->minor_linker_version;
        optional->SizeOfCode = SCC_LE32(builder->sizeof_code);
        optional->SizeOfInitializedData = SCC_LE32(builder->sizeof_init_data);
        optional->SizeOfUninitializedData =
            SCC_LE32(builder->sizeof_uninit_data);
        optional->AddressOfEntryPoint =
            SCC_LE32(config->address_of_entry_point);
        optional->BaseOfCode = SCC_LE32(builder->baseof_code);
        optional->ImageBase = SCC_LE64(config->image_base);
        optional->SectionAlignment = SCC_LE32(builder->section_alignment);
        optional->FileAlignment = SCC_LE32(builder->file_alignment);
        optional->MajorOperatingSystemVersion =
            SCC_LE16(config->major_operating_system_version);
        optional->MinorOperatingSystemVersion =
            SCC_LE16(config->minor_operating_system_version);
        optional->MajorImageVersion = SCC_LE16(config->major_image_version);
        optional->MinorImageVersion = SCC_LE16(config->minor_image_version);
        optional->MajorSubsystemVersion =
            SCC_LE16(config->major_subsystem_version);
        optional->MinorSubsystemVersion =
            SCC_LE16(config->minor_subsystem_version);
        optional->Win32VersionValue = SCC_LE32(0);
        optional->SizeOfImage = SCC_LE32(builder->virtual_offset);
        optional->SizeOfHeaders = SCC_LE32(builder->sizeof_headers);
        optional->CheckSum = SCC_LE32(0);
        optional->Subsystem = SCC_LE16(config->subsystem);
        optional->DllCharacteristics = SCC_LE16(config->dll_characteristics);
        optional->SizeOfStackReserve = SCC_LE64(config->size_of_stack_reserve);
        optional->SizeOfStackCommit = SCC_LE64(config->size_of_stack_commit);
        optional->SizeOfHeapReserve = SCC_LE64(config->size_of_heap_reserve);
        optional->SizeOfHeapCommit = SCC_LE64(config->size_of_heap_commit);
        optional->LoaderFlags = SCC_LE32(0);
        optional->NumberOfRvaAndSizes =
            SCC_LE32(scc_vec_size(builder->image_data_directory_vec));
    } else {
        IMAGE_OPTIONAL_HEADER32 *optional = &nt_headers32.OptionalHeader;
        optional->Magic = SCC_LE16(IMAGE_NT_OPTIONAL_HDR32_MAGIC);
        optional->MajorLinkerVersion = config->major_linker_version;
        optional->MinorLinkerVersion = config->minor_linker_version;
        optional->SizeOfCode = SCC_LE32(builder->sizeof_code);
        optional->SizeOfInitializedData = SCC_LE32(builder->sizeof_init_data);
        optional->SizeOfUninitializedData =
            SCC_LE32(builder->sizeof_uninit_data);
        optional->AddressOfEntryPoint =
            SCC_LE32(config->address_of_entry_point);
        optional->BaseOfCode = SCC_LE32(builder->baseof_code);
        optional->BaseOfData = SCC_LE32(builder->baseof_data); // 32位特有的字段
        optional->ImageBase =
            SCC_LE32((uint32_t)(config->image_base & 0xFFFFFFFF));
        optional->SectionAlignment = SCC_LE32(builder->section_alignment);
        optional->FileAlignment = SCC_LE32(builder->file_alignment);
        optional->MajorOperatingSystemVersion =
            SCC_LE16(config->major_operating_system_version);
        optional->MinorOperatingSystemVersion =
            SCC_LE16(config->minor_operating_system_version);
        optional->MajorImageVersion = SCC_LE16(config->major_image_version);
        optional->MinorImageVersion = SCC_LE16(config->minor_image_version);
        optional->MajorSubsystemVersion =
            SCC_LE16(config->major_subsystem_version);
        optional->MinorSubsystemVersion =
            SCC_LE16(config->minor_subsystem_version);
        optional->Win32VersionValue = SCC_LE32(0);
        optional->SizeOfImage = SCC_LE32(builder->virtual_offset);
        optional->SizeOfHeaders = SCC_LE32(builder->sizeof_headers);
        optional->CheckSum = SCC_LE32(0);
        optional->Subsystem = SCC_LE16(config->subsystem);
        optional->DllCharacteristics = SCC_LE16(config->dll_characteristics);
        optional->SizeOfStackReserve =
            SCC_LE32((uint32_t)(config->size_of_stack_reserve & 0xFFFFFFFF));
        optional->SizeOfStackCommit =
            SCC_LE32((uint32_t)(config->size_of_stack_commit & 0xFFFFFFFF));
        optional->SizeOfHeapReserve =
            SCC_LE32((uint32_t)(config->size_of_heap_reserve & 0xFFFFFFFF));
        optional->SizeOfHeapCommit =
            SCC_LE32((uint32_t)(config->size_of_heap_commit & 0xFFFFFFFF));
        optional->LoaderFlags = SCC_LE32(0);
        optional->NumberOfRvaAndSizes =
            SCC_LE32(scc_vec_size(builder->image_data_directory_vec));
    }

    if (builder->is_64) {
        for (usize i = 0; i < sizeof(nt_headers64); i++) {
            scc_vec_push(builder->buffer, ((u8 *)(&nt_headers64))[i]);
        }
    } else {
        for (usize i = 0; i < sizeof(nt_headers32); i++) {
            scc_vec_push(builder->buffer, ((u8 *)(&nt_headers32))[i]);
        }
    }

    scc_vec_foreach(builder->image_data_directory_vec, i) {
        IMAGE_DATA_DIRECTORY *data =
            &scc_vec_at(builder->image_data_directory_vec, i);
        for (usize i = 0; i < sizeof(IMAGE_DATA_DIRECTORY); i++) {
            scc_vec_push(builder->buffer, ((u8 *)(data))[i]);
        }
    }
}

void scc_pe_write_section_headers(scc_pe_builder_t *builder) {
    Assert(builder->section_header_num ==
           builder->image_section_header_vec.size);
    scc_vec_foreach(builder->image_section_header_vec, i) {
        IMAGE_SECTION_HEADER *section_header =
            &scc_vec_at(builder->image_section_header_vec, i);
        for (usize j = 0; j < sizeof(IMAGE_SECTION_HEADER); j += 1) {
            scc_vec_push(builder->buffer, ((u8 *)section_header)[j]);
        }
    }
}

void scc_pe_write_header(scc_pe_builder_t *builder, scc_pe_config_t *config) {
    scc_vec_init(builder->buffer);
    scc_pe_write_dos_header_and_stub(builder);
    padding_util(builder, builder->nt_headers_offset);
    scc_pe_write_nt_header_with_config(builder, config);
    scc_pe_write_section_headers(builder);
}

scc_pe_section_range scc_pe_reserve_section_header(scc_pe_builder_t *builder,
                                                   BYTE name[8],
                                                   u32 characteristics,
                                                   u32 virtual_size,
                                                   u32 data_size) {
    scc_pe_section_range range = {0};
    range.virual_address =
        reserve_virtual(builder, virtual_size, builder->section_alignment);
    range.virual_size = virtual_size;

    range.file_size = reserve_align(data_size, builder->file_alignment);
    range.file_offset =
        range.file_size != 0
            ? reserve_file(builder, range.file_size, builder->file_alignment)
            : 0;

    u32 aligned_virtual_size =
        reserve_align(virtual_size, builder->file_alignment);

    if (characteristics & IMAGE_SCN_CNT_CODE) {
        if (!builder->baseof_code) {
            builder->baseof_code = range.virual_address;
        }
        builder->sizeof_code += aligned_virtual_size;
    } else if (characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA) {
        if (!builder->baseof_data) {
            builder->baseof_data = range.virual_address;
        }
        builder->sizeof_init_data += aligned_virtual_size;
    } else if (characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA) {
        if (!builder->baseof_data) {
            builder->baseof_data = range.virual_address;
        }
        builder->sizeof_uninit_data += aligned_virtual_size;
    }

    IMAGE_SECTION_HEADER section_header = (IMAGE_SECTION_HEADER){
        .Name[0] = name[0],
        .Name[1] = name[1],
        .Name[2] = name[2],
        .Name[3] = name[3],
        .Name[4] = name[4],
        .Name[5] = name[5],
        .Name[6] = name[6],
        .Name[7] = name[7],
        .Misc.VirtualSize = SCC_LE32(range.virual_size),
        .VirtualAddress = SCC_LE32(range.virual_address),
        .SizeOfRawData = SCC_LE32(range.file_size),
        .PointerToRawData = SCC_LE32(range.file_offset),
        .PointerToRelocations = SCC_LE32(0),
        .PointerToLinenumbers = SCC_LE32(0),
        .NumberOfRelocations = SCC_LE16(0),
        .NumberOfLinenumbers = SCC_LE16(0),
        .Characteristics = SCC_LE32(characteristics),
    };
    scc_vec_push(builder->image_section_header_vec, section_header);
    return range;
}

void scc_pe_write_section(scc_pe_builder_t *builder,
                          scc_pe_section_range *range, u8 *data,
                          usize data_size) {
    if (range == null || data == null || data_size == 0) {
        return;
    }
    padding_util(builder, range->file_offset);
    for (usize i = 0; i < data_size; i += 1) {
        scc_vec_push(builder->buffer, data[i]);
    }
    usize size = scc_vec_size(builder->buffer);
    size = reserve_align(size, builder->file_alignment);
    padding_util(builder, size);
}

void scc_pe_dump_to_file(scc_pe_builder_t *builder, const char *file_path) {
    scc_file_t fp = scc_fopen(file_path, SCC_FILE_WRITE);
    scc_fwrite(fp, builder->buffer.data, builder->buffer.size);
    scc_fclose(fp);
}