#define NOMINMAX

#include <iostream>
#include <Windows.h>
#include <GL/glew.h>
#include <SDL.h>
#include <SDL_syswm.h>
#include <Extras/OVR_Math.h>
#include <OVR_CAPI.h>
#include <OVR_CAPI_GL.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_gl_interop.h>
#include <sl/Camera.hpp>
#include "Shader.hpp"

// Добавляем перечисление для типов входных данных
enum class InputType {
    CAMERA_USB,
    STREAM,
    VIDEO_FILE
};

// Структура для конфигурации входных данных
struct InputConfig {
    InputType type;
    std::string source; // IP:port для стрима, путь к файлу для видео
};

// Модифицированная структура данных потока
struct ThreadData {
    sl::Camera zed;
    sl::Mat zed_image[2];
    sl::Mat depth_map; // Добавляем карту глубины
    std::mutex mtx;
    bool run;
    bool new_frame;
    bool show_depth; // Флаг для переключения между обычным видео и картой глубины
};

// Функция для настройки параметров входных данных
sl::InitParameters configureInputParameters(const InputConfig& config) {
    sl::InitParameters init_parameters;
    init_parameters.depth_mode = sl::DEPTH_MODE::ULTRA; // Включаем режим глубины
    init_parameters.coordinate_units = sl::UNIT::METER;
    init_parameters.depth_minimum_distance = 0.3f; // Минимальная дистанция в метрах
    init_parameters.camera_resolution = sl::RESOLUTION::HD720;

    switch (config.type) {
        case InputType::CAMERA_USB:
            // Настройки для USB камеры - используются параметры по умолчанию
            break;
        case InputType::STREAM:
            init_parameters.input.setFromStream(config.source.c_str());
            break;
        case InputType::VIDEO_FILE:
            init_parameters.input.setFromSVOFile(config.source.c_str());
            break;
    }

    return init_parameters;
}

// Модифицированная функция потока для захвата изображения
void __zed_runner__(ThreadData &thread_data) {
    sl::RuntimeParameters runtime_parameters;
    runtime_parameters.enable_depth = true; // Включаем получение данных глубины
    runtime_parameters.confidence_threshold = 50;
    runtime_parameters.texture_confidence_threshold = 100;

    cuCtxSetCurrent(thread_data.zed.getCUDAContext());

    while (thread_data.run) {
        if (thread_data.zed.grab(runtime_parameters) == sl::ERROR_CODE::SUCCESS) {
            thread_data.mtx.lock();
            
            // Получаем обычные изображения
            thread_data.zed.retrieveImage(thread_data.zed_image[0], sl::VIEW::LEFT, sl::MEM::GPU);
            thread_data.zed.retrieveImage(thread_data.zed_image[1], sl::VIEW::RIGHT, sl::MEM::GPU);
            
            // Получаем карту глубины
            if (thread_data.show_depth) {
                thread_data.zed.retrieveMeasure(thread_data.depth_map, sl::MEASURE::DEPTH, sl::MEM::GPU);
                // Конвертируем карту глубины в цветное изображение для визуализации
                thread_data.zed.normalizeDepthColorization(thread_data.depth_map);
            }
            
            thread_data.mtx.unlock();
            thread_data.new_frame = true;
        } else {
            sl::sleep_ms(2);
        }
    }

    // Освобождаем память
    for (int eye = 0; eye < 2; eye++)
        thread_data.zed_image[eye].free();
    thread_data.depth_map.free();
    thread_data.zed.close();
}

int main(int argc, char **argv) {
    // Получаем конфигурацию от пользователя
    InputConfig input_config;
    std::cout << "Select input source (1: USB Camera, 2: Stream, 3: Video File): ";
    int choice;
    std::cin >> choice;
    
    switch (choice) {
        case 1:
            input_config.type = InputType::CAMERA_USB;
            break;
        case 2:
            std::cout << "Enter stream address (IP:port): ";
            std::cin >> input_config.source;
            input_config.type = InputType::STREAM;
            break;
        case 3:
            std::cout << "Enter video file path: ";
            std::cin >> input_config.source;
            input_config.type = InputType::VIDEO_FILE;
            break;
        default:
            std::cout << "Invalid choice. Using USB Camera." << std::endl;
            input_config.type = InputType::CAMERA_USB;
    }

    // Инициализация SDL2
    SDL_Init(SDL_INIT_VIDEO);
    
    // Инициализация Oculus
    ovrResult result = ovr_Initialize(nullptr);
    if (OVR_FAILURE(result)) {
        std::cout << "ERROR: Failed to initialize libOVR" << std::endl;
        SDL_Quit();
        return -1;
    }

    // [Остальной код инициализации Oculus и SDL остается без изменений]

    // Инициализация камеры ZED с новыми параметрами
    ThreadData thread_data;
    sl::InitParameters init_parameters = configureInputParameters(input_config);
    
    sl::ERROR_CODE err_ = thread_data.zed.open(init_parameters);
    if (err_ != sl::ERROR_CODE::SUCCESS) {
        std::cout << "ERROR: " << sl::toString(err_) << std::endl;
        thread_data.zed.close();
        ovr_Destroy(session);
        ovr_Shutdown();
        SDL_GL_DeleteContext(glContext);
        SDL_DestroyWindow(window);
        SDL_Quit();
        return -1;
    }

    // Добавляем обработку клавиши D для переключения режима отображения глубины
    while (!end) {
        while (SDL_PollEvent(&events)) {
            if (events.type == SDL_KEYUP) {
                if (events.key.keysym.scancode == SDL_SCANCODE_Q)
                    end = true;
                else if (events.key.keysym.scancode == SDL_SCANCODE_D)
                    thread_data.show_depth = !thread_data.show_depth;
            }
        }
        
        // [Остальной код главного цикла остается без изменений]
    }

    // [Код очистки ресурсов остается без изменений]
    
    return 0;
}