#include "../lighthouse_driver_wrapper.h"
#include "../imu_data_provider.h"
#include "../sauna_device_driver.h"
#include "test_utils.h"
#include <openvr_driver.h>
#include <iostream>
#include <fstream>
#include <memory>
#include <thread>
#include <chrono>
#include <vector>
#include <numeric>
#include <algorithm>
#include <iomanip>
#include <ctime>

#ifdef _WIN32
#include <windows.h>
#include <psapi.h>
#elif defined(__linux__)
#include <sys/resource.h>
#include <unistd.h>
#elif defined(__APPLE__)
#include <mach/mach.h>
#endif

using namespace SaunaTest;
using namespace sauna;

// Utility class to measure CPU and memory usage
class SystemMetrics {
public:
    // Get current CPU usage (percentage)
    static double GetCPUUsage() {
#ifdef _WIN32
        static ULARGE_INTEGER lastCPU, lastSysCPU, lastUserCPU;
        static int numProcessors;
        static HANDLE self;
        static bool initialized = false;

        if (!initialized) {
            SYSTEM_INFO sysInfo;
            FILETIME ftime, fsys, fuser;

            GetSystemInfo(&sysInfo);
            numProcessors = sysInfo.dwNumberOfProcessors;

            self = GetCurrentProcess();
            GetProcessTimes(self, &ftime, &ftime, &fsys, &fuser);
            lastSysCPU.LowPart = fsys.dwLowDateTime;
            lastSysCPU.HighPart = fsys.dwHighDateTime;
            lastUserCPU.LowPart = fuser.dwLowDateTime;
            lastUserCPU.HighPart = fuser.dwHighDateTime;

            initialized = true;
        }

        FILETIME ftime, fsys, fuser;
        ULARGE_INTEGER now, sys, user;
        double percent;

        GetProcessTimes(self, &ftime, &ftime, &fsys, &fuser);
        sys.LowPart = fsys.dwLowDateTime;
        sys.HighPart = fsys.dwHighDateTime;
        user.LowPart = fuser.dwLowDateTime;
        user.HighPart = fuser.dwHighDateTime;
        now.QuadPart = GetTickCount64() * 10000;

        percent = (sys.QuadPart - lastSysCPU.QuadPart) +
                  (user.QuadPart - lastUserCPU.QuadPart);
        percent /= (now.QuadPart - lastCPU.QuadPart);
        percent /= numProcessors;
        percent *= 100;

        lastCPU = now;
        lastSysCPU = sys;
        lastUserCPU = user;

        return percent;
#elif defined(__linux__)
        // Linux implementation
        struct rusage usage;
        getrusage(RUSAGE_SELF, &usage);
        
        static struct rusage lastUsage;
        static std::chrono::high_resolution_clock::time_point lastTime;
        static bool initialized = false;
        
        if (!initialized) {
            lastUsage = usage;
            lastTime = std::chrono::high_resolution_clock::now();
            initialized = true;
            return 0.0;
        }
        
        auto now = std::chrono::high_resolution_clock::now();
        double timeElapsed = std::chrono::duration<double>(now - lastTime).count();
        
        double userTime = (usage.ru_utime.tv_sec + usage.ru_utime.tv_usec / 1000000.0) -
                         (lastUsage.ru_utime.tv_sec + lastUsage.ru_utime.tv_usec / 1000000.0);
        double systemTime = (usage.ru_stime.tv_sec + usage.ru_stime.tv_usec / 1000000.0) -
                           (lastUsage.ru_stime.tv_sec + lastUsage.ru_stime.tv_usec / 1000000.0);
        
        double cpuUsage = (userTime + systemTime) / timeElapsed * 100.0;
        
        lastUsage = usage;
        lastTime = now;
        
        return cpuUsage;
#elif defined(__APPLE__)
        // macOS implementation
        static bool initialized = false;
        static mach_port_t self;
        static std::chrono::high_resolution_clock::time_point lastTime;
        static double lastTotalTime = 0.0;
        
        if (!initialized) {
            self = mach_task_self();
            lastTime = std::chrono::high_resolution_clock::now();
            initialized = true;
            return 0.0;
        }
        
        task_info_data_t tinfo;
        mach_msg_type_number_t task_info_count = TASK_INFO_MAX;
        
        if (task_info(self, TASK_BASIC_INFO, (task_info_t)tinfo, &task_info_count) != KERN_SUCCESS)
            return 0.0;
        
        struct task_basic_info* ti = (struct task_basic_info*)tinfo;
        
        thread_array_t thread_list;
        mach_msg_type_number_t thread_count;
        
        if (task_threads(self, &thread_list, &thread_count) != KERN_SUCCESS)
            return 0.0;
        
        double totalTime = 0.0;
        
        for (int i = 0; i < thread_count; i++) {
            thread_info_data_t thinfo;
            mach_msg_type_number_t thread_info_count = THREAD_INFO_MAX;
            
            if (thread_info(thread_list[i], THREAD_BASIC_INFO, (thread_info_t)thinfo, &thread_info_count) != KERN_SUCCESS)
                continue;
            
            struct thread_basic_info* tbi = (struct thread_basic_info*)thinfo;
            
            totalTime += tbi->user_time.seconds + tbi->user_time.microseconds / 1000000.0;
            totalTime += tbi->system_time.seconds + tbi->system_time.microseconds / 1000000.0;
        }
        
        auto now = std::chrono::high_resolution_clock::now();
        double timeElapsed = std::chrono::duration<double>(now - lastTime).count();
        
        double cpuUsage = (totalTime - lastTotalTime) / timeElapsed * 100.0;
        
        lastTotalTime = totalTime;
        lastTime = now;
        
        return cpuUsage;
#else
        return 0.0; // Unsupported platform
#endif
    }

    // Get current memory usage (in MB)
    static double GetMemoryUsage() {
#ifdef _WIN32
        PROCESS_MEMORY_COUNTERS_EX pmc;
        GetProcessMemoryInfo(GetCurrentProcess(), (PROCESS_MEMORY_COUNTERS*)&pmc, sizeof(pmc));
        return static_cast<double>(pmc.WorkingSetSize) / (1024 * 1024);
#elif defined(__linux__)
        // Linux implementation
        FILE* file = fopen("/proc/self/status", "r");
        if (file == nullptr)
            return 0.0;
        
        char line[128];
        int vmSize = 0;
        
        while (fgets(line, 128, file) != nullptr) {
            if (strncmp(line, "VmRSS:", 6) == 0) {
                char* p = line + 6;
                while (*p == ' ' || *p == '\t') p++;
                vmSize = atoi(p);
                break;
            }
        }
        
        fclose(file);
        return vmSize / 1024.0; // Convert KB to MB
#elif defined(__APPLE__)
        // macOS implementation
        task_info_data_t tinfo;
        mach_msg_type_number_t task_info_count = TASK_INFO_MAX;
        
        if (task_info(mach_task_self(), TASK_BASIC_INFO, (task_info_t)tinfo, &task_info_count) != KERN_SUCCESS)
            return 0.0;
        
        struct task_basic_info* ti = (struct task_basic_info*)tinfo;
        return ti->resident_size / (1024 * 1024.0);
#else
        return 0.0; // Unsupported platform
#endif
    }
};

// Performance test for the sauna driver
class SaunaDriverPerformanceTest {
private:
    std::unique_ptr<LighthouseDriverWrapper> lighthouseWrapper;
    std::unique_ptr<IMUDataProvider> imuProvider;
    std::ofstream reportFile;
    
    // Test parameters
    const int warmupFrames = 100;
    const int testFrames = 1000;
    const int imuSamplesPerTest = 1000;
    
public:
    SaunaDriverPerformanceTest() {
        lighthouseWrapper = std::make_unique<LighthouseDriverWrapper>();
        imuProvider = std::make_unique<IMUDataProvider>();
        
        // Create a timestamped report file
        std::time_t now = std::time(nullptr);
        std::tm* timeinfo = std::localtime(&now);
        char buffer[80];
        std::strftime(buffer, sizeof(buffer), "performance_report_%Y%m%d_%H%M%S.txt", timeinfo);
        
        reportFile.open(buffer);
        if (!reportFile.is_open()) {
            std::cerr << "Failed to create report file" << std::endl;
        }
    }
    
    ~SaunaDriverPerformanceTest() {
        // Clean up
        lighthouseWrapper->Shutdown();
        imuProvider->Shutdown();
        
        if (reportFile.is_open()) {
            reportFile.close();
        }
    }
    
    void writeToReport(const std::string& text) {
        std::cout << text << std::endl;
        if (reportFile.is_open()) {
            reportFile << text << std::endl;
        }
    }
    
    TestResult testCPUUsage() {
        writeToReport("\n=== CPU Usage Test ===");
        writeToReport("This test measures CPU usage with and without the sauna driver.");
        
        // Initialize components
        bool lighthouseInitResult = lighthouseWrapper->Initialize();
        bool imuInitResult = imuProvider->Initialize();
        
        if (!lighthouseInitResult) {
            writeToReport("Lighthouse driver initialization failed - skipping test");
            return TestResult(true, "CPU usage test skipped");
        }
        
        // Warm up
        writeToReport("Warming up...");
        for (int i = 0; i < warmupFrames; i++) {
            lighthouseWrapper->RunFrame();
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
        
        // Measure baseline CPU usage (lighthouse driver only)
        writeToReport("Measuring baseline CPU usage (lighthouse driver only)...");
        std::vector<double> baselineCPU;
        
        for (int i = 0; i < testFrames; i++) {
            double cpuUsage = SystemMetrics::GetCPUUsage();
            baselineCPU.push_back(cpuUsage);
            
            lighthouseWrapper->RunFrame();
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
        
        // Calculate baseline statistics
        double baselineAvg = std::accumulate(baselineCPU.begin(), baselineCPU.end(), 0.0) / baselineCPU.size();
        double baselineMax = *std::max_element(baselineCPU.begin(), baselineCPU.end());
        
        // Register a test device with the IMU provider
        uint32_t testDeviceId = 0; // HMD is usually device 0
        imuProvider->RegisterDevice(testDeviceId);
        
        // Measure CPU usage with sauna driver active
        writeToReport("Measuring CPU usage with sauna driver active...");
        std::vector<double> saunaCPU;
        
        for (int i = 0; i < testFrames; i++) {
            // Add an IMU sample
            vr::ImuSample_t sample;
            sample.fSampleTime = static_cast<float>(i) * 0.01f;
            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;
            
            imuProvider->AddIMUSample(testDeviceId, sample);
            
            double cpuUsage = SystemMetrics::GetCPUUsage();
            saunaCPU.push_back(cpuUsage);
            
            lighthouseWrapper->RunFrame();
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
        
        // Calculate sauna driver statistics
        double saunaAvg = std::accumulate(saunaCPU.begin(), saunaCPU.end(), 0.0) / saunaCPU.size();
        double saunaMax = *std::max_element(saunaCPU.begin(), saunaCPU.end());
        
        // Report results
        writeToReport("CPU Usage Results:");
        writeToReport("  Baseline (lighthouse driver only):");
        writeToReport("    Average: " + std::to_string(baselineAvg) + "%");
        writeToReport("    Maximum: " + std::to_string(baselineMax) + "%");
        writeToReport("  With Sauna Driver:");
        writeToReport("    Average: " + std::to_string(saunaAvg) + "%");
        writeToReport("    Maximum: " + std::to_string(saunaMax) + "%");
        writeToReport("  Difference:");
        writeToReport("    Average: " + std::to_string(saunaAvg - baselineAvg) + "%");
        writeToReport("    Maximum: " + std::to_string(saunaMax - baselineMax) + "%");
        
        return TestResult(true, "CPU usage test completed");
    }
    
    TestResult testMemoryUsage() {
        writeToReport("\n=== Memory Usage Test ===");
        writeToReport("This test measures memory usage with and without the sauna driver.");
        
        // Measure baseline memory usage (before initialization)
        double baselineMemory = SystemMetrics::GetMemoryUsage();
        writeToReport("Baseline memory usage (before initialization): " + 
                     std::to_string(baselineMemory) + " MB");
        
        // Initialize lighthouse driver
        bool lighthouseInitResult = lighthouseWrapper->Initialize();
        if (!lighthouseInitResult) {
            writeToReport("Lighthouse driver initialization failed - skipping test");
            return TestResult(true, "Memory usage test skipped");
        }
        
        // Measure memory usage with lighthouse driver
        double lighthouseMemory = SystemMetrics::GetMemoryUsage();
        writeToReport("Memory usage with lighthouse driver: " + 
                     std::to_string(lighthouseMemory) + " MB");
        writeToReport("Lighthouse driver memory impact: " + 
                     std::to_string(lighthouseMemory - baselineMemory) + " MB");
        
        // Initialize IMU provider
        bool imuInitResult = imuProvider->Initialize();
        
        // Measure memory usage with sauna driver
        double saunaMemory = SystemMetrics::GetMemoryUsage();
        writeToReport("Memory usage with sauna driver: " + 
                     std::to_string(saunaMemory) + " MB");
        writeToReport("Sauna driver additional memory impact: " + 
                     std::to_string(saunaMemory - lighthouseMemory) + " MB");
        writeToReport("Total memory impact: " + 
                     std::to_string(saunaMemory - baselineMemory) + " MB");
        
        // Register a test device and add IMU samples
        uint32_t testDeviceId = 0;
        imuProvider->RegisterDevice(testDeviceId);
        
        for (int i = 0; i < imuSamplesPerTest; i++) {
            vr::ImuSample_t sample;
            sample.fSampleTime = static_cast<float>(i) * 0.01f;
            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;
            
            imuProvider->AddIMUSample(testDeviceId, sample);
        }
        
        // Measure memory usage after adding IMU samples
        double withSamplesMemory = SystemMetrics::GetMemoryUsage();
        writeToReport("Memory usage after adding " + std::to_string(imuSamplesPerTest) + 
                     " IMU samples: " + std::to_string(withSamplesMemory) + " MB");
        writeToReport("IMU samples memory impact: " + 
                     std::to_string(withSamplesMemory - saunaMemory) + " MB");
        
        return TestResult(true, "Memory usage test completed");
    }
    
    TestResult testIMUThroughput() {
        writeToReport("\n=== IMU Throughput Test ===");
        writeToReport("This test measures the IMU data throughput.");
        
        // Initialize components
        bool lighthouseInitResult = lighthouseWrapper->Initialize();
        bool imuInitResult = imuProvider->Initialize();
        
        if (!lighthouseInitResult || !imuInitResult) {
            writeToReport("Initialization failed - skipping test");
            return TestResult(true, "IMU throughput test skipped");
        }
        
        // Register a test device
        uint32_t testDeviceId = 0;
        imuProvider->RegisterDevice(testDeviceId);
        
        // Measure time to add samples
        writeToReport("Measuring time to add " + std::to_string(imuSamplesPerTest) + " IMU samples...");
        
        auto startAddTime = std::chrono::high_resolution_clock::now();
        
        for (int i = 0; i < imuSamplesPerTest; i++) {
            vr::ImuSample_t sample;
            sample.fSampleTime = static_cast<float>(i) * 0.01f;
            sample.vAccel.v[0] = static_cast<float>(i % 10) * 0.1f;
            sample.vAccel.v[1] = 9.81f;
            sample.vAccel.v[2] = static_cast<float>((i + 5) % 10) * 0.1f;
            sample.vGyro.v[0] = static_cast<float>(i % 5) * 0.01f;
            sample.vGyro.v[1] = static_cast<float>((i + 2) % 5) * 0.01f;
            sample.vGyro.v[2] = static_cast<float>((i + 4) % 5) * 0.01f;
            
            imuProvider->AddIMUSample(testDeviceId, sample);
        }
        
        auto endAddTime = std::chrono::high_resolution_clock::now();
        auto addDuration = std::chrono::duration_cast<std::chrono::microseconds>(endAddTime - startAddTime);
        
        double addThroughput = static_cast<double>(imuSamplesPerTest) / (addDuration.count() / 1000000.0);
        
        // Measure time to retrieve samples
        writeToReport("Measuring time to retrieve IMU samples...");
        
        auto startGetTime = std::chrono::high_resolution_clock::now();
        int samplesRetrieved = 0;
        
        while (imuProvider->IsIMUDataAvailable(testDeviceId) && samplesRetrieved < imuSamplesPerTest) {
            vr::ImuSample_t sample;
            if (imuProvider->GetLatestIMUSample(testDeviceId, &sample)) {
                samplesRetrieved++;
            }
        }
        
        auto endGetTime = std::chrono::high_resolution_clock::now();
        auto getDuration = std::chrono::duration_cast<std::chrono::microseconds>(endGetTime - startGetTime);
        
        double getThroughput = static_cast<double>(samplesRetrieved) / (getDuration.count() / 1000000.0);
        
        // Report results
        writeToReport("IMU Throughput Results:");
        writeToReport("  Add operation:");
        writeToReport("    Samples: " + std::to_string(imuSamplesPerTest));
        writeToReport("    Time: " + std::to_string(addDuration.count() / 1000.0) + " ms");
        writeToReport("    Throughput: " + std::to_string(addThroughput) + " samples/second");
        writeToReport("  Get operation:");
        writeToReport("    Samples: " + std::to_string(samplesRetrieved));
        writeToReport("    Time: " + std::to_string(getDuration.count() / 1000.0) + " ms");
        writeToReport("    Throughput: " + std::to_string(getThroughput) + " samples/second");
        
        return TestResult(true, "IMU throughput test completed");
    }
    
    TestResult testLatencyImpact() {
        writeToReport("\n=== Latency Impact Test ===");
        writeToReport("This test measures the latency impact of the sauna driver.");
        
        // Initialize components
        bool lighthouseInitResult = lighthouseWrapper->Initialize();
        bool imuInitResult = imuProvider->Initialize();
        
        if (!lighthouseInitResult || !imuInitResult) {
            writeToReport("Initialization failed - skipping test");
            return TestResult(true, "Latency impact test skipped");
        }
        
        // Measure baseline frame time (lighthouse driver only)
        writeToReport("Measuring baseline frame time (lighthouse driver only)...");
        std::vector<double> baselineFrameTimes;
        
        for (int i = 0; i < testFrames; i++) {
            auto startTime = std::chrono::high_resolution_clock::now();
            
            lighthouseWrapper->RunFrame();
            
            auto endTime = std::chrono::high_resolution_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::microseconds>(endTime - startTime);
            
            baselineFrameTimes.push_back(duration.count());
            
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
        
        // Calculate baseline statistics
        double baselineAvg = std::accumulate(baselineFrameTimes.begin(), baselineFrameTimes.end(), 0.0) / baselineFrameTimes.size();
        double baselineMax = *std::max_element(baselineFrameTimes.begin(), baselineFrameTimes.end());
        
        // Register a test device with the IMU provider
        uint32_t testDeviceId = 0;
        imuProvider->RegisterDevice(testDeviceId);
        
        // Measure frame time with sauna driver active
        writeToReport("Measuring frame time with sauna driver active...");
        std::vector<double> saunaFrameTimes;
        
        for (int i = 0; i < testFrames; i++) {
            // Add an IMU sample
            vr::ImuSample_t sample;
            sample.fSampleTime = static_cast<float>(i) * 0.01f;
            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;
            
            auto startTime = std::chrono::high_resolution_clock::now();
            
            imuProvider->AddIMUSample(testDeviceId, sample);
            lighthouseWrapper->RunFrame();
            
            auto endTime = std::chrono::high_resolution_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::microseconds>(endTime - startTime);
            
            saunaFrameTimes.push_back(duration.count());
            
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
        
        // Calculate sauna driver statistics
        double saunaAvg = std::accumulate(saunaFrameTimes.begin(), saunaFrameTimes.end(), 0.0) / saunaFrameTimes.size();
        double saunaMax = *std::max_element(saunaFrameTimes.begin(), saunaFrameTimes.end());
        
        // Report results
        writeToReport("Latency Impact Results:");
        writeToReport("  Baseline (lighthouse driver only):");
        writeToReport("    Average frame time: " + std::to_string(baselineAvg) + " µs");
        writeToReport("    Maximum frame time: " + std::to_string(baselineMax) + " µs");
        writeToReport("  With Sauna Driver:");
        writeToReport("    Average frame time: " + std::to_string(saunaAvg) + " µs");
        writeToReport("    Maximum frame time: " + std::to_string(saunaMax) + " µs");
        writeToReport("  Difference:");
        writeToReport("    Average: " + std::to_string(saunaAvg - baselineAvg) + " µs");
        writeToReport("    Maximum: " + std::to_string(saunaMax - baselineMax) + " µs");
        
        return TestResult(true, "Latency impact test completed");
    }
    
public:
    void runAllTests() {
        // Print header
        writeToReport("=== Sauna Driver Performance Tests ===");
        writeToReport("Date: " + getCurrentDateTime());
        writeToReport("System Information:");
        writeToReport("  OS: " + getOSInfo());
        writeToReport("  CPU: " + getCPUInfo());
        writeToReport("  Memory: " + getMemoryInfo());
        
        TestSuite suite("Sauna Driver Performance Tests");
        
        suite.addTest("CPU Usage", [this]() { return testCPUUsage(); });
        suite.addTest("Memory Usage", [this]() { return testMemoryUsage(); });
        suite.addTest("IMU Throughput", [this]() { return testIMUThroughput(); });
        suite.addTest("Latency Impact", [this]() { return testLatencyImpact(); });
        
        suite.runAll();
        
        // Print summary
        writeToReport("\n=== Performance Test Summary ===");
        writeToReport("All performance tests completed.");
        writeToReport("Results have been saved to the performance report file.");
    }
    
private:
    // Helper functions to get system information
    std::string getCurrentDateTime() {
        std::time_t now = std::time(nullptr);
        std::tm* timeinfo = std::localtime(&now);
        char buffer[80];
        std::strftime(buffer, sizeof(buffer), "%Y-%m-%d %H:%M:%S", timeinfo);
        return std::string(buffer);
    }
    
    std::string getOSInfo() {
#ifdef _WIN32
        return "Windows";
#elif defined(__linux__)
        return "Linux";
#elif defined(__APPLE__)
        return "macOS";
#else
        return "Unknown";
#endif
    }
    
    std::string getCPUInfo() {
        // This is a simplified implementation
#ifdef _WIN32
        return "Windows CPU";
#elif defined(__linux__)
        return "Linux CPU";
#elif defined(__APPLE__)
        return "macOS CPU";
#else
        return "Unknown CPU";
#endif
    }
    
    std::string getMemoryInfo() {
        // Get total system memory
#ifdef _WIN32
        MEMORYSTATUSEX memInfo;
        memInfo.dwLength = sizeof(MEMORYSTATUSEX);
        GlobalMemoryStatusEx(&memInfo);
        double totalPhysMem = static_cast<double>(memInfo.ullTotalPhys) / (1024 * 1024 * 1024);
        return std::to_string(totalPhysMem) + " GB";
#elif defined(__linux__)
        return "Linux Memory";
#elif defined(__APPLE__)
        return "macOS Memory";
#else
        return "Unknown Memory";
#endif
    }
};

int main(int argc, char** argv) {
    std::cout << "=== Sauna Driver Performance Tests ===" << std::endl;
    std::cout << "These tests measure the performance impact of the sauna driver." << std::endl;
    std::cout << "Results will be saved to a performance report file." << std::endl;
    
    SaunaDriverPerformanceTest test;
    test.runAllTests();
    
    std::cout << "\nAll performance tests completed." << std::endl;
    std::cout << "Press any key to exit..." << std::endl;
    std::cin.get();
    
    return 0;
}