#pragma once
#include <windows.h>
#include <combaseapi.h>
#include <unicorn/unicorn.h>
#include "sandbox.h"

// 辅助函数声明
void read_from_unicorn(uc_engine* uc, uint64_t address, void* buffer,
                       size_t size);
void write_to_unicorn(uc_engine* uc, uint64_t address, const void* buffer,
                      size_t size);
void set_return_value(uc_engine* uc, uint32_t value);

// COM 组件接口定义
struct ITaskService;
struct ITaskFolder;
struct ITaskDefinition;
struct IRegisteredTask;

// COM 组件模拟器基类
class ComObjectSimulator {
   public:
    virtual ~ComObjectSimulator() = default;
    virtual HRESULT QueryInterface(REFIID riid, void** ppv) = 0;
    virtual ULONG AddRef() = 0;
    virtual ULONG Release() = 0;
};

// 计划任务COM组件模拟器
class TaskServiceSimulator : public ComObjectSimulator {
   public:
    virtual HRESULT Connect(VARIANT ServerName, VARIANT User, VARIANT Domain,
                            VARIANT Password) = 0;
    virtual HRESULT GetFolder(BSTR path, ITaskFolder** ppFolder) = 0;
    virtual HRESULT NewTask(DWORD flags, ITaskDefinition** ppDefinition) = 0;
};

// COM API 模拟函数声明
void Api_CoInitializeEx(void* sandbox, uc_engine* uc, uint64_t address);
void Api_CoCreateInstance(void* sandbox, uc_engine* uc, uint64_t address);
void Api_VariantInit(void* sandbox, uc_engine* uc, uint64_t address);
void Api_VariantClear(void* sandbox, uc_engine* uc, uint64_t address);
void Api_SysAllocString(void* sandbox, uc_engine* uc, uint64_t address);

// COM 组件工厂
class ComObjectFactory {
   public:
    static ComObjectSimulator* CreateInstance(const CLSID& clsid);

   private:
    static bool IsTaskSchedulerCLSID(const CLSID& clsid);
};