#include "preprocess.hpp"
#include <opencv2/imgproc.hpp>

namespace pipeline {

// 声明GPU预处理函数
namespace detail {
bool preprocessGPU(const std::vector<cv::Mat>& images,
                  detail::CudaBuffer& output_buffer,
                  const std::vector<std::unique_ptr<detail::CudaBuffer>>& temp_buffers,
                  cudaStream_t stream,
                  float scale);
}

Preprocessor::Preprocessor(const Config& config)
    : config_(config),
      stream_(std::make_unique<detail::CudaStream>()) {
    
    // 预分配临时缓冲区
    if (config_.use_cuda) {
        const size_t single_size = config_.input_shape[0] * config_.input_shape[1] * 
                                 config_.input_shape[2] * sizeof(float);
        for (int i = 0; i < config_.max_batch_size; ++i) {
            temp_buffers_.emplace_back(std::make_unique<detail::CudaBuffer>(single_size));
        }
    }
}

Preprocessor::~Preprocessor() = default;

bool Preprocessor::process(const std::vector<cv::Mat>& images, detail::CudaBuffer& output_buffer) {
    try {
        // 检查输入
        if (!checkInput(images)) {
            Logger::error("Invalid input for preprocessing");
            return false;
        }

        // 根据配置选择处理方式
        return config_.use_cuda ? processGPU(images, output_buffer)
                              : processCPU(images, output_buffer);
    }
    catch (const std::exception& e) {
        Logger::error("Error in preprocessing: " + std::string(e.what()));
        return false;
    }
}

bool Preprocessor::checkInput(const std::vector<cv::Mat>& images) const {
    if (images.empty() || images.size() > static_cast<size_t>(config_.max_batch_size)) {
        return false;
    }

    for (const auto& img : images) {
        if (img.empty() || img.type() != CV_8UC3) {
            return false;
        }
    }

    return true;
}

bool Preprocessor::processCPU(const std::vector<cv::Mat>& images, detail::CudaBuffer& output_buffer) {
    const int batch_size = static_cast<int>(images.size());
    const int target_height = config_.input_shape[1];
    const int target_width = config_.input_shape[2];
    
    // 获取输出缓冲区的主机内存指针
    float* output_ptr = static_cast<float*>(output_buffer.hostPtr());
    if (!output_ptr) {
        Logger::error("Failed to get host pointer from output buffer");
        return false;
    }

    // 处理每张图片
    for (int i = 0; i < batch_size; ++i) {
        float* current_output = output_ptr + i * target_height * target_width * 3;
        if (!preprocessImageCPU(images[i], current_output, target_height, target_width)) {
            Logger::error("Failed to preprocess image " + std::to_string(i));
            return false;
        }
    }

    // 将处理后的数据拷贝到设备内存
    if (!output_buffer.copyH2D(stream_->get())) {
        Logger::error("Failed to copy preprocessed data to device");
        return false;
    }

    return stream_->sync();
}

bool Preprocessor::preprocessImageCPU(const cv::Mat& input, float* output,
                                    int target_height, int target_width) {
    try {
        // 1. 缩放图像
        cv::Mat resized;
        cv::resize(input, resized, cv::Size(target_width, target_height), 0, 0, cv::INTER_LINEAR);

        // 2. BGR转RGB并归一化
        for (int h = 0; h < target_height; ++h) {
            for (int w = 0; w < target_width; ++w) {
                const cv::Vec3b& pixel = resized.at<cv::Vec3b>(h, w);
                const int base_idx = (h * target_width + w) * 3;
                
                // BGR转RGB并归一化到0-1
                output[base_idx + 0] = pixel[2] * config_.scale;  // R
                output[base_idx + 1] = pixel[1] * config_.scale;  // G
                output[base_idx + 2] = pixel[0] * config_.scale;  // B
            }
        }

        return true;
    }
    catch (const std::exception& e) {
        Logger::error("Error in CPU preprocessing: " + std::string(e.what()));
        return false;
    }
}

bool Preprocessor::processGPU(const std::vector<cv::Mat>& images, detail::CudaBuffer& output_buffer) {
    try {
        // 1. 缩放图像到目标尺寸
        std::vector<cv::Mat> resized_images;
        resized_images.reserve(images.size());
        
        for (const auto& img : images) {
            cv::Mat resized;
            cv::resize(img, resized, cv::Size(config_.input_shape[2], config_.input_shape[1]),
                      0, 0, cv::INTER_LINEAR);
            resized_images.push_back(resized);
        }
        
        // 2. 使用CUDA进行BGR转RGB和归一化
        if (!detail::preprocessGPU(resized_images, output_buffer, temp_buffers_,
                                 stream_->get(), config_.scale)) {
            Logger::error("GPU preprocessing failed");
            return false;
        }
        
        return stream_->sync();
    }
    catch (const std::exception& e) {
        Logger::error("Error in GPU preprocessing: " + std::string(e.what()));
        return false;
    }
}

} // namespace pipeline