#include <scc_pe_idata.h>
#include <sccf2pe.h>
#include <sccf_utils.h>
#include <x86/scc_x86_patch.h>

typedef struct {
    scc_hashtable_t str2libsym;
    scc_pe_idata_lib_vec_t idata_libs;
    scc_strpool_t pool; /* owns all interned symbol names and library names */
    const char *find_path;
} pe_idata_lib_ctx_t;

static void load_from_def(pe_idata_lib_ctx_t *ctx, const char *file_path,
                          const char *dll_name) {
    /*
    LIBRARY <path name>
    EXPORTS
    name @number
    ...
    */
    scc_str_t fpath = scc_str_from_cstr(file_path);
    scc_str_append_ch(&fpath, '/');
    scc_str_append_cstr(&fpath, dll_name, scc_strlen(dll_name));
    scc_str_append_cstr(&fpath, ".def", 4);
    const char *fname = scc_str_as_cstr(&fpath);
    scc_file_t fp = scc_fopen(fname, SCC_FILE_READ);
    if (fp == nullptr) {
        scc_str_drop(&fpath);
        return;
    }
    LOG_TRACE("load_from_def file read sucessful: %s", fname);

    usize fsize = scc_fsize(fp);
    if (fsize == 0) {
        scc_fclose(fp);
        scc_str_drop(&fpath);
        return;
    }

    char *buffer = scc_malloc(fsize);
    Assert(buffer != nullptr);

    usize read_size = scc_fread(fp, buffer, fsize);
    Assert(read_size == fsize);
    scc_fclose(fp);

    usize line = 0;
    buffer[fsize - 1] = '\0';
    const char *dll_pooled = scc_strpool_intern(&ctx->pool, dll_name);
    for (usize i = 0; i < fsize; i += 1) {
        if (buffer[i] == '\n') {
            line += 1;
        }
        if (line < 2) {
            continue;
        }
        if (buffer[i] == ' ') {
            continue;
        }
        for (usize j = i; j < fsize; j += 1) {
            if (buffer[j] == ' ') {
                buffer[j] = '\0';
                break;
            }
        }
        const char *sym_pooled = scc_strpool_intern(&ctx->pool, buffer + i);
        scc_hashtable_set(&ctx->str2libsym, sym_pooled, (void *)dll_pooled);
    }
    scc_free(buffer);
    scc_str_drop(&fpath);
}

static void pe_idata_lib_init(pe_idata_lib_ctx_t *ctx, const char *find_path) {
    // Got .dll.def

    scc_hashtable_cstr_init(&ctx->str2libsym);
    scc_strpool_init(&ctx->pool);
    scc_vec_init(ctx->idata_libs);
    ctx->find_path = find_path;
    load_from_def(ctx, ctx->find_path, "msvcrt.dll");
}

static cbool pe_idata_get(pe_idata_lib_ctx_t *ctx, const char *name) {
    const char *lib_name = scc_hashtable_get(&ctx->str2libsym, name);
    if (lib_name == nullptr) {
        return false;
    }

    scc_pe_idata_lib_t *lib = nullptr;
    scc_vec_foreach(ctx->idata_libs, i) {
        scc_pe_idata_lib_t *idata_lib = &scc_vec_at(ctx->idata_libs, i);
        if (scc_strcmp(lib_name, idata_lib->name) == 0) {
            lib = idata_lib;
            break;
        }
    }
    if (lib == nullptr) {
        scc_pe_idata_lib_t new_lib;
        new_lib.name = lib_name;
        scc_vec_init(new_lib.symbol_names);
        scc_vec_push(ctx->idata_libs, new_lib);
        lib = &scc_vec_at(ctx->idata_libs, scc_vec_size(ctx->idata_libs) - 1);
    }

    scc_vec_push(lib->symbol_names, name);
    return true;
}

void sccf2pe(scc_pe_builder_t *builder, const sccf_t *sccf) {
    scc_pe_builder_init(builder, true, 4096, 512);
    sccf_strtab_t strtab;
    scc_vec_init(strtab);
    sccf_reloc_vec_t relocs;
    scc_vec_init(relocs);
    sccf_sym_vec_t symtab;
    scc_vec_init(symtab);

    sccf_sect_data_t *code_data = nullptr;
    sccf_sect_data_t *data_data = nullptr;
    int num_of_section = 1;
    scc_vec_foreach(sccf->sect_headers, i) {
        sccf_sect_header_t *sect_header = &scc_vec_at(sccf->sect_headers, i);
        sccf_sect_data_t *sect_data = &scc_vec_at(sccf->sect_datas, i);
        if (sect_header->sccf_sect_type == SCCF_SECT_CODE) {
            if (scc_vec_size(*sect_data) != 0) {
                code_data = sect_data;
                num_of_section += 1;
            }
        } else if (sect_header->sccf_sect_type == SCCF_SECT_DATA) {
            if (scc_vec_size(*sect_data) != 0) {
                data_data = sect_data;
                num_of_section += 1;
            }
        } else if (sect_header->sccf_sect_type == SCCF_SECT_STRTAB) {
            scc_vec_unsafe_from_buffer(
                strtab, (char *)scc_vec_unsafe_get_data(*sect_data),
                scc_vec_size(*sect_data));
        } else if (sect_header->sccf_sect_type == SCCF_SECT_RELOCS) {
            scc_vec_unsafe_from_buffer(
                relocs, (sccf_reloc_t *)scc_vec_unsafe_get_data(*sect_data),
                scc_vec_size(*sect_data));
        } else if (sect_header->sccf_sect_type == SCCF_SECT_SYMTAB) {
            scc_vec_unsafe_from_buffer(
                symtab, (sccf_sym_t *)scc_vec_unsafe_get_data(*sect_data),
                scc_vec_size(*sect_data));
        }
    }

    scc_pe_reserve_header(builder, num_of_section);

    scc_pe_section_range code_range = {0};
    scc_pe_section_range data_range = {0};
    scc_pe_section_range idata_range = {0};

    if (code_data) {
        code_range = scc_pe_reserve_text_section_header(
            builder, scc_vec_size(*code_data));
    }
    if (data_data) {
        data_range = scc_pe_reserve_data_section_header(
            builder, scc_vec_size(*data_data));
    }

    pe_idata_lib_ctx_t idata_lib_ctx;
    pe_idata_lib_init(&idata_lib_ctx, "./.scc_dll");
    scc_vec_foreach(symtab, i) {
        sccf_sym_t *sym = &scc_vec_at(symtab, i);
        if (sym->sccf_sym_type == SCCF_SYM_TYPE_EXTERN) {
            const char *name =
                (const char *)&scc_vec_at(strtab, sym->name_offset);
            if (pe_idata_get(&idata_lib_ctx, name) == false) {
                LOG_ERROR("link error: symbol [%s] not found", name);
            }
        }
    }
    scc_pe_idata_builder_t idata_builder;
    scc_pe_idata_builder_init(&idata_builder, &idata_lib_ctx.idata_libs);
    u32 idata_size = scc_pe_reserve_idata(&idata_builder);
    idata_range = scc_pe_reserve_idata_section_header(builder, idata_size);
    scc_pe_buffer_t idata_buffer =
        scc_pe_construct_idata(&idata_builder, &idata_range);

    u32 entry_point_offset = sccf->header.entry_point;
    Assert(code_data != nullptr &&
           entry_point_offset < scc_vec_size(*code_data));
    u64 base_address = 0x140000000;
    u32 entry_point = code_range.virual_address + entry_point_offset;
    scc_pe_config_t config = (scc_pe_config_t){
        .machine = IMAGE_FILE_MACHINE_AMD64,
        .time_date_stamp = 0,
        .characteristics =
            IMAGE_FILE_EXECUTABLE_IMAGE | IMAGE_FILE_LARGE_ADDRESS_AWARE,
        .major_linker_version = 14,
        .minor_linker_version = 0,
        .address_of_entry_point = entry_point,
        .image_base = base_address,
        .major_operating_system_version = 6,
        .minor_operating_system_version = 0,
        .major_image_version = 0,
        .minor_image_version = 0,
        .major_subsystem_version = 6,
        .minor_subsystem_version = 0,
        .subsystem = IMAGE_SUBSYSTEM_WINDOWS_CUI,
        .dll_characteristics = IMAGE_DLLCHARACTERISTICS_HIGH_ENTROPY_VA |
                               IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE |
                               IMAGE_DLLCHARACTERISTICS_NX_COMPAT,
        .size_of_stack_reserve = 0x100000,
        .size_of_stack_commit = 0x1000,
        .size_of_heap_reserve = 0x100000,
        .size_of_heap_commit = 0x1000,
    };

    scc_mcode_t mcode;
    scc_mcode_init(&mcode, SCC_MCODE_ARCH_X86_64);
    // FIXME hack
    mcode.code = *(scc_mcode_buff_t *)code_data;
    scc_vec_foreach(relocs, i) {
        sccf_reloc_t *reloc = &scc_vec_at(relocs, i);
        if (reloc->reloc_type == SCCF_RELOC_TYPE_EMPTY) {
            continue;
        }
        Assert(reloc->sect_type == SCCF_SECT_CODE);
        sccf_sym_t *sym = &scc_vec_at(symtab, reloc->sym_idx);
        const char *name = &scc_vec_at(strtab, sym->name_offset);

        u32 rva = 0;
        if (sym->sccf_sym_type == SCCF_SYM_TYPE_EXTERN) {
            rva = scc_pe_idata_get_symbol_rva(&idata_builder, name);
        } else if (sym->sccf_sect_type == SCCF_SECT_DATA) {
            rva = data_range.virual_address + sym->sccf_sect_offset;
        } else if (sym->sccf_sect_type == SCCF_SECT_CODE) {
            rva = code_range.virual_address + sym->sccf_sect_offset;
        } else {
            Panic("unsupported reloc symbol type");
        }

        Assert(rva != 0);
        Assert(code_data != nullptr);
        // FIXME patch type
        // if (reloc->reloc_type == SCCF_RELOC_TYPE_ABS) {
        //     TODO();
        // }
        scc_x86_patch(&mcode, SCC_X86_PATCH_PC32, reloc->offset + reloc->addend,
                      (i64)rva - (i64)code_range.virual_address);
    }
    *(scc_mcode_buff_t *)code_data = mcode.code;

    scc_pe_write_header(builder, &config);
    if (code_data != nullptr) {
        scc_pe_write_section(builder, &code_range,
                             (u8 *)scc_vec_unsafe_get_data(*code_data),
                             scc_vec_size(*code_data));
    }
    if (data_data != nullptr) {
        scc_pe_write_section(builder, &data_range,
                             (u8 *)scc_vec_unsafe_get_data(*data_data),
                             scc_vec_size(*data_data));
    }
    scc_pe_write_section(builder, &idata_range,
                         scc_vec_unsafe_get_data(idata_buffer),
                         scc_vec_size(idata_buffer));

    // Cleanup
    scc_vec_foreach(idata_lib_ctx.idata_libs, i) {
        scc_vec_free(scc_vec_at(idata_lib_ctx.idata_libs, i).symbol_names);
    }
    scc_vec_free(idata_lib_ctx.idata_libs);
    scc_hashtable_drop(&idata_lib_ctx.str2libsym);
    scc_strpool_drop(&idata_lib_ctx.pool);
}