#pragma once
#include <functional>
#include <map>

#include "head.h"
#define PAGE_SIZE 0x1000
#define CF_MASK (1 << 0)
#define PF_MASK (1 << 2)
#define ZF_MASK (1 << 6)
#define SF_MASK (1 << 7)
#define OF_MASK (1 << 11)
#define ALL_MASK (OF_MASK | SF_MASK | ZF_MASK | PF_MASK | CF_MASK)
#define STACK_BASE_64 0x7ffffffde000
#define STACK_BASE_32 0xfffdd000
#define STACK_SIZE_64 0x40000
#define STACK_SIZE_32 0x21000
#define HEAP_ADDRESS_64 0x500000000
#define HEAP_SIZE_64 0x5000000
#define HEAP_ADDRESS_32 0x5000000
#define HEAP_SIZE_32 0x5000000

#define PEB_BASE 0x90000
#define TEB_BASE 0x90000
#define CMDLINE_ADDRESS 0x100000   // 命令行字符串的固定地址
#define CMDLINEW_ADDRESS 0x110000  // 宽字符命令行字符串的固定地址

#define X86_GDT_ADDR 0x30000
#define X86_GDT_LIMIT 0x1000
#define X86_GDT_ENTRY_SIZE 0x8

#define API_FUNCTION_SIZE 8
#define PAGE_ALIGN(Va) (ULONG_PTR)(Va) & ~(PAGE_SIZE - 1)
#define PAGE_ALIGN_64(Va) (Va) & ~(0x1000ull - 1)
#define PAGE_ALIGN_64k(Va) ((Va)) & ~(0x10000ull - 1)
#define AlignSize(Size, Align) (Size + Align - 1) / Align* Align

enum class WinVer {
    kWin7 = 0x0610,
    kWin7SP1 = 0x0611,
    kWin8 = 0x0620,
    kWin81 = 0x0630,
    kWin10 = 0x0A00,
    kWin10RS1 = 0x0A01,   // Anniversary update
    kWin10RS2 = 0x0A02,   // Creators update
    kWin10RS3 = 0x0A03,   // Fall creators update
    kWin10RS4 = 0x0A04,   // Spring creators update
    kWin10RS5 = 0x0A05,   // October 2018 update
    kWin1019H1 = 0x0A06,  // May 2019 update 19H1
    kWin1019H2 = 0x0A07,  // November 2019 update 19H2
    kWin1020H1 = 0x0A08   // April 2020 update 20H1
};
struct _fakeApi {
    std::function<void(void*, uc_engine*, uint64_t)> func;
    uint32_t paramCount;
};

// 添加堆管理相关的结构定义
struct HeapBlock {
    uint64_t address;  // 块的起始地址
    size_t size;       // 块的大小
    bool is_free;      // 是否是空闲块
    HeapBlock* next;   // 下一个块
    HeapBlock* prev;   // 上一个块
};

struct HeapSegment {
    uint64_t base;      // 堆段的基址
    size_t size;        // 堆段的总大小
    HeapBlock* blocks;  // 块链表
};

class Sandbox {
    friend class cFixImprot;  // 声明cFixImprot为友元类
   public:
    Sandbox();
    ~Sandbox();

    // Public methods
    auto InitEnv(std::shared_ptr<BasicPeInfo> peInfo) -> void;
    auto Run() -> void;
    auto GetCapstoneHandle() const -> csh { return m_csHandle; }
    auto GetUnicornHandle() const -> uc_engine* { return m_ucEngine; }
    auto GetPeInfo() const -> std::shared_ptr<BasicPeInfo> { return m_peInfo; }
    auto GetModuleList() const -> std::vector<std::shared_ptr<struct_moudle>> {
        return m_moduleList;
    }
    auto EmulateApi(uc_engine* uc, uint64_t address, uint64_t rip,
                    std::string ApiName) -> void;
    auto GetPeb32() -> X32PEB* { return &m_peb32; }
    auto GetPeb64() -> X64PEB* { return &m_peb64; }
    auto GetTeb32() -> X32TEB* { return &m_teb32; }
    auto GetTeb64() -> X64TEB* { return &m_teb64; }
    auto GetCommandLine() const -> const char* { return m_commandLine.c_str(); }
    auto GetCommandLineAddress() const -> uint64_t { return CMDLINE_ADDRESS; }
    auto GetCommandLineWAddress() const -> uint64_t { return CMDLINEW_ADDRESS; }
    auto InitCommandLine() -> void;

    // 堆管理相关的公共方法
    auto CreateHeapSegment(uint64_t base, size_t size) -> HeapSegment*;
    auto AllocateFromSegment(HeapSegment* segment, size_t size) -> uint64_t;
    auto FreeBlock(uint64_t address) -> bool;
    auto FindHeapSegment(uint64_t address) -> HeapSegment*;
    auto MergeBlocks(HeapBlock* block) -> void;
    auto SplitBlock(HeapBlock* block, size_t size) -> void;
    std::map<uint64_t, HeapSegment*> m_heapSegments;  // 堆段映射表

   private:
    std::shared_ptr<BasicPeInfo> m_peInfo;
    uint64_t m_gsBase;
    uint64_t m_pebBase;
    uint64_t m_pebEnd;
    uint64_t m_tebBase;
    uint64_t m_tebEnd;
    PVOID m_stackBuffer;  // 没有释放
    uint64_t m_stackBase;
    uint64_t m_stackSize;
    uint64_t m_stackEnd;
    uint64_t m_heapBase;
    uint64_t m_heapSize;
    uint64_t m_heapEnd;
    uint64_t m_fakeBase;
    struct_gs_base m_gsBaseStruct = {0};
    X64TEB m_teb64 = {0};
    X64PEB m_peb64 = {0};
    X32TEB m_teb32 = {0};
    X32PEB m_peb32 = {0};
    csh m_csHandle;         // Capstone handle
    uc_engine* m_ucEngine;  // Unicorn engine handle
    std::vector<std::shared_ptr<moudle_import>> m_impFuncDict;
    std::vector<std::shared_ptr<moudle_export>> m_exportFuncDict;
    std::vector<std::shared_ptr<struct_moudle>> m_moduleList;
    std::map<std::string, std::shared_ptr<_fakeApi>> api_map;
    std::string m_commandLine;  // 存储命令行字符串

    auto ResoveImport() -> void;
    auto SetupVirtualMachine() -> void;
    auto PushModuleToVM(const char* dllName, uint64_t moduleBase) -> void;
    auto processImportModule(const moudle_import* importModule) -> void;
    auto ResolveImportExports() -> void;
    auto ResolveExport(uint64_t moduleBase)
        -> std::vector<std::shared_ptr<moudle_export>>;
    auto InitApiHooks() -> void;
    auto InitCommandLine(std::string commandLine) -> void;
};