#ifndef __SCC_WINPE_H__
#define __SCC_WINPE_H__

#include "scc_pe_def.h"

typedef SCC_VEC(IMAGE_DATA_DIRECTORY) scc_winpe_image_data_directory_vec_t;
typedef SCC_VEC(IMAGE_SECTION_HEADER) scc_winpe_image_section_header_vec_t;

/**
 * @brief
 *
 */
typedef struct {
    scc_pe_buffer_t buffer; ///< 镜像文件(File Data)数据

    u32 section_alignment; ///< 文件镜像(Virtual Address)各个段的对齐方式
    u32 file_alignment;    ///< PE文件内(File Adderss)各个段对齐方式

    u32 baseof_code;        ///< 代码段基地址(Virtual Address)
    u32 sizeof_code;        ///< 代码段大小
    u32 baseof_data;        ///< 数据段基地址(Virtual Address)
    u32 sizeof_init_data;   ///< 初始化数据段大小
    u32 sizeof_uninit_data; ///< 未初始化数据段大小

    u32 file_offset;       ///< 内部使用变量,用于追踪文件内偏移地址
    u32 virtual_offset;    ///< 内部使用变量,用于追踪镜像内偏移地址
    u32 nt_headers_offset; ///< 内部使用变量,用于追踪NT头偏移地址
    u32 sizeof_headers;    ///< 内部使用变量,用于追踪完整镜像头大小

    scc_winpe_image_data_directory_vec_t image_data_directory_vec;
    scc_winpe_image_section_header_vec_t image_section_header_vec;
    u16 section_header_num;

    u32 symbol_offset;
    u32 symbol_num;

    cbool is_64;
} scc_pe_builder_t;

typedef struct {
    u16 machine;
    u32 time_date_stamp;
    u16 characteristics;
    u8 major_linker_version;
    u8 minor_linker_version;
    u32 address_of_entry_point;
    u64 image_base;
    u16 major_operating_system_version;
    u16 minor_operating_system_version;
    u16 major_image_version;
    u16 minor_image_version;
    u16 major_subsystem_version;
    u16 minor_subsystem_version;
    u16 subsystem;
    u16 dll_characteristics;
    u64 size_of_stack_reserve;
    u64 size_of_stack_commit;
    u64 size_of_heap_reserve;
    u64 size_of_heap_commit;
} scc_pe_config_t;

void scc_pe_builder_init(scc_pe_builder_t *builder, bool is_64,
                         u32 section_alignment, u32 file_alignment);

void scc_pe_reserve_header(scc_pe_builder_t *builder, u32 num_of_section);
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);

static inline scc_pe_section_range
scc_pe_reserve_text_section_header(scc_pe_builder_t *builder, u32 data_size) {
    return scc_pe_reserve_section_header(
        builder, (BYTE *)".text\0\0",
        IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
        data_size, data_size);
}

static inline scc_pe_section_range
scc_pe_reserve_data_section_header(scc_pe_builder_t *builder, u32 data_size) {
    return scc_pe_reserve_section_header(builder, (BYTE *)".data\0\0",
                                         IMAGE_SCN_CNT_INITIALIZED_DATA |
                                             IMAGE_SCN_MEM_READ |
                                             IMAGE_SCN_MEM_WRITE,
                                         data_size, data_size);
}

static inline scc_pe_section_range
scc_pe_reserve_rdata_section_header(scc_pe_builder_t *builder, u32 data_size) {
    return scc_pe_reserve_section_header(builder, (BYTE *)".rdata\0",
                                         IMAGE_SCN_CNT_INITIALIZED_DATA |
                                             IMAGE_SCN_MEM_READ,
                                         data_size, data_size);
}

static inline scc_pe_section_range
scc_pe_reserve_bss_section_header(scc_pe_builder_t *builder, u32 data_size) {
    return scc_pe_reserve_section_header(builder, (BYTE *)".bss\0\0\0",
                                         IMAGE_SCN_CNT_UNINITIALIZED_DATA |
                                             IMAGE_SCN_MEM_READ |
                                             IMAGE_SCN_MEM_WRITE,
                                         data_size, 0);
}

static inline scc_pe_section_range
scc_pe_reserve_idata_section_header(scc_pe_builder_t *builder, u32 data_size) {
    scc_pe_section_range range = scc_pe_reserve_section_header(
        builder, (BYTE *)".idata\0",
        IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ |
            IMAGE_SCN_MEM_WRITE,
        data_size, data_size);
    scc_vec_at(builder->image_data_directory_vec,
               IMAGE_DIRECTORY_ENTRY_IMPORT) = (IMAGE_DATA_DIRECTORY){
        .VirtualAddress = range.virual_address, .Size = range.virual_size};
    return range;
}

static inline scc_pe_section_range
scc_pe_reserve_reloc_section_header(scc_pe_builder_t *builder, u32 data_size) {
    scc_pe_section_range range = scc_pe_reserve_section_header(
        builder, (BYTE *)".reloc\0\0",
        IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ |
            IMAGE_SCN_MEM_DISCARDABLE,
        data_size, data_size);
    scc_vec_at(builder->image_data_directory_vec,
               IMAGE_DIRECTORY_ENTRY_BASERELOC) = (IMAGE_DATA_DIRECTORY){
        .VirtualAddress = range.virual_address, .Size = range.virual_size};
    return range;
}

void scc_pe_write_header(scc_pe_builder_t *builder, scc_pe_config_t *config);

void scc_pe_write_section(scc_pe_builder_t *builder,
                          scc_pe_section_range *range, u8 *data,
                          usize data_size);

void scc_pe_dump_to_file(scc_pe_builder_t *builder, const char *file_path);

#endif /* __SCC_WINPE_H__ */