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

using namespace SaunaTest;
using namespace sauna;

// Mock implementation of ITrackedDeviceServerDriver for testing
class MockTrackedDeviceDriver : public vr::ITrackedDeviceServerDriver {
private:
    bool activated;
    uint32_t deviceId;
    vr::DriverPose_t pose;
    
public:
    MockTrackedDeviceDriver() 
        : activated(false)
        , deviceId(vr::k_unTrackedDeviceIndexInvalid) {
        // Initialize the pose to identity
        pose = {};
        pose.qRotation.w = 1.0;
        pose.result = vr::TrackingResult_Running_OK;
        pose.poseIsValid = true;
        pose.deviceIsConnected = true;
    }
    
    virtual vr::EVRInitError Activate(uint32_t unObjectId) override {
        activated = true;
        deviceId = unObjectId;
        return vr::VRInitError_None;
    }
    
    virtual void Deactivate() override {
        activated = false;
        deviceId = vr::k_unTrackedDeviceIndexInvalid;
    }
    
    virtual void EnterStandby() override {
        // Nothing to do
    }
    
    virtual void *GetComponent(const char *pchComponentNameAndVersion) override {
        // We don't provide any components
        return nullptr;
    }
    
    virtual void DebugRequest(const char *pchRequest, char *pchResponseBuffer, uint32_t unResponseBufferSize) override {
        // Nothing to do
    }
    
    virtual vr::DriverPose_t GetPose() override {
        return pose;
    }
    
    // Mock-specific methods for test control
    void SetPose(const vr::DriverPose_t& newPose) {
        pose = newPose;
    }
    
    void SimulateOpticalTrackingLoss() {
        pose.result = vr::TrackingResult_Running_OutOfRange;
        pose.poseIsValid = false;
    }
    
    void SimulateOpticalTrackingRecovery() {
        pose.result = vr::TrackingResult_Running_OK;
        pose.poseIsValid = true;
    }
    
    bool IsActivated() const {
        return activated;
    }
    
    uint32_t GetDeviceId() const {
        return deviceId;
    }
};

// Test fixture for SaunaDeviceDriver tests
class SaunaDeviceDriverTest {
private:
    std::unique_ptr<IMUDataProvider> imuProvider;
    std::unique_ptr<MockTrackedDeviceDriver> mockDriver;
    std::unique_ptr<SaunaDeviceDriver> saunaDriver;

public:
public:
    SaunaDeviceDriverTest() {
        imuProvider = std::make_unique<IMUDataProvider>();
        imuProvider->Initialize();
        
        mockDriver = std::make_unique<MockTrackedDeviceDriver>();
        
        saunaDriver = std::make_unique<SaunaDeviceDriver>(mockDriver.get(), imuProvider.get());
    }

    ~SaunaDeviceDriverTest() {
        // Clean up
        if (mockDriver->IsActivated()) {
            saunaDriver->Deactivate();
        }
        imuProvider->Shutdown();
    }

    TestResult testActivate() {
        // Activate the driver
        uint32_t testDeviceId = 1;
        vr::EVRInitError error = saunaDriver->Activate(testDeviceId);
        
        // Check if activation succeeded
        if (error != vr::VRInitError_None) {
            return TestResult(false, "Activation failed with error: " + std::to_string(error));
        }
        
        // Check if the mock driver was activated
        if (!mockDriver->IsActivated()) {
            return TestResult(false, "Mock driver was not activated");
        }
        
        // Check if the device ID was passed correctly
        if (mockDriver->GetDeviceId() != testDeviceId) {
            return TestResult(false, "Device ID was not passed correctly to mock driver");
        }
        
        return TestResult(true, "Activation succeeded");
    }

    TestResult testDeactivate() {
        // Activate the driver first
        uint32_t testDeviceId = 2;
        saunaDriver->Activate(testDeviceId);
        
        // Deactivate the driver
        saunaDriver->Deactivate();
        
        // Check if the mock driver was deactivated
        if (mockDriver->IsActivated()) {
            return TestResult(false, "Mock driver was not deactivated");
        }
        
        return TestResult(true, "Deactivation succeeded");
    }

    TestResult testGetComponent() {
        // Activate the driver
        saunaDriver->Activate(3);
        
        // Try to get the IMU component
        void* component = saunaDriver->GetComponent(IVRIMUComponent_Version);
        
        // Check if the component was returned
        if (!component) {
            return TestResult(false, "IMU component was not returned");
        }
        
        // Check if the component is of the correct type
        IVRIMUComponent* imuComponent = static_cast<IVRIMUComponent*>(component);
        if (!imuComponent) {
            return TestResult(false, "Component is not an IMU component");
        }
        
        return TestResult(true, "GetComponent returned the IMU component");
    }

    TestResult testGetPose() {
        // Activate the driver
        saunaDriver->Activate(4);
        
        // Set a test pose in the mock driver
        vr::DriverPose_t testPose = {};
        testPose.qRotation.w = 0.5;
        testPose.qRotation.x = 0.5;
        testPose.qRotation.y = 0.5;
        testPose.qRotation.z = 0.5;
        testPose.vecPosition[0] = 1.0;
        testPose.vecPosition[1] = 2.0;
        testPose.vecPosition[2] = 3.0;
        testPose.result = vr::TrackingResult_Running_OK;
        testPose.poseIsValid = true;
        testPose.deviceIsConnected = true;
        
        mockDriver->SetPose(testPose);
        
        // Get the pose from the sauna driver
        vr::DriverPose_t returnedPose = saunaDriver->GetPose();
        
        // Check if the pose was passed through correctly
        if (returnedPose.qRotation.w != testPose.qRotation.w ||
            returnedPose.qRotation.x != testPose.qRotation.x ||
            returnedPose.qRotation.y != testPose.qRotation.y ||
            returnedPose.qRotation.z != testPose.qRotation.z ||
            returnedPose.vecPosition[0] != testPose.vecPosition[0] ||
            returnedPose.vecPosition[1] != testPose.vecPosition[1] ||
            returnedPose.vecPosition[2] != testPose.vecPosition[2] ||
            returnedPose.result != testPose.result ||
            returnedPose.poseIsValid != testPose.poseIsValid ||
            returnedPose.deviceIsConnected != testPose.deviceIsConnected) {
            return TestResult(false, "Pose was not passed through correctly");
        }
        
        return TestResult(true, "GetPose returned the correct pose");
    }

    TestResult testIMUDataAccess() {
        // Activate the driver
        uint32_t testDeviceId = 5;
        saunaDriver->Activate(testDeviceId);
        
        // Add an IMU sample
        vr::ImuSample_t sampleIn;
        sampleIn.fSampleTime = 1.0;
        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;
        
        imuProvider->AddIMUSample(testDeviceId, sampleIn);
        
        // Get the IMU component
        IVRIMUComponent* imuComponent = static_cast<IVRIMUComponent*>(
            saunaDriver->GetComponent(IVRIMUComponent_Version));
        
        if (!imuComponent) {
            return TestResult(false, "Failed to get IMU component");
        }
        
        // Check if IMU data is available
        bool dataAvailable = imuComponent->IsIMUDataAvailable();
        if (!dataAvailable) {
            return TestResult(false, "IMU data should be available");
        }
        
        // Get the IMU sample
        vr::ImuSample_t sampleOut;
        bool result = imuComponent->GetLatestIMUSample(&sampleOut);
        
        if (!result) {
            return TestResult(false, "Failed to get IMU sample");
        }
        
        // Verify the sample data
        if (sampleOut.fSampleTime != sampleIn.fSampleTime ||
            sampleOut.vAccel.v[0] != sampleIn.vAccel.v[0] ||
            sampleOut.vAccel.v[1] != sampleIn.vAccel.v[1] ||
            sampleOut.vAccel.v[2] != sampleIn.vAccel.v[2] ||
            sampleOut.vGyro.v[0] != sampleIn.vGyro.v[0] ||
            sampleOut.vGyro.v[1] != sampleIn.vGyro.v[1] ||
            sampleOut.vGyro.v[2] != sampleIn.vGyro.v[2]) {
            return TestResult(false, "IMU sample data mismatch");
        }
        
        return TestResult(true, "IMU data access succeeded");
    }

    TestResult testFallbackMode() {
        // Activate the driver
        uint32_t testDeviceId = 6;
        saunaDriver->Activate(testDeviceId);
        
        // Add an IMU sample
        vr::ImuSample_t sampleIn;
        sampleIn.fSampleTime = 2.0;
        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;
        
        imuProvider->AddIMUSample(testDeviceId, sampleIn);
        
        // Get the IMU component
        IVRIMUComponent* imuComponent = static_cast<IVRIMUComponent*>(
            saunaDriver->GetComponent(IVRIMUComponent_Version));
        
        if (!imuComponent) {
            return TestResult(false, "Failed to get IMU component");
        }
        
        // Simulate optical tracking loss
        mockDriver->SimulateOpticalTrackingLoss();
        
        // Get IMU data in fallback mode
        vr::ImuSample_t sampleOut;
        bool result = imuComponent->GetIMUDataInFallbackMode(&sampleOut);
        
        if (!result) {
            return TestResult(false, "Failed to get IMU data in fallback mode");
        }
        
        // Verify the sample data
        if (sampleOut.fSampleTime != sampleIn.fSampleTime ||
            sampleOut.vAccel.v[0] != sampleIn.vAccel.v[0] ||
            sampleOut.vAccel.v[1] != sampleIn.vAccel.v[1] ||
            sampleOut.vAccel.v[2] != sampleIn.vAccel.v[2] ||
            sampleOut.vGyro.v[0] != sampleIn.vGyro.v[0] ||
            sampleOut.vGyro.v[1] != sampleIn.vGyro.v[1] ||
            sampleOut.vGyro.v[2] != sampleIn.vGyro.v[2]) {
            return TestResult(false, "IMU sample data mismatch in fallback mode");
        }
        
        // Simulate optical tracking recovery
        mockDriver->SimulateOpticalTrackingRecovery();
        
        // Try to get IMU data in fallback mode again
        result = imuComponent->GetIMUDataInFallbackMode(&sampleOut);
        
        // This should fail because we're no longer in fallback mode
        if (result) {
            return TestResult(false, "GetIMUDataInFallbackMode should fail when not in fallback mode");
        }
        
        return TestResult(true, "Fallback mode test succeeded");
    }

    void runAllTests() {
        TestSuite suite("SaunaDeviceDriver Tests");
        
        suite.addTest("Activate", [this]() { return testActivate(); });
        suite.addTest("Deactivate", [this]() { return testDeactivate(); });
        suite.addTest("GetComponent", [this]() { return testGetComponent(); });
        suite.addTest("GetPose", [this]() { return testGetPose(); });
        suite.addTest("IMUDataAccess", [this]() { return testIMUDataAccess(); });
        suite.addTest("FallbackMode", [this]() { return testFallbackMode(); });
        
        suite.runAll();
    }
};

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