#include "sandbox.h"
#include "sandbox_callbacks.h"
#include "sandbox_api_winhttp.h"
#include <tlhelp32.h>
#include "sandbox_api_com.h"

// COM 初始化状态跟踪
static bool g_comInitialized = false;

// 计划任务COM组件具体实现
class TaskServiceImpl : public TaskServiceSimulator {
   private:
    ULONG m_refCount = 1;

   public:
    HRESULT QueryInterface(REFIID riid, void** ppv) override {
        // 这里只模拟基本的ITaskService接口
        *ppv = this;
        AddRef();
        return S_OK;
    }

    ULONG AddRef() override { return ++m_refCount; }

    ULONG Release() override {
        ULONG ref = --m_refCount;
        if (ref == 0) {
            delete this;
        }
        return ref;
    }

    HRESULT Connect(VARIANT ServerName, VARIANT User, VARIANT Domain,
                    VARIANT Password) override {
        // 模拟连接成功
        return S_OK;
    }

    HRESULT GetFolder(BSTR path, ITaskFolder** ppFolder) override {
        // 模拟获取文件夹成功
        *ppFolder = nullptr;  // 实际使用时需要创建ITaskFolder实现
        return S_OK;
    }

    HRESULT NewTask(DWORD flags, ITaskDefinition** ppDefinition) override {
        // 模拟创建新任务成功
        *ppDefinition = nullptr;  // 实际使用时需要创建ITaskDefinition实现
        return S_OK;
    }
};

// COM 组件工厂实现
bool ComObjectFactory::IsTaskSchedulerCLSID(const CLSID& clsid) {
    // 检查是否是Task Scheduler 2.0 的CLSID
    static const CLSID CLSID_TaskScheduler = {
        0x0f87369f,
        0xa4e5,
        0x4cfc,
        {0xbd, 0x3e, 0x73, 0xe6, 0x15, 0x45, 0x72, 0xdd}};

    return IsEqualCLSID(clsid, CLSID_TaskScheduler);
}

ComObjectSimulator* ComObjectFactory::CreateInstance(const CLSID& clsid) {
    if (IsTaskSchedulerCLSID(clsid)) {
        return new TaskServiceImpl();
    }
    return nullptr;
}

// COM API 实现
void Api_CoInitializeEx(void* sandbox, uc_engine* uc, uint64_t address) {
    Sandbox* sb = static_cast<Sandbox*>(sandbox);

    // 获取参数
    uint32_t pvReserved = 0;
    uint32_t dwCoInit = 0;
    uc_mem_read(uc, address + 4, &pvReserved, sizeof(pvReserved));
    uc_mem_read(uc, address + 8, &dwCoInit, sizeof(dwCoInit));

    // 设置COM初始化状态
    g_comInitialized = true;

    // 返回成功
    uc_reg_write(uc, UC_X86_REG_EAX, &(uint32_t){S_OK});
}

void Api_CoCreateInstance(void* sandbox, uc_engine* uc, uint64_t address) {
    Sandbox* sb = static_cast<Sandbox*>(sandbox);

    if (!g_comInitialized) {
        uint32_t result = CO_E_NOTINITIALIZED;
        uc_reg_write(uc, UC_X86_REG_EAX, &result);
        return;
    }

    // 获取参数
    CLSID rclsid;
    uint32_t pUnkOuter = 0;
    uint32_t dwClsContext = 0;
    IID riid;
    uint32_t ppv = 0;

    uc_mem_read(uc, address + 4, &rclsid, sizeof(rclsid));
    uc_mem_read(uc, address + 20, &pUnkOuter, sizeof(pUnkOuter));
    uc_mem_read(uc, address + 24, &dwClsContext, sizeof(dwClsContext));
    uc_mem_read(uc, address + 28, &riid, sizeof(riid));
    uc_mem_read(uc, address + 44, &ppv, sizeof(ppv));

    // 创建COM对象
    ComObjectSimulator* obj = ComObjectFactory::CreateInstance(rclsid);
    if (obj == nullptr) {
        uint32_t result = CLASS_E_CLASSNOTAVAILABLE;
        uc_reg_write(uc, UC_X86_REG_EAX, &result);
        return;
    }

    // 写入对象指针
    uint32_t objPtr = reinterpret_cast<uint32_t>(obj);
    uc_mem_write(uc, ppv, &objPtr, sizeof(objPtr));

    uint32_t result = S_OK;
    uc_reg_write(uc, UC_X86_REG_EAX, &result);
}

void Api_VariantInit(void* sandbox, uc_engine* uc, uint64_t address) {
    Sandbox* sb = static_cast<Sandbox*>(sandbox);

    // 获取VARIANT指针
    uint32_t pvarg = 0;
    uc_mem_read(uc, address + 4, &pvarg, sizeof(pvarg));

    // 初始化VARIANT结构体
    VARIANT v;
    VariantInit(&v);

    // 写回初始化后的VARIANT
    uc_mem_write(uc, pvarg, &v, sizeof(VARIANT));

    uint32_t result = S_OK;
    uc_reg_write(uc, UC_X86_REG_EAX, &result);
}

void Api_VariantClear(void* sandbox, uc_engine* uc, uint64_t address) {
    Sandbox* sb = static_cast<Sandbox*>(sandbox);

    // 获取VARIANT指针
    uint32_t pvarg = 0;
    uc_mem_read(uc, address + 4, &pvarg, sizeof(pvarg));

    // 读取VARIANT结构体
    VARIANT v;
    uc_mem_read(uc, pvarg, &v, sizeof(VARIANT));

    // 清理VARIANT
    VariantClear(&v);

    // 写回清理后的VARIANT
    uc_mem_write(uc, pvarg, &v, sizeof(VARIANT));

    uint32_t result = S_OK;
    uc_reg_write(uc, UC_X86_REG_EAX, &result);
}

void Api_SysAllocString(void* sandbox, uc_engine* uc, uint64_t address) {
    Sandbox* sb = static_cast<Sandbox*>(sandbox);

    // 获取字符串指针
    uint32_t psz = 0;
    uc_mem_read(uc, address + 4, &psz, sizeof(psz));

    // 读取字符串
    wchar_t buffer[MAX_PATH];
    uc_mem_read(uc, psz, buffer, sizeof(buffer));

    // 分配BSTR
    BSTR bstr = SysAllocString(buffer);

    // 返回BSTR指针
    uint32_t result = reinterpret_cast<uint32_t>(bstr);
    uc_reg_write(uc, UC_X86_REG_EAX, &result);
}