/*
 * Copyright (c) 2026 Mykhailo Mamedov. All rights reserved.
 *
 * RESEARCH PREVIEW / REFERENCE ONLY:
 * This source code is provided solely for the purpose of reviewing
 * the author's research methods and implementation.
 *
 * NO LICENSE GRANTED:
 * This code is NOT for distribution, modification, or use in any
 * project (commercial or otherwise). Unauthorized copying or
 * use of this code is strictly prohibited.
 *
 * For inquiries regarding use or licensing, contact: ua.modin@gmail.com
 * 
 * Minimal test.
 * 
 */

#include "pch.h"
#include "common.h"
#include "pipelineTracker.hpp"

using namespace M_CORE;

// Fake hardware state registries
static std::unordered_map<VkSemaphore, uint64_t> g_mockSems;
static uint64_t g_nextMockHandle = 0x1000;

// --- VULKAN OVERRIDES VIA HOOKING
extern "C" {
    VkResult vkGetSemaphoreCounterValue(VkDevice device, VkSemaphore semaphore, uint64_t* pValue) {
        if (pValue && g_mockSems.count(semaphore)) {
            *pValue = g_mockSems[semaphore];
        }
        else if (pValue) {
            *pValue = 0; // Default fallback
        }
        return VK_SUCCESS;
    }

    VkResult vkSignalSemaphore(VkDevice device, const VkSemaphoreSignalInfo* pSignalInfo) {
        if (pSignalInfo && pSignalInfo->semaphore) {
            g_mockSems[pSignalInfo->semaphore] = pSignalInfo->value;
        }
        return VK_SUCCESS;
    }

    VkResult vkCreateSemaphore(VkDevice device, const VkSemaphoreCreateInfo* pCreateInfo, 
        const VkAllocationCallbacks* pAllocator, VkSemaphore* pSemaphore) {
            if (!pSemaphore) return VK_ERROR_INITIALIZATION_FAILED;

            // Generate a fake handle.
            g_nextMockHandle += 8;
            VkSemaphore fakeHandle = reinterpret_cast<VkSemaphore>(g_nextMockHandle);

            // Parse the initial value.
            uint64_t initialValue = 0;
            if (pCreateInfo && pCreateInfo->pNext) {
                auto* typeInfo = reinterpret_cast<const VkSemaphoreTypeCreateInfo*>(pCreateInfo->pNext);
                if (typeInfo->sType == VK_STRUCTURE_TYPE_SEMAPHORE_TYPE_CREATE_INFO) {
                    initialValue = typeInfo->initialValue;
                }
            }

            // Register into map.
            g_mockSems[fakeHandle] = initialValue;

            *pSemaphore = fakeHandle;
            return VK_SUCCESS;
    }

    void vkDestroySemaphore(
        VkDevice device,
        VkSemaphore semaphore,
        const VkAllocationCallbacks* pAllocator) {}

    VkResult vkCreateQueryPool(VkDevice device, const VkQueryPoolCreateInfo* pCreateInfo,
        const VkAllocationCallbacks* pAllocator, VkQueryPool* pQueryPool) {
            *pQueryPool = (VkQueryPool)0x00123;
            return VK_SUCCESS;
    }
}

class PipelineTrackerTestFixture : public ::testing::Test {
protected:
    VulkanPipelineTracker* tracker = nullptr;
    mCore m_vkCore;
    mDevice devA;
    mDevice devB;

    void SetUp() override {
        g_mockSems.clear();
        devA.virtualDevice = (VkDevice)0x1;
        devB.virtualDevice = (VkDevice)0x2;
        std::vector<StageConfig> stageMapping = {
            {&devA, nullptr,    "Stage0"},
            {&devA, &devB ,     "Stage1"},
            {&devB, &devA,      "Stage2"},
            {&devB, nullptr,    "Stage3"}
        };
        tracker = new VulkanPipelineTracker(&m_vkCore, stageMapping);
    }

    void TearDown() override {
        delete tracker;
        g_mockSems.clear();
    }
};

/**
 * @brief Verifies stage completion status and sem evaluations.
 *          Validates that the tracking matrix accurately catches matching,
 *          mismatched, or uninitialized host counters against mock
 *          GPU exec timeline registers before flagging readiness.
 */
TEST_F(PipelineTrackerTestFixture, VerifyStageCompletion) {
    u32 stageId = 0;
    u32 consumerStageId = 1;
    u64 latestStageSignaled = 0;
    EXPECT_TRUE(tracker->IsStageComplete(stageId, latestStageSignaled));

    StageContext* context0 = tracker->GetStageContext(stageId);
    context0->m_gpuExecutionCounter = 1;
    g_mockSems[context0->localTimelineSem.semaphore] = 1;
    EXPECT_TRUE(tracker->IsStageComplete(stageId, latestStageSignaled));

    context0->m_gpuExecutionCounter = 2;
    EXPECT_FALSE(tracker->IsStageComplete(stageId, latestStageSignaled));

    EXPECT_TRUE(tracker->IsIndexSafeToWrite(consumerStageId, 0));
    u64 targetLocalTimelineValue = 0;
    sem* localSemaphore = nullptr;
    tracker->AdvanceStageSync(stageId, targetLocalTimelineValue, localSemaphore, 0);
    EXPECT_EQ(targetLocalTimelineValue, 3);

    EXPECT_EQ(tracker->GetStageDevice(stageId), devA.virtualDevice);

    u64 latestProducerFrame = 0;
    EXPECT_EQ(tracker->getSafeIndex(stageId, latestProducerFrame), 1);
    EXPECT_EQ(latestProducerFrame, 1);
 
}

/**
 * @brief Verifies cross-device value replication.
 *          Validates that a multi-value producer jump forces the CPU
 *          mirroring loop to step the consumer timeline forward strictly
 *          by exactly +1 per call/frame without skipping.
 */
TEST_F(PipelineTrackerTestFixture, VerifySignalMirroring) {    
    u32 consumerStageId = 1;
    StageContext* context1 = tracker->GetStageContext(consumerStageId);
    g_mockSems[context1->localTimelineSem.semaphore] = 3;
    tracker->MirrorSignals();
    tracker->MirrorSignals();
    tracker->MirrorSignals();
    EXPECT_EQ(g_mockSems[context1->localTimelineSem.semaphoreRemote], 3);
}

int main(int argc, char** argv) {
    testing::InitGoogleTest(&argc, argv);
    assert(putenv(cfgTest::envSetting) == 0);
    return RUN_ALL_TESTS();
}