#include <scc_pe_idata.h>

// RVA Relative Virtual Address
static inline IMAGE_IMPORT_DESCRIPTOR
image_import_descriptor_init(u32 import_lookup_table_rva, u32 name_rva,
                             u32 import_address_table_rva) {
    IMAGE_IMPORT_DESCRIPTOR iid = {0};
    iid.DUMMYUNIONNAME.OriginalFirstThunk = SCC_LE32(import_lookup_table_rva);
    iid.TimeDateStamp = SCC_LE32(0);
    iid.ForwarderChain = SCC_LE32(0);
    iid.Name = SCC_LE32(name_rva);
    iid.FirstThunk = SCC_LE32(import_address_table_rva);
    return iid;
};

static void scc_winpe_hnt_builder_push(scc_pe_hnt_builder_t *builder,
                                       const char *name, u16 hint);
static void scc_winpe_hnt_builder_init(scc_pe_hnt_builder_t *builder) {
    scc_vec_init(builder->data);
    builder->section_offset = 0;
    scc_hashtable_cstr_init(&builder->str_map);
}

static void scc_winpe_hnt_builder_push(scc_pe_hnt_builder_t *builder,
                                       const char *name, u16 hint) {
    ///< 可选哈希表去重
    if (scc_hashtable_get(&builder->str_map, name)) {
        return;
    }
    scc_hashtable_set(&builder->str_map, name, (void *)(u64)builder->data.size);

    /// FIXME WARNING需要转换为little endian
    scc_vec_push(builder->data, hint & 0xff);
    scc_vec_push(builder->data, hint >> 8);

    while (*name) {
        scc_vec_push(builder->data, *name);
        name++;
    }
    scc_vec_push(builder->data, 0);

    if (scc_vec_size(builder->data) % 2 != 0) {
        scc_vec_push(builder->data, 0);
    }

    // return scc_vec_size(builder->data);
    Assert(scc_vec_size(builder->data) % 2 == 0);
}

static u64 scc_winpe_hnt_get_offset(scc_pe_hnt_builder_t *builder,
                                    const char *name) {
    return (u64)scc_hashtable_get(&builder->str_map, name);
}

void scc_pe_idata_builder_init(scc_pe_idata_builder_t *builder,
                               scc_pe_idata_lib_vec_t *idata_libs) {
    scc_vec_init(builder->buffer);
    builder->idata_libs = *idata_libs;

    scc_winpe_hnt_builder_init(&builder->hnt_builder);

    scc_hashtable_cstr_init(&builder->iat_map);
}

u32 scc_pe_reserve_idata(scc_pe_idata_builder_t *builder) {
    /* TODO overflow check: idata_size 为 u32, 乘法可能溢出,
       设计未定暂不插入 stdint 类型, 仅做简单防护 */
    u32 idata_size = (scc_vec_size(builder->idata_libs) + 1) *
                     sizeof(IMAGE_IMPORT_DESCRIPTOR);

    scc_vec_foreach(builder->idata_libs, i) {
        scc_pe_idata_lib_t *lib = &scc_vec_at(builder->idata_libs, i);
        u32 thunk_size = (scc_vec_size(lib->symbol_names) + 1) * 2 *
                         sizeof(IMAGE_THUNK_DATA64);
        if (idata_size + thunk_size < idata_size) {
            LOG_FATAL("scc_pe_reserve_idata: u32 overflow");
        }
        idata_size += thunk_size;
        scc_winpe_hnt_builder_push(&builder->hnt_builder, lib->name, 0);

        scc_vec_foreach(lib->symbol_names, j) {
            scc_winpe_hnt_builder_push(&builder->hnt_builder,
                                       scc_vec_at(lib->symbol_names, j), 0);
        }
    }
    builder->hnt_builder.section_offset = idata_size;
    if (idata_size + scc_vec_size(builder->hnt_builder.data) < idata_size) {
        LOG_FATAL("scc_pe_reserve_idata: u32 overflow on hnt data");
    }
    idata_size += scc_vec_size(builder->hnt_builder.data);

    scc_vec_realloc(builder->buffer, idata_size);
    scc_vec_size(builder->buffer) = idata_size;
    return idata_size;
}

scc_pe_buffer_t scc_pe_construct_idata(scc_pe_idata_builder_t *builder,
                                       scc_pe_section_range *idata_range) {
    u32 idata_rva = idata_range->virual_address;
    u32 hnt_offset = builder->hnt_builder.section_offset;

    usize import_file_count = scc_vec_size(builder->idata_libs);
    usize current_offset =
        (import_file_count + 1) * sizeof(IMAGE_IMPORT_DESCRIPTOR);
    scc_vec_foreach(builder->idata_libs, i) {
        scc_pe_idata_lib_t *lib = &scc_vec_at(builder->idata_libs, i);
        scc_winpe_lookup_table_vec_t lookup_table;
        scc_vec_init(lookup_table);

        scc_vec_foreach(lib->symbol_names, j) {
            u64 name_offset = scc_winpe_hnt_get_offset(
                &builder->hnt_builder, scc_vec_at(lib->symbol_names, j));
            u64 name_rva = idata_rva + hnt_offset + name_offset;

            IMAGE_THUNK_DATA64 lookup_table_entry = {.u1.AddressOfData =
                                                         name_rva};
            scc_vec_push(lookup_table, lookup_table_entry);
        }
        scc_vec_push(lookup_table, (IMAGE_THUNK_DATA64){0});
        usize table_size =
            scc_vec_size(lookup_table) * sizeof(IMAGE_THUNK_DATA64);

        usize ilt_offset = current_offset;
        usize iat_offset = ilt_offset + table_size;
        current_offset += table_size * 2;

        scc_vec_foreach(lib->symbol_names, j) {
            // 构造IAT链接
            scc_hashtable_set(
                &builder->iat_map, scc_vec_at(lib->symbol_names, j),
                (void *)(usize)(idata_rva + j * sizeof(IMAGE_THUNK_DATA64)) +
                    iat_offset);
        }

        usize name_rva = idata_rva + hnt_offset +
                         scc_winpe_hnt_get_offset(&builder->hnt_builder,
                                                  lib->name); ///< hashtable get

        // 构造并写入导入目录项
        // The address of an ASCII string that contains the name of the DLL.
        // This address is relative to the image base.
        // https://learn.microsoft.com/en-us/windows/win32/debug/pe-format#import-directory-table
        // so add 2 to skip the two bytes of the Hint
        IMAGE_IMPORT_DESCRIPTOR dir_entry = image_import_descriptor_init(
            ilt_offset + idata_rva, name_rva + 2, iat_offset + idata_rva);
        scc_memcpy(
            &scc_vec_at(builder->buffer, i * sizeof(IMAGE_IMPORT_DESCRIPTOR)),
            &dir_entry,
            sizeof(IMAGE_IMPORT_DESCRIPTOR)); // FIXME 大小端

        Assert(table_size ==
               scc_vec_size(lookup_table) * sizeof(IMAGE_THUNK_DATA64));

        // 转换为字节序列并写入
        // FIXME 大小端
        // 写入ILT表
        scc_memcpy(&scc_vec_at(builder->buffer, ilt_offset),
                   &scc_vec_at(lookup_table, 0), table_size);
        // 写入IAT表
        scc_memcpy(&scc_vec_at(builder->buffer, iat_offset),
                   &scc_vec_at(lookup_table, 0), table_size);
    }

    // 添加nullptr终止的目录项
    IMAGE_IMPORT_DESCRIPTOR nullptr_entry =
        image_import_descriptor_init(0, 0, 0);
    scc_memcpy(&scc_vec_at(builder->buffer,
                           import_file_count * sizeof(IMAGE_IMPORT_DESCRIPTOR)),
               &nullptr_entry, sizeof(IMAGE_IMPORT_DESCRIPTOR));

    // 填充Hint/Name表
    scc_memcpy(&scc_vec_at(builder->buffer, hnt_offset),
               &scc_vec_at(builder->hnt_builder.data, 0),
               scc_vec_size(builder->hnt_builder.data));

    return builder->buffer;
}