#include "../imu_data_provider.h"
#include "test_utils.h"
#include <openvr_driver.h>
#include <iostream>
#include <memory>

using namespace SaunaTest;
using namespace sauna;

// Test fixture for IMUDataProvider tests
class IMUDataProviderTest {
private:
    std::unique_ptr<IMUDataProvider> provider;

public:
    IMUDataProviderTest() {
        provider = std::make_unique<IMUDataProvider>();
    }

    ~IMUDataProviderTest() {
        provider->Shutdown();
    }

    TestResult testInitialize() {
        bool result = provider->Initialize();
        return assertTrue(result, "IMUDataProvider initialization failed");
    }

    TestResult testRegisterDevice() {
        provider->Initialize();
        
        // Register a test device
        uint32_t testDeviceId = 1;
        provider->RegisterDevice(testDeviceId);
        
        // Check if the device is registered by trying to add a sample
        vr::ImuSample_t sample;
        sample.nTimeInMicroSec = 1000000;
        sample.vAccel.v[0] = 0.0f;
        sample.vAccel.v[1] = 9.81f;
        sample.vAccel.v[2] = 0.0f;
        sample.vGyro.v[0] = 0.0f;
        sample.vGyro.v[1] = 0.0f;
        sample.vGyro.v[2] = 0.0f;
        
        provider->AddIMUSample(testDeviceId, sample);
        
        bool dataAvailable = provider->IsIMUDataAvailable(testDeviceId);
        return assertTrue(dataAvailable, "Device registration failed");
    }

    TestResult testAddAndGetIMUSample() {
        provider->Initialize();
        
        // Register a test device
        uint32_t testDeviceId = 2;
        provider->RegisterDevice(testDeviceId);
        
        // Create a test sample
        vr::ImuSample_t sampleIn;
        sampleIn.nTimeInMicroSec = 2000000;
        sampleIn.vAccel.v[0] = 1.0f;
        sampleIn.vAccel.v[1] = 2.0f;
        sampleIn.vAccel.v[2] = 3.0f;
        sampleIn.vGyro.v[0] = 0.1f;
        sampleIn.vGyro.v[1] = 0.2f;
        sampleIn.vGyro.v[2] = 0.3f;
        
        // Add the sample
        provider->AddIMUSample(testDeviceId, sampleIn);
        
        // Get the sample back
        vr::ImuSample_t sampleOut;
        bool result = provider->GetLatestIMUSample(testDeviceId, &sampleOut);
        
        if (!result) {
            return TestResult(false, "Failed to get IMU sample");
        }
        
        // Verify the sample data
        bool timeMatches = (sampleIn.nTimeInMicroSec == sampleOut.nTimeInMicroSec);
        bool accelMatches = (sampleIn.vAccel.v[0] == sampleOut.vAccel.v[0]) &&
                           (sampleIn.vAccel.v[1] == sampleOut.vAccel.v[1]) &&
                           (sampleIn.vAccel.v[2] == sampleOut.vAccel.v[2]);
        bool gyroMatches = (sampleIn.vGyro.v[0] == sampleOut.vGyro.v[0]) &&
                          (sampleIn.vGyro.v[1] == sampleOut.vGyro.v[1]) &&
                          (sampleIn.vGyro.v[2] == sampleOut.vGyro.v[2]);
        
        if (!timeMatches) {
            return TestResult(false, "Time mismatch in IMU sample");
        }
        if (!accelMatches) {
            return TestResult(false, "Accelerometer data mismatch in IMU sample");
        }
        if (!gyroMatches) {
            return TestResult(false, "Gyroscope data mismatch in IMU sample");
        }
        
        return TestResult(true, "IMU sample add and get successful");
    }

    TestResult testQueueBehavior() {
        provider->Initialize();
        
        // Register a test device
        uint32_t testDeviceId = 3;
        provider->RegisterDevice(testDeviceId);
        
        // Add multiple samples
        const int numSamples = 150; // More than the max queue size (100)
        for (int i = 0; i < numSamples; i++) {
            vr::ImuSample_t sample;
            sample.nTimeInMicroSec = 3000000 + i;
            sample.vAccel.v[0] = static_cast<float>(i);
            sample.vAccel.v[1] = 0.0f;
            sample.vAccel.v[2] = 0.0f;
            sample.vGyro.v[0] = 0.0f;
            sample.vGyro.v[1] = 0.0f;
            sample.vGyro.v[2] = 0.0f;
            
            provider->AddIMUSample(testDeviceId, sample);
        }
        
        // Get the first sample, which should be the oldest one that's still in the queue
        // Since the queue size is limited to 100, the first 50 samples should have been discarded
        vr::ImuSample_t firstSample;
        bool result = provider->GetLatestIMUSample(testDeviceId, &firstSample);
        
        if (!result) {
            return TestResult(false, "Failed to get first IMU sample");
        }
        
        // The first sample should have a timestamp of 3000000 + (numSamples - 100)
        // since the first (numSamples - 100) samples should have been discarded
        uint64_t expectedTime = 3000000 + (numSamples - 100);
        if (firstSample.nTimeInMicroSec != expectedTime) {
            std::string errorMsg = "Queue behavior incorrect. Expected time: " + 
                                  std::to_string(expectedTime) + 
                                  ", Actual time: " + 
                                  std::to_string(firstSample.nTimeInMicroSec);
            return TestResult(false, errorMsg);
        }
        
        return TestResult(true, "Queue behavior correct");
    }

    TestResult testMultipleDevices() {
        provider->Initialize();
        
        // Register multiple test devices
        uint32_t device1 = 4;
        uint32_t device2 = 5;
        provider->RegisterDevice(device1);
        provider->RegisterDevice(device2);
        
        // Add samples to both devices
        vr::ImuSample_t sample1;
        sample1.nTimeInMicroSec = 4000000;
        sample1.vAccel.v[0] = 1.0f;
        sample1.vAccel.v[1] = 0.0f;
        sample1.vAccel.v[2] = 0.0f;
        sample1.vGyro.v[0] = 0.1f;
        sample1.vGyro.v[1] = 0.0f;
        sample1.vGyro.v[2] = 0.0f;
        
        vr::ImuSample_t sample2;
        sample2.nTimeInMicroSec = 5000000;
        sample2.vAccel.v[0] = 0.0f;
        sample2.vAccel.v[1] = 2.0f;
        sample2.vAccel.v[2] = 0.0f;
        sample2.vGyro.v[0] = 0.0f;
        sample2.vGyro.v[1] = 0.2f;
        sample2.vGyro.v[2] = 0.0f;
        
        provider->AddIMUSample(device1, sample1);
        provider->AddIMUSample(device2, sample2);
        
        // Verify that each device has its own data
        vr::ImuSample_t outSample1;
        vr::ImuSample_t outSample2;
        bool result1 = provider->GetLatestIMUSample(device1, &outSample1);
        bool result2 = provider->GetLatestIMUSample(device2, &outSample2);
        
        if (!result1 || !result2) {
            return TestResult(false, "Failed to get IMU samples for multiple devices");
        }
        
        // Verify the sample data
        if (outSample1.nTimeInMicroSec != sample1.nTimeInMicroSec ||
            outSample1.vAccel.v[0] != sample1.vAccel.v[0] ||
            outSample1.vGyro.v[0] != sample1.vGyro.v[0]) {
            return TestResult(false, "Data mismatch for device 1");
        }
        
        if (outSample2.nTimeInMicroSec != sample2.nTimeInMicroSec ||
            outSample2.vAccel.v[1] != sample2.vAccel.v[1] ||
            outSample2.vGyro.v[1] != sample2.vGyro.v[1]) {
            return TestResult(false, "Data mismatch for device 2");
        }
        
        return TestResult(true, "Multiple device support verified");
    }

    TestResult testShutdown() {
        provider->Initialize();
        
        // Register a device and add a sample
        uint32_t testDeviceId = 6;
        provider->RegisterDevice(testDeviceId);
        
        vr::ImuSample_t sample;
        sample.nTimeInMicroSec = 6000000;
        provider->AddIMUSample(testDeviceId, sample);
        
        // Verify the sample is available
        bool dataAvailableBefore = provider->IsIMUDataAvailable(testDeviceId);
        
        // Shutdown the provider
        provider->Shutdown();
        
        // Verify the sample is no longer available
        bool dataAvailableAfter = provider->IsIMUDataAvailable(testDeviceId);
        
        if (!dataAvailableBefore) {
            return TestResult(false, "Data should be available before shutdown");
        }
        
        if (dataAvailableAfter) {
            return TestResult(false, "Data should not be available after shutdown");
        }
        
        return TestResult(true, "Shutdown behavior correct");
    }

    void runAllTests() {
        TestSuite suite("IMUDataProvider Tests");
        
        suite.addTest("Initialize", [this]() { return testInitialize(); });
        suite.addTest("RegisterDevice", [this]() { return testRegisterDevice(); });
        suite.addTest("AddAndGetIMUSample", [this]() { return testAddAndGetIMUSample(); });
        suite.addTest("QueueBehavior", [this]() { return testQueueBehavior(); });
        suite.addTest("MultipleDevices", [this]() { return testMultipleDevices(); });
        suite.addTest("Shutdown", [this]() { return testShutdown(); });
        
        suite.runAll();
    }
};

int main(int argc, char** argv) {
    std::cout << "Running IMUDataProvider tests..." << std::endl;
    
    IMUDataProviderTest test;
    test.runAllTests();
    
    return 0;
}