#include "camera.h"

#include <d3d9.h>
#include <mfapi.h>
#include <mfidl.h>

#include "avs/game.h"
#include "external/rapidjson/document.h"
#include "external/rapidjson/pointer.h"
#include "external/rapidjson/prettywriter.h"
#include "games/iidx/iidx.h"
#include "util/detour.h"
#include "util/fileutils.h"
#include "util/libutils.h"
#include "util/logging.h"
#include "util/sigscan.h"
#include "util/utils.h"

static std::filesystem::path dll_path;
static HMODULE iidx_module;
static uintptr_t addr_hook_a = 0;
static uintptr_t addr_hook_b = 0;
static uintptr_t addr_textures = 0;
static uintptr_t addr_device_offset = 0;

typedef void **(__fastcall *camera_hook_a_t)(PBYTE);
typedef LPDIRECT3DTEXTURE9 (*camera_hook_b_t)(void*, int);
static camera_hook_a_t camera_hook_a_orig = nullptr;
static camera_hook_b_t camera_hook_b_orig = nullptr;
auto static hook_a_init = std::once_flag {};

static LPDIRECT3DDEVICE9EX device;
static LPDIRECT3DTEXTURE9 *texture_a = nullptr;
static LPDIRECT3DTEXTURE9 *texture_b = nullptr;

struct PredefinedHook {
    std::string     pe_identifier;
    uintptr_t       hook_a;
    uintptr_t       hook_b;
    uintptr_t       hook_textures;
    uintptr_t       hook_device_offset;
};

PredefinedHook g_predefinedHooks[] =
{
    // 27 003
    { "5f713b52_6d4090", 0x005971b0, 0x005fb2c0, 0x04e49898, 0x000000d0 },
    // 27 010
    { "5f713946_7f52b0", 0x006b89e0, 0x0071c950, 0x04fd08b8, 0x000000d0 },
};

const DWORD   g_predefinedHooksLength = ARRAYSIZE(g_predefinedHooks);

namespace games::iidx {
    static IIDXLocalCamera *top_camera = nullptr;       // camera id #0
    static IIDXLocalCamera *front_camera = nullptr;     // camera id #1
    std::vector<IIDXLocalCamera*> LOCAL_CAMERA_LIST = {};
    static IDirect3DDeviceManager9 *s_pD3DManager = nullptr;

    std::filesystem::path CAMERA_CONFIG_PATH = std::filesystem::path(_wgetenv(L"APPDATA")) / L"spicetools_camera_control.json";
    bool CAMERA_READY = false;

    bool parse_cmd_params() {
        const auto remove_spaces = [](const char& c) {
            return c == ' ';
        };

        log_misc("iidx::tdjcam", "Using user-supplied addresses (-iidxtdjcamhookoffset)");

        auto s = TDJ_CAMERA_OVERRIDE.value();
        s.erase(std::remove_if(s.begin(), s.end(), remove_spaces), s.end());

        if (sscanf(
                s.c_str(),
                "%i,%i,%i,%i",
                (int*)&addr_hook_a, (int*)&addr_hook_b, (int*)&addr_textures, (int*)&addr_device_offset) == 4) {

            return true;

        } else {
            log_fatal("iidx::tdjcam", "failed to parse -iidxtdjcamhookoffset");
            return false;
        }
    }

    bool find_camera_hooks() {
        std::string dll_name = avs::game::DLL_NAME;
        dll_path = MODULE_PATH / dll_name; 

        iidx_module = libutils::try_module(dll_path);
        if (!iidx_module) {
            log_warning("iidx::tdjcam", "failed to hook game DLL");
            return FALSE;
        }

        uint32_t time_date_stamp = 0;
        uint32_t address_of_entry_point = 0;

        // there are three different ways we try to hook the camera entry points
        // 1) user-provided override via cmd line
        // 2) predefined offsets using PE header matching (needed for iidx27 and few others)
        // 3) signature match (should work for most iidx28-31)

        // method 1: user provided
        if (TDJ_CAMERA_OVERRIDE.has_value() && parse_cmd_params()) {
            return TRUE;
        }

        // method 2: predefined offsets
        get_pe_identifier(dll_path, &time_date_stamp, &address_of_entry_point);
        auto pe = fmt::format("{:x}_{:x}", time_date_stamp, address_of_entry_point);
        log_info("iidx::tdjcam", "Locating predefined hook addresses for LDJ-{}", pe);

        for (DWORD i = 0; i < g_predefinedHooksLength; i++) {
            if (pe.compare(g_predefinedHooks[i].pe_identifier) == 0) {
                log_misc("iidx::tdjcam", "Found predefined addresses");
                addr_hook_a        = g_predefinedHooks[i].hook_a;
                addr_hook_b        = g_predefinedHooks[i].hook_b;
                addr_textures      = g_predefinedHooks[i].hook_textures;
                addr_device_offset = g_predefinedHooks[i].hook_device_offset;
                return TRUE;
            }
        }

        // method 3: signature match
        log_info("iidx::tdjcam", "Did not find predefined hook address, try signature match");

        // --- addr_hook_a ---
        uint8_t *addr_hook_a_ptr = reinterpret_cast<uint8_t *>(find_pattern(
            iidx_module,
            "488BC4565741564883EC5048C740D8FEFFFFFF",
            "XXXXXXXXXXXXXXXXXXX",
            0, 0));

        if (addr_hook_a_ptr == nullptr) {
            log_warning("iidx::tdjcam", "failed to find hook: addr_hook_a");
            log_warning("iidx::tdjcam", "hint: this feature REQUIRES a compatible DLL!");
            return FALSE;
        }

        addr_hook_a = (uint64_t) (addr_hook_a_ptr - (uint8_t*) iidx_module);

        // addr_hook_b
        uint8_t* addr_hook_b_ptr = reinterpret_cast<uint8_t *>(find_pattern(
            iidx_module,
            "E8000000004885C07439E8",
            "X????XXXXXX",
            0, 0));

        if (addr_hook_b_ptr == nullptr) {
            log_warning("iidx::tdjcam", "failed to find hook: addr_hook_b");
            return FALSE;
        }

        // displace with the content of wildcard bytes
        int32_t disp_b = *((int32_t *) (addr_hook_b_ptr + 1));

        addr_hook_b = (addr_hook_b_ptr - (uint8_t*) iidx_module) + disp_b + 5;

        // --- addr_textures ---
        uint8_t* addr_textures_ptr = reinterpret_cast<uint8_t *>(find_pattern(
            iidx_module,
            "E800000000488BC8488BD34883C420",
            "X????XXXXXXXXXX",
            0, 0));
        
        if (addr_textures_ptr == nullptr) {
            log_warning("iidx::tdjcam", "failed to find hook: addr_textures (part 1)");
            return FALSE;
        }

        // displace with the content of wildcard bytes
        int32_t disp_textures = *((int32_t *) (addr_textures_ptr + 1));
        addr_textures_ptr += disp_textures + 5;

        uint32_t search_from = (addr_textures_ptr - (uint8_t*) iidx_module);

        addr_textures_ptr = reinterpret_cast<uint8_t *>(find_pattern_from(
            iidx_module,
            "488D0D",
            "XXX",
            0, 0, search_from));

        if (addr_textures_ptr == nullptr) {
            log_warning("iidx::tdjcam", "failed to find hook: addr_textures (part 2)");
            return FALSE;
        }

        // displace again with the next integer
        disp_textures = *((int32_t *) (addr_textures_ptr + 3));

        addr_textures = (addr_textures_ptr - (uint8_t*) iidx_module) + disp_textures + 7;

        // addr_device_offset
        search_from = (addr_hook_a_ptr - (uint8_t*) iidx_module);
        uint8_t *addr_device_ptr = reinterpret_cast<uint8_t *>(find_pattern_from(
                iidx_module,
                "488B89",
                "XXX",
                3, 0, search_from));
        addr_device_offset = *addr_device_ptr;

        if (addr_textures_ptr == nullptr) {
            log_warning("iidx::tdjcam", "failed to find hook: addr_textures (part 3)");
            return FALSE;
        }

        return TRUE;
    }

    static void **__fastcall camera_hook_a(PBYTE a1) {
        std::call_once(hook_a_init, [&]{           
            device = *reinterpret_cast<LPDIRECT3DDEVICE9EX*>(a1 + addr_device_offset);
            auto const textures = *reinterpret_cast<LPDIRECT3DTEXTURE9**>((uint8_t*)iidx_module + addr_textures);

            texture_a = textures;
            texture_b = textures + 2;

            init_local_camera();
        });
        return camera_hook_a_orig(a1);
    }

    static LPDIRECT3DTEXTURE9 camera_hook_b(void* a1, int idx) {
        if (idx == 0 && top_camera) {
            return top_camera->GetTexture();
        }
        if (idx == 1 && front_camera) {
            return front_camera->GetTexture();
        }
        return camera_hook_b_orig(a1, idx);
    }

    bool create_hooks() {
        auto *hook_a_ptr = reinterpret_cast<uint16_t *>(
                reinterpret_cast<intptr_t>(iidx_module) + addr_hook_a);

        if (!detour::trampoline(
                    reinterpret_cast<camera_hook_a_t>(hook_a_ptr),
                    camera_hook_a,
                    &camera_hook_a_orig)) {
            log_warning("iidx::tdjcam", "failed to trampoline hook_a");
            return FALSE;
        }

        auto *hook_b_ptr = reinterpret_cast<uint16_t *>(
                reinterpret_cast<intptr_t>(iidx_module) + addr_hook_b);

        if (!detour::trampoline(
                    reinterpret_cast<camera_hook_b_t>(hook_b_ptr),
                    camera_hook_b,
                    &camera_hook_b_orig)) {
            log_warning("iidx::tdjcam", "failed to trampoline hook_b");
            return FALSE;
        }

        return TRUE;
    }

    HRESULT create_d3d_device_manager() {
        UINT resetToken = 0;
        IDirect3DDeviceManager9 *pD3DManager = nullptr;

        HRESULT hr = DXVA2CreateDirect3DDeviceManager9(&resetToken, &pD3DManager);
        if (FAILED(hr)) { goto done; }

        hr = pD3DManager->ResetDevice(device, resetToken);
        if (FAILED(hr)) { goto done; }

        s_pD3DManager = pD3DManager;
        (s_pD3DManager)->AddRef();

    done:
        if (SUCCEEDED(hr)) {
            log_misc("iidx::tdjcam", "Created DeviceManager for DXVA");
        } else {
            log_warning("iidx::tdjcam", "Cannot create DXVA DeviceManager: {}", hr);
        }

        SafeRelease(&pD3DManager);
        return hr;
    }

    bool init_local_camera() {
        HRESULT hr = S_OK;
        IMFAttributes *pAttributes = nullptr;
        IMFActivate **ppDevices = nullptr;
        uint32_t numDevices;

        // Initialize an attribute store to specify enumeration parameters.
        hr = MFCreateAttributes(&pAttributes, 1);

        // Ask for source type = video capture devices.
        if (SUCCEEDED(hr)) {
            hr = pAttributes->SetGUID(
                MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE, 
                MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_VIDCAP_GUID
            );
        }
        
        // Enumerate devices.
        if (SUCCEEDED(hr)) {
            hr = MFEnumDeviceSources(pAttributes, &ppDevices, &numDevices);
        }

        log_info("iidx::tdjcam", "MFEnumDeviceSources returned {} device(s)", numDevices);
        if (numDevices > 0) {
            hr = create_d3d_device_manager();
        }      

        if (SUCCEEDED(hr)) {
            int cameraIndex = 0;
            for (size_t i = 0; i < numDevices && LOCAL_CAMERA_LIST.size() < 2; i++) {
                // get camera activation
                auto pActivate = ppDevices[i];

                if ((cameraIndex == 0 && !FLIP_CAMS) || (cameraIndex == 1 && FLIP_CAMS)) {
                    top_camera = new IIDXLocalCamera("top", TDJ_CAMERA_PREFER_16_9, pActivate, s_pD3DManager, device, texture_a);
                    if (top_camera->m_initialized) {
                        LOCAL_CAMERA_LIST.push_back(top_camera);
                        cameraIndex++;
                    } else {
                        top_camera->Release();
                        top_camera = nullptr;
                    }
                } else {
                    front_camera = new IIDXLocalCamera("front", TDJ_CAMERA_PREFER_16_9, pActivate, s_pD3DManager, device, texture_b);
                    if (front_camera->m_initialized) {
                        LOCAL_CAMERA_LIST.push_back(front_camera);
                        cameraIndex++;
                    } else {
                        front_camera->Release();
                        front_camera = nullptr;
                    }
                }
            }
        }

        if (LOCAL_CAMERA_LIST.size() == 0) {
            goto done;
        }

        camera_config_load();
        if (top_camera) {
            top_camera->StartCapture();
        }
        if (front_camera) {
            front_camera->StartCapture();
        }

        CAMERA_READY = true;

    done:
        SafeRelease(&pAttributes);

        for (DWORD i = 0; i < numDevices; i++) {
            SafeRelease(&ppDevices[i]);
        }
        CoTaskMemFree(ppDevices);

        return SUCCEEDED(hr);
    }

    bool init_camera_hooks() {
        bool result = find_camera_hooks();
        if (result) {
            log_misc(
                "iidx::tdjcam",
                "found hooks:\n    hook_a        0x{:x}\n    hook_b        0x{:x}\n    textures      0x{:x}\n    device_offset 0x{:x}",
                addr_hook_a, addr_hook_b, addr_textures, addr_device_offset);
            result = create_hooks();
        }
        
        return result;
    }

    void camera_release() {
        if (top_camera) {
            top_camera->Release();
            top_camera = nullptr;
        }
        if (front_camera) {
            front_camera->Release();
            front_camera = nullptr;
        }
        SafeRelease(&s_pD3DManager);
    }

    bool camera_config_load() {
        log_info("iidx::tdjcam", "loading config");

        try {
            // read config file
            std::string config = fileutils::text_read(CAMERA_CONFIG_PATH);
            if (!config.empty()) {
                // parse document
                rapidjson::Document doc;
                doc.Parse(config.c_str());

                // check parse error
                auto error = doc.GetParseError();
                if (error) {
                    log_warning("iidx::tdjcam", "config parse error: {}", error);
                    return false;
                }

                // verify root is a dict
                if (!doc.IsObject()) {
                    log_warning("iidx::tdjcam", "config not found");
                    return false;
                }

                const std::string root("/sp2x_cameras");
                auto numCameras = LOCAL_CAMERA_LIST.size();
                for (size_t i = 0; i < numCameras; i++) {
                    auto *camera = LOCAL_CAMERA_LIST.at(i);
                    auto symLink = camera->GetSymLink();

                    const auto cameraNode = rapidjson::Pointer(root + "/" + symLink).Get(doc);
                    if (cameraNode && cameraNode->IsObject()) {
                        log_info("iidx::tdjcam", "Parsing config for: {}", symLink);

                        // Media type
                        auto mediaTypePointer = rapidjson::Pointer(root + "/" + symLink + "/MediaType").Get(doc);
                        if (mediaTypePointer && mediaTypePointer->IsString()) {
                            std::string mediaType = mediaTypePointer->GetString();
                            if (mediaType.length() > 0) {
                                camera->m_selectedMediaTypeDescription = mediaType;
                                camera->m_useAutoMediaType = false;
                            } else {
                                camera->m_useAutoMediaType = true;
                            }
                        }

                        // Draw mode
                        auto drawModePointer = rapidjson::Pointer(root + "/" + symLink + "/DrawMode").Get(doc);
                        if (drawModePointer && drawModePointer->IsString()) {
                            std::string drawModeString = drawModePointer->GetString();
                            for (int j = 0; j < DRAW_MODE_SIZE; j++) {
                                if (DRAW_MODE_LABELS[j].compare(drawModeString) == 0) {
                                    camera->m_drawMode = (LocalCameraDrawMode) j;
                                    break;
                                }
                            }
                        }

                        // Allow manual control
                        auto manualControlPointer = rapidjson::Pointer(root + "/" + symLink + "/AllowManualControl").Get(doc);
                        if (manualControlPointer && manualControlPointer->IsBool() && manualControlPointer->GetBool()) {
                            camera->m_allowManualControl = true;
                        }

                        // Camera control
                        // if m_allowManualControl is false, SetCameraControlProp will fail, but run through this code
                        // anyway to exercise parsing code
                        for (int propIndex = 0; propIndex < CAMERA_CONTROL_PROP_SIZE; propIndex++) {
                            std::string label = CAMERA_CONTROL_LABELS[propIndex];
                            CameraControlProp prop = {};
                            camera->GetCameraControlProp(propIndex, &prop);
                            auto valuePointer = rapidjson::Pointer(root + "/" + symLink + "/" + label + "/value").Get(doc);
                            auto flagsPointer = rapidjson::Pointer(root + "/" + symLink + "/" + label + "/flags").Get(doc);
                            if (valuePointer && valuePointer->IsInt() && flagsPointer && flagsPointer->IsInt()) {
                                prop.value = (long)valuePointer->GetInt();
                                prop.valueFlags = (long)flagsPointer->GetInt();
                                camera->SetCameraControlProp(propIndex, prop.value, prop.valueFlags);
                            }

                        }
                        log_misc("iidx::tdjcam", "  >> done");
                    } else {
                        log_misc("iidx::tdjcam", "No previous config for: {}", symLink);
                    }
                }
            }
        } catch (const std::exception& e) {
            log_warning("iidx::tdjcam", "exception occurred while config: {}", e.what());
            return false;
        }

        return true;
    }

    bool camera_config_save() {
        log_info("iidx::tdjcam", "saving config");

        // create document
        rapidjson::Document doc;
        std::string config = fileutils::text_read(CAMERA_CONFIG_PATH);
        if (!config.empty()) {
            doc.Parse(config.c_str());
            log_misc("iidx::tdjcam", "existing config file found");
        }
        if (!doc.IsObject()) {
            log_misc("iidx::tdjcam", "clearing out config file");
            doc.SetObject();
        }

        auto numCameras = LOCAL_CAMERA_LIST.size();
        for (size_t i = 0; i < numCameras; i++) {
            auto *camera = LOCAL_CAMERA_LIST.at(i);
            auto symLink = camera->GetSymLink();
            std::string root("/sp2x_cameras/" + symLink + "/");

            // Media type
            if (camera->m_useAutoMediaType) {
                rapidjson::Pointer(root + "MediaType").Set(doc, "");
            } else {
                rapidjson::Pointer(root + "MediaType").Set(doc, camera->m_selectedMediaTypeDescription);
            }

            // Draw Mode
            rapidjson::Pointer(root + "DrawMode").Set(doc, DRAW_MODE_LABELS[camera->m_drawMode]);

            // Manual control
            rapidjson::Pointer(root + "AllowManualControl").Set(doc, camera->m_allowManualControl);

            // Camera control
            for (int propIndex = 0; propIndex < CAMERA_CONTROL_PROP_SIZE; propIndex++) {
                std::string label = CAMERA_CONTROL_LABELS[propIndex];
                CameraControlProp prop = {};
                camera->GetCameraControlProp(propIndex, &prop);
                rapidjson::Pointer(root + label + "/value").Set(doc, (int)prop.value);
                rapidjson::Pointer(root + label + "/flags").Set(doc, (int)prop.valueFlags);
            }
        }

        // build JSON
        rapidjson::StringBuffer buffer;
        rapidjson::PrettyWriter<rapidjson::StringBuffer> writer(buffer);
        doc.Accept(writer);

        // save to file
        if (fileutils::text_write(CAMERA_CONFIG_PATH, buffer.GetString())) {
        } else {
            log_warning("iidx::tdjcam", "unable to save config file to {}", CAMERA_CONFIG_PATH.string());
        }

        return true;
    }

    bool camera_config_reset() {
        return false;
    }

}