#include "sauna_device_driver.h"
#include <openvr_driver.h>
#include <cstring>
#include <cmath>

namespace sauna
{

SaunaDeviceDriver::SaunaDeviceDriver(vr::ITrackedDeviceServerDriver *pWrappedDriver, IMUDataProvider *pIMUDataProvider)
    : m_pWrappedDriver(pWrappedDriver)
    , m_pIMUDataProvider(pIMUDataProvider)
    , m_unDeviceId(vr::k_unTrackedDeviceIndexInvalid)
    , m_bActivated(false)
    , m_lastImuTime(0.0)
    , m_hasLastImuSample(false)
{
    // Initialize the last pose to identity
    m_lastPose.qRotation.w = 1.0;
    m_lastPose.qRotation.x = 0.0;
    m_lastPose.qRotation.y = 0.0;
    m_lastPose.qRotation.z = 0.0;
    
    for (int i = 0; i < 3; i++)
    {
        m_lastPose.vecPosition[i] = 0.0;
        m_lastPose.vecVelocity[i] = 0.0;
        m_lastPose.vecAcceleration[i] = 0.0;
        m_lastPose.vecAngularVelocity[i] = 0.0;
        m_lastPose.vecAngularAcceleration[i] = 0.0;
    }
    
    m_lastPose.result = vr::TrackingResult_Running_OK;
    m_lastPose.poseIsValid = true;
    m_lastPose.deviceIsConnected = true;
}

SaunaDeviceDriver::~SaunaDeviceDriver()
{
    // We don't own these pointers, so we don't delete them
    m_pWrappedDriver = nullptr;
    m_pIMUDataProvider = nullptr;
}

vr::EVRInitError SaunaDeviceDriver::Activate(uint32_t unObjectId)
{
    m_unDeviceId = unObjectId;
    m_bActivated = true;
    
    // Register this device with the IMU data provider
    if (m_pIMUDataProvider)
    {
        m_pIMUDataProvider->RegisterDevice(m_unDeviceId);
    }
    
    // Forward the call to the wrapped driver
    return m_pWrappedDriver->Activate(unObjectId);
}

void SaunaDeviceDriver::Deactivate()
{
    m_bActivated = false;
    m_unDeviceId = vr::k_unTrackedDeviceIndexInvalid;
    
    // Forward the call to the wrapped driver
    m_pWrappedDriver->Deactivate();
}

void SaunaDeviceDriver::EnterStandby()
{
    // Forward the call to the wrapped driver
    m_pWrappedDriver->EnterStandby();
}

void *SaunaDeviceDriver::GetComponent(const char *pchComponentNameAndVersion)
{
    // First check if the wrapped driver provides this component
    void *pComponent = m_pWrappedDriver->GetComponent(pchComponentNameAndVersion);
    if (pComponent)
    {
        return pComponent;
    }
    
    // If the component is our custom IMU component, return it
    if (strcmp(pchComponentNameAndVersion, IVRIMUComponent_Version) == 0)
    {
        return static_cast<IVRIMUComponent*>(this);
    }
    
    return nullptr;
}

void SaunaDeviceDriver::DebugRequest(const char *pchRequest, char *pchResponseBuffer, uint32_t unResponseBufferSize)
{
    // Forward the call to the wrapped driver
    m_pWrappedDriver->DebugRequest(pchRequest, pchResponseBuffer, unResponseBufferSize);
}

// Helper function to normalize a quaternion
void NormalizeQuaternion(vr::HmdQuaternion_t &q)
{
    double magnitude = sqrt(q.w * q.w + q.x * q.x + q.y * q.y + q.z * q.z);
    if (magnitude > 0.0)
    {
        q.w /= magnitude;
        q.x /= magnitude;
        q.y /= magnitude;
        q.z /= magnitude;
    }
    else
    {
        q.w = 1.0;
        q.x = 0.0;
        q.y = 0.0;
        q.z = 0.0;
    }
}

// Helper function to multiply two quaternions
vr::HmdQuaternion_t MultiplyQuaternions(const vr::HmdQuaternion_t &q1, const vr::HmdQuaternion_t &q2)
{
    vr::HmdQuaternion_t result;
    result.w = q1.w * q2.w - q1.x * q2.x - q1.y * q2.y - q1.z * q2.z;
    result.x = q1.w * q2.x + q1.x * q2.w + q1.y * q2.z - q1.z * q2.y;
    result.y = q1.w * q2.y - q1.x * q2.z + q1.y * q2.w + q1.z * q2.x;
    result.z = q1.w * q2.z + q1.x * q2.y - q1.y * q2.x + q1.z * q2.w;
    return result;
}

// Helper function to create a quaternion from axis-angle representation
vr::HmdQuaternion_t QuaternionFromAxisAngle(const double axis[3], double angle)
{
    vr::HmdQuaternion_t q;
    double halfAngle = angle * 0.5;
    double sinHalfAngle = sin(halfAngle);
    
    q.w = cos(halfAngle);
    q.x = axis[0] * sinHalfAngle;
    q.y = axis[1] * sinHalfAngle;
    q.z = axis[2] * sinHalfAngle;
    
    NormalizeQuaternion(q);
    return q;
}

vr::DriverPose_t SaunaDeviceDriver::GetPose()
{
    // Get the pose from the wrapped driver
    vr::DriverPose_t pose = m_pWrappedDriver->GetPose();
    
    // If optical tracking is lost, we can still provide IMU data
    // This is indicated by the tracking result
    if (pose.result == vr::TrackingResult_Fallback_RotationOnly ||
        pose.result == vr::TrackingResult_Calibrating_OutOfRange ||
        pose.result == vr::TrackingResult_Running_OutOfRange)
    {
        // Get the latest IMU data
        vr::ImuSample_t imuSample;
        if (m_pIMUDataProvider && m_pIMUDataProvider->GetLatestIMUSample(m_unDeviceId, &imuSample))
        {
            // Integrate the IMU data to update the pose
            IntegrateIMUData(pose, imuSample);
        }
    }
    else
    {
        // If we have good tracking, update our last known good pose
        m_lastPose = pose;
        
        // Also update our IMU calibration if we have IMU data
        vr::ImuSample_t imuSample;
        if (m_pIMUDataProvider && m_pIMUDataProvider->GetLatestIMUSample(m_unDeviceId, &imuSample))
        {
            m_lastImuSample = imuSample;
            m_lastImuTime = imuSample.fSampleTime;
            m_hasLastImuSample = true;
        }
    }
    
    return pose;
}

void SaunaDeviceDriver::IntegrateIMUData(vr::DriverPose_t &pose, const vr::ImuSample_t &imuSample)
{
    // If this is our first IMU sample, just store it and return
    if (!m_hasLastImuSample)
    {
        m_lastImuSample = imuSample;
        m_lastImuTime = imuSample.fSampleTime;
        m_hasLastImuSample = true;
        return;
    }
    
    // Calculate the time delta
    double dt = imuSample.fSampleTime - m_lastImuTime;
    if (dt <= 0.0)
    {
        // Invalid time delta, skip this sample
        return;
    }
    
    // Start with the last known good pose
    pose = m_lastPose;
    
    // Update the angular velocity from the gyroscope
    pose.vecAngularVelocity[0] = imuSample.vGyro.v[0];
    pose.vecAngularVelocity[1] = imuSample.vGyro.v[1];
    pose.vecAngularVelocity[2] = imuSample.vGyro.v[2];
    
    // Update the acceleration from the accelerometer
    pose.vecAcceleration[0] = imuSample.vAccel.v[0];
    pose.vecAcceleration[1] = imuSample.vAccel.v[1];
    pose.vecAcceleration[2] = imuSample.vAccel.v[2];
    
    // Integrate the gyroscope data to update the orientation
    // Convert angular velocity to axis-angle representation
    double angularVelocityMagnitude = sqrt(
        pose.vecAngularVelocity[0] * pose.vecAngularVelocity[0] +
        pose.vecAngularVelocity[1] * pose.vecAngularVelocity[1] +
        pose.vecAngularVelocity[2] * pose.vecAngularVelocity[2]);
    
    if (angularVelocityMagnitude > 0.0)
    {
        // Normalize the angular velocity to get the rotation axis
        double rotationAxis[3];
        rotationAxis[0] = pose.vecAngularVelocity[0] / angularVelocityMagnitude;
        rotationAxis[1] = pose.vecAngularVelocity[1] / angularVelocityMagnitude;
        rotationAxis[2] = pose.vecAngularVelocity[2] / angularVelocityMagnitude;
        
        // Calculate the rotation angle
        double rotationAngle = angularVelocityMagnitude * dt;
        
        // Create a quaternion from the axis-angle representation
        vr::HmdQuaternion_t deltaRotation = QuaternionFromAxisAngle(rotationAxis, rotationAngle);
        
        // Apply the rotation to the current orientation
        pose.qRotation = MultiplyQuaternions(pose.qRotation, deltaRotation);
        NormalizeQuaternion(pose.qRotation);
    }
    
    // Integrate the accelerometer data to update the velocity and position
    // First, remove gravity from the acceleration
    // Note: This is a simplified approach. In a real implementation, you would
    // need to transform the acceleration from the IMU frame to the world frame.
    double gravity[3] = { 0.0, 9.81, 0.0 }; // Assuming Y is up
    double linearAccel[3];
    linearAccel[0] = imuSample.vAccel.v[0] - gravity[0];
    linearAccel[1] = imuSample.vAccel.v[1] - gravity[1];
    linearAccel[2] = imuSample.vAccel.v[2] - gravity[2];
    
    // Update the velocity using the acceleration
    pose.vecVelocity[0] += linearAccel[0] * dt;
    pose.vecVelocity[1] += linearAccel[1] * dt;
    pose.vecVelocity[2] += linearAccel[2] * dt;
    
    // Apply a simple velocity decay to prevent drift
    const double velocityDecay = 0.95;
    pose.vecVelocity[0] *= velocityDecay;
    pose.vecVelocity[1] *= velocityDecay;
    pose.vecVelocity[2] *= velocityDecay;
    
    // Update the position using the velocity
    pose.vecPosition[0] += pose.vecVelocity[0] * dt;
    pose.vecPosition[1] += pose.vecVelocity[1] * dt;
    pose.vecPosition[2] += pose.vecVelocity[2] * dt;
    
    // Store the updated pose and IMU data for the next integration
    m_lastPose = pose;
    m_lastImuSample = imuSample;
    m_lastImuTime = imuSample.fSampleTime;
}

bool SaunaDeviceDriver::GetLatestIMUSample(vr::ImuSample_t *pSample)
{
    if (!m_bActivated || !m_pIMUDataProvider || !pSample)
    {
        return false;
    }
    
    return m_pIMUDataProvider->GetLatestIMUSample(m_unDeviceId, pSample);
}

bool SaunaDeviceDriver::IsIMUDataAvailable()
{
    if (!m_bActivated || !m_pIMUDataProvider)
    {
        return false;
    }
    
    return m_pIMUDataProvider->IsIMUDataAvailable(m_unDeviceId);
}

bool SaunaDeviceDriver::GetIMUDataInFallbackMode(vr::ImuSample_t *pSample)
{
    if (!m_bActivated || !m_pIMUDataProvider || !pSample)
    {
        return false;
    }
    
    // Get the pose to check if we're in fallback mode
    vr::DriverPose_t pose = m_pWrappedDriver->GetPose();
    
    // Check if optical tracking is lost
    if (pose.result == vr::TrackingResult_Fallback_RotationOnly ||
        pose.result == vr::TrackingResult_Calibrating_OutOfRange ||
        pose.result == vr::TrackingResult_Running_OutOfRange)
    {
        // Get the IMU data
        return m_pIMUDataProvider->GetLatestIMUSample(m_unDeviceId, pSample);
    }
    
    return false;
}

} // namespace sauna