#include <gtest/gtest.h>
#include "pipeline/inference/preprocess.hpp"
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

using namespace pipeline;

class PreprocessorTest : public ::testing::Test {
protected:
    void SetUp() override {
        // 创建测试图像
        test_image_ = cv::Mat(480, 640, CV_8UC3, cv::Scalar(0, 0, 0));
        // 绘制一些测试图案
        cv::circle(test_image_, cv::Point(320, 240), 100, cv::Scalar(255, 0, 0), -1);
        cv::rectangle(test_image_, cv::Rect(100, 100, 200, 200), cv::Scalar(0, 255, 0), -1);
        
        // 创建预处理器配置
        config_.input_shape = {3, 640, 640};  // [channels, height, width]
        config_.use_cuda = false;             // 先测试CPU版本
        config_.max_batch_size = 4;
        config_.scale = 1.0f/255.0f;
    }

    void TearDown() override {
        test_image_.release();
    }

    // 辅助函数：检查预处理结果
    bool checkPreprocessedData(const float* data, int height, int width) {
        // 检查是否在[0,1]范围内
        for (int i = 0; i < height * width * 3; ++i) {
            if (data[i] < 0.0f || data[i] > 1.0f) {
                return false;
            }
        }
        return true;
    }

protected:
    cv::Mat test_image_;
    Preprocessor::Config config_;
};

// 测试基本功能
TEST_F(PreprocessorTest, BasicFunctionality) {
    // 创建预处理器
    Preprocessor preprocessor(config_);
    
    // 创建输入图像列表
    std::vector<cv::Mat> images{test_image_};
    
    // 创建输出缓冲区
    const size_t output_size = config_.input_shape[0] * config_.input_shape[1] * 
                              config_.input_shape[2] * sizeof(float);
    detail::CudaBuffer output_buffer(output_size);
    
    // 执行预处理
    EXPECT_TRUE(preprocessor.process(images, output_buffer));
    
    // 检查结果
    float* output_data = static_cast<float*>(output_buffer.hostPtr());
    EXPECT_TRUE(output_data != nullptr);
    EXPECT_TRUE(checkPreprocessedData(output_data, config_.input_shape[1], config_.input_shape[2]));
}

// 测试输入验证
TEST_F(PreprocessorTest, InputValidation) {
    Preprocessor preprocessor(config_);
    detail::CudaBuffer output_buffer(1024);  // 随意大小，反正会失败
    
    // 测试空输入
    {
        std::vector<cv::Mat> empty_images;
        EXPECT_FALSE(preprocessor.process(empty_images, output_buffer));
    }
    
    // 测试超出批大小
    {
        std::vector<cv::Mat> too_many_images(config_.max_batch_size + 1, test_image_);
        EXPECT_FALSE(preprocessor.process(too_many_images, output_buffer));
    }
    
    // 测试无效图像
    {
        cv::Mat empty_image;
        std::vector<cv::Mat> invalid_images{empty_image};
        EXPECT_FALSE(preprocessor.process(invalid_images, output_buffer));
    }
}

// 测试批处理
TEST_F(PreprocessorTest, BatchProcessing) {
    Preprocessor preprocessor(config_);
    
    // 创建多张测试图像
    std::vector<cv::Mat> images;
    for (int i = 0; i < config_.max_batch_size; ++i) {
        images.push_back(test_image_.clone());
    }
    
    // 创建输出缓冲区
    const size_t output_size = config_.max_batch_size * config_.input_shape[0] * 
                              config_.input_shape[1] * config_.input_shape[2] * sizeof(float);
    detail::CudaBuffer output_buffer(output_size);
    
    // 执行批处理
    EXPECT_TRUE(preprocessor.process(images, output_buffer));
    
    // 检查结果
    float* output_data = static_cast<float*>(output_buffer.hostPtr());
    EXPECT_TRUE(output_data != nullptr);
    
    // 检查每个批次的结果
    const size_t single_size = config_.input_shape[0] * config_.input_shape[1] * config_.input_shape[2];
    for (int i = 0; i < config_.max_batch_size; ++i) {
        EXPECT_TRUE(checkPreprocessedData(output_data + i * single_size, 
                                        config_.input_shape[1], 
                                        config_.input_shape[2]));
    }
}

// 测试内存管理
TEST_F(PreprocessorTest, MemoryManagement) {
    // 创建和销毁多个预处理器，检查内存泄漏
    for (int i = 0; i < 10; ++i) {
        Preprocessor preprocessor(config_);
        std::vector<cv::Mat> images{test_image_};
        detail::CudaBuffer output_buffer(config_.input_shape[0] * config_.input_shape[1] * 
                                       config_.input_shape[2] * sizeof(float));
        EXPECT_TRUE(preprocessor.process(images, output_buffer));
    }
}

// 测试错误处理
TEST_F(PreprocessorTest, ErrorHandling) {
    Preprocessor preprocessor(config_);
    
    // 测试无效的输出缓冲区大小
    {
        std::vector<cv::Mat> images{test_image_};
        detail::CudaBuffer small_buffer(1);  // 明显太小的缓冲区
        EXPECT_FALSE(preprocessor.process(images, small_buffer));
    }
    
    // 测试无效的图像格式
    {
        cv::Mat gray_image;
        cv::cvtColor(test_image_, gray_image, cv::COLOR_BGR2GRAY);
        std::vector<cv::Mat> invalid_images{gray_image};
        detail::CudaBuffer output_buffer(config_.input_shape[0] * config_.input_shape[1] * 
                                       config_.input_shape[2] * sizeof(float));
        EXPECT_FALSE(preprocessor.process(invalid_images, output_buffer));
    }
}

// 测试性能
TEST_F(PreprocessorTest, Performance) {
    Preprocessor preprocessor(config_);
    std::vector<cv::Mat> images{test_image_};
    detail::CudaBuffer output_buffer(config_.input_shape[0] * config_.input_shape[1] * 
                                   config_.input_shape[2] * sizeof(float));
    
    // 预热
    preprocessor.process(images, output_buffer);
    
    // 计时测试
    const int num_iterations = 100;
    auto start = std::chrono::high_resolution_clock::now();
    
    for (int i = 0; i < num_iterations; ++i) {
        EXPECT_TRUE(preprocessor.process(images, output_buffer));
    }
    
    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    
    float avg_time = static_cast<float>(duration) / num_iterations;
    std::cout << "Average preprocessing time: " << avg_time << "ms" << std::endl;
    
    // 性能要求：单张图像处理时间不超过10ms
    EXPECT_LT(avg_time, 10.0f);
}

// 测试GPU加速
TEST_F(PreprocessorTest, GPUAcceleration) {
    // 启用CUDA
    config_.use_cuda = true;
    Preprocessor preprocessor(config_);
    
    // 创建输入图像列表
    std::vector<cv::Mat> images{test_image_};
    
    // 创建输出缓冲区
    const size_t output_size = config_.input_shape[0] * config_.input_shape[1] * 
                              config_.input_shape[2] * sizeof(float);
    detail::CudaBuffer output_buffer(output_size);
    
    // 执行GPU预处理
    EXPECT_TRUE(preprocessor.process(images, output_buffer));
    
    // 检查结果
    float* output_data = static_cast<float*>(output_buffer.hostPtr());
    EXPECT_TRUE(output_data != nullptr);
    EXPECT_TRUE(checkPreprocessedData(output_data, config_.input_shape[1], config_.input_shape[2]));
}

// 测试GPU批处理
TEST_F(PreprocessorTest, GPUBatchProcessing) {
    // 启用CUDA
    config_.use_cuda = true;
    Preprocessor preprocessor(config_);
    
    // 创建多张测试图像
    std::vector<cv::Mat> images;
    for (int i = 0; i < config_.max_batch_size; ++i) {
        images.push_back(test_image_.clone());
    }
    
    // 创建输出缓冲区
    const size_t output_size = config_.max_batch_size * config_.input_shape[0] * 
                              config_.input_shape[1] * config_.input_shape[2] * sizeof(float);
    detail::CudaBuffer output_buffer(output_size);
    
    // 执行GPU批处理
    EXPECT_TRUE(preprocessor.process(images, output_buffer));
    
    // 检查结果
    float* output_data = static_cast<float*>(output_buffer.hostPtr());
    EXPECT_TRUE(output_data != nullptr);
    
    // 检查每个批次的结果
    const size_t single_size = config_.input_shape[0] * config_.input_shape[1] * config_.input_shape[2];
    for (int i = 0; i < config_.max_batch_size; ++i) {
        EXPECT_TRUE(checkPreprocessedData(output_data + i * single_size, 
                                        config_.input_shape[1], 
                                        config_.input_shape[2]));
    }
}

// 测试GPU性能
TEST_F(PreprocessorTest, GPUPerformance) {
    // 启用CUDA
    config_.use_cuda = true;
    Preprocessor preprocessor(config_);
    
    // 创建输入数据
    std::vector<cv::Mat> images;
    for (int i = 0; i < config_.max_batch_size; ++i) {
        images.push_back(test_image_.clone());
    }
    
    // 创建输出缓冲区
    const size_t output_size = config_.max_batch_size * config_.input_shape[0] * 
                              config_.input_shape[1] * config_.input_shape[2] * sizeof(float);
    detail::CudaBuffer output_buffer(output_size);
    
    // 预热GPU
    for (int i = 0; i < 5; ++i) {
        preprocessor.process(images, output_buffer);
    }
    
    // 性能测试
    const int num_iterations = 100;
    auto start = std::chrono::high_resolution_clock::now();
    
    for (int i = 0; i < num_iterations; ++i) {
        EXPECT_TRUE(preprocessor.process(images, output_buffer));
    }
    
    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    
    float avg_time = static_cast<float>(duration) / num_iterations;
    float avg_time_per_image = avg_time / config_.max_batch_size;
    
    std::cout << "GPU Performance Results:" << std::endl;
    std::cout << "Total batch size: " << config_.max_batch_size << std::endl;
    std::cout << "Average batch processing time: " << avg_time << "ms" << std::endl;
    std::cout << "Average time per image: " << avg_time_per_image << "ms" << std::endl;
    
    // GPU处理应该比CPU快
    EXPECT_LT(avg_time_per_image, 5.0f);  // 每张图像处理时间应小于5ms
}

// 测试CPU和GPU结果一致性
TEST_F(PreprocessorTest, CPUGPUConsistency) {
    // 创建两个预处理器
    config_.use_cuda = false;
    Preprocessor cpu_preprocessor(config_);
    
    config_.use_cuda = true;
    Preprocessor gpu_preprocessor(config_);
    
    // 创建输入数据
    std::vector<cv::Mat> images{test_image_};
    
    // 创建输出缓冲区
    const size_t output_size = config_.input_shape[0] * config_.input_shape[1] * 
                              config_.input_shape[2] * sizeof(float);
    detail::CudaBuffer cpu_output(output_size);
    detail::CudaBuffer gpu_output(output_size);
    
    // 执行预处理
    EXPECT_TRUE(cpu_preprocessor.process(images, cpu_output));
    EXPECT_TRUE(gpu_preprocessor.process(images, gpu_output));
    
    // 比较结果
    float* cpu_data = static_cast<float*>(cpu_output.hostPtr());
    float* gpu_data = static_cast<float*>(gpu_output.hostPtr());
    
    const float epsilon = 1e-5f;  // 允许的误差范围
    const int total_elements = config_.input_shape[0] * config_.input_shape[1] * config_.input_shape[2];
    
    for (int i = 0; i < total_elements; ++i) {
        EXPECT_NEAR(cpu_data[i], gpu_data[i], epsilon)
            << "Mismatch at index " << i;
    }
}