#include "trt_inference.hpp"
#include "cuda_helper.hpp"
#include <fstream>
#include <NvOnnxParser.h>
#include <iostream>
#include <filesystem>
#include "../common/logger.hpp"
#include <opencv2/imgproc.hpp>

namespace pipeline {

TrtInference::TrtInference(const InferenceConfig& config)
    : config_(config), 
      logger_(std::make_unique<detail::Logger>()),
      stream_(std::make_unique<detail::CudaStream>()) {
    
    // 创建预处理器配置
    Preprocessor::Config preprocess_config;
    preprocess_config.input_shape = config.model.input_shape;
    preprocess_config.use_cuda = true;  // 默认使用CUDA加速
    preprocess_config.max_batch_size = config.max_batch_size;
    preprocess_config.scale = 1.0f/255.0f;  // 归一化系数
    
    // 创建预处理器实例
    preprocessor_ = std::make_unique<Preprocessor>(preprocess_config);
}

TrtInference::~TrtInference() {
    destroy();
}

bool TrtInference::convertOnnxToEngine(const std::string& onnx_path, const std::string& engine_path) {
    try {
        Logger::info("=== Starting ONNX to TensorRT Conversion ===");
        Logger::info("ONNX file: " + onnx_path);
        
        // 检查文件是否存在
        if (!std::filesystem::exists(onnx_path)) {
            Logger::error("ONNX file does not exist: " + onnx_path);
            return false;
        }

        // 创建 builder
        std::unique_ptr<nvinfer1::IBuilder, void(*)(nvinfer1::IBuilder*)> builder(
            nvinfer1::createInferBuilder(*logger_),
            [](nvinfer1::IBuilder* b) { delete b; });
        if (!builder) {
            Logger::error("Failed to create TensorRT Builder");
            return false;
        }

        // 创建网络定义
        #pragma GCC diagnostic push
        #pragma GCC diagnostic ignored "-Wdeprecated-declarations"
        const auto explicit_batch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
        #pragma GCC diagnostic pop
        std::unique_ptr<nvinfer1::INetworkDefinition, void(*)(nvinfer1::INetworkDefinition*)> network(
            builder->createNetworkV2(explicit_batch),
            [](nvinfer1::INetworkDefinition* n) { delete n; });
        if (!network) {
            Logger::error("Failed to create network definition");
            return false;
        }

        // 创建 ONNX 解析器
        std::unique_ptr<nvonnxparser::IParser, void(*)(nvonnxparser::IParser*)> parser(
            nvonnxparser::createParser(*network, *logger_),
            [](nvonnxparser::IParser* p) { delete p; });
        if (!parser) {
            Logger::error("Failed to create ONNX parser");
            return false;
        }

        // 解析 ONNX 文件
        if (!parser->parseFromFile(onnx_path.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kINFO))) {
            Logger::error("Failed to parse ONNX file");
            return false;
        }

        // 创建构建配置
        std::unique_ptr<nvinfer1::IBuilderConfig, void(*)(nvinfer1::IBuilderConfig*)> config(
            builder->createBuilderConfig(),
            [](nvinfer1::IBuilderConfig* c) { delete c; });
        if (!config) {
            Logger::error("Failed to create builder config");
            return false;
        }

        // 修改优化配置部分
        auto profile = builder->createOptimizationProfile();
        auto input = network->getInput(0);
        if (!input) {
            Logger::error("Failed to get network input");
            return false;
        }

        // 获取输入名称和维度
        const char* input_name = input->getName();
        nvinfer1::Dims dims = input->getDimensions();

        // 确保所有维度都是正数
        for (int i = 0; i < dims.nbDims; i++) {
            if (dims.d[i] < 0) {
                if (i == 0) {
                    // batch维度
                    dims.d[i] = 1;
                } else {
                    // 其他维度使用配置中的值
                    dims.d[i] = config_.model.input_shape[i-1];
                }
            }
        }

        // 设置最小、最优和最大维度
        nvinfer1::Dims min_dims = dims;
        nvinfer1::Dims opt_dims = dims;
        nvinfer1::Dims max_dims = dims;

        // 只修改batch维度
        min_dims.d[0] = 1;  // 最小batch size
        opt_dims.d[0] = 1;  // 最优batch size
        max_dims.d[0] = config_.max_batch_size;  // 最大batch size

        // 设置优化配置
        profile->setDimensions(input_name, nvinfer1::OptProfileSelector::kMIN, min_dims);
        profile->setDimensions(input_name, nvinfer1::OptProfileSelector::kOPT, opt_dims);
        profile->setDimensions(input_name, nvinfer1::OptProfileSelector::kMAX, max_dims);

        if (!profile->isValid()) {
            Logger::error("Invalid optimization profile");
            return false;
        }

        config->addOptimizationProfile(profile);

        // 设置工作空间大小
        config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, config_.workspace_size);

        // 设置精度
        #pragma GCC diagnostic push
        #pragma GCC diagnostic ignored "-Wdeprecated-declarations"
        if (builder->platformHasFastFp16()) {
            Logger::info("Enabling FP16 precision");
            config->setFlag(nvinfer1::BuilderFlag::kFP16);
        }
        #pragma GCC diagnostic pop

        // 构建引擎
        std::unique_ptr<nvinfer1::IHostMemory, void(*)(nvinfer1::IHostMemory*)> serializedEngine(
            builder->buildSerializedNetwork(*network, *config),
            [](nvinfer1::IHostMemory* m) { delete m; });
        if (!serializedEngine) {
            Logger::error("Failed to build TensorRT engine");
            return false;
        }

        // 保存引擎文件
        std::ofstream engineFile(engine_path, std::ios::binary);
        if (!engineFile) {
            Logger::error("Failed to open engine file for writing");
            return false;
        }

        engineFile.write(static_cast<const char*>(serializedEngine->data()), serializedEngine->size());
        engineFile.close();

        Logger::info("Successfully saved TensorRT engine to: " + engine_path);
        return true;
    }
    catch (const std::exception& e) {
        Logger::error("Error in convertOnnxToEngine: " + std::string(e.what()));
        return false;
    }
}

bool TrtInference::loadEngine() {
    try {
        // 确保CUDA上下文初始化
        cudaFree(0);  // 强制初始化CUDA上下文
        
        // 检查文件是否存在
        if (!std::filesystem::exists(config_.model.engine_path)) {
            Logger::error("Engine file not found: " + config_.model.engine_path);
            return false;
        }

        // 读取引擎文件
        std::ifstream file(config_.model.engine_path, std::ios::binary);
        if (!file) {
            Logger::error("Failed to open engine file");
            return false;
        }

        file.seekg(0, std::ios::end);
        size_t size = file.tellg();
        file.seekg(0, std::ios::beg);

        std::vector<char> engineData(size);
        file.read(engineData.data(), size);
        file.close();

        // 创建 runtime
        runtime_ = nvinfer1::createInferRuntime(*logger_);
        if (!runtime_) {
            Logger::error("Failed to create TensorRT Runtime");
            return false;
        }

        // 反序列化引擎
        engine_ = runtime_->deserializeCudaEngine(engineData.data(), size);
        if (!engine_) {
            Logger::error("Failed to deserialize TensorRT engine");
            return false;
        }

        // 创建执行上下文
        if (!createContext()) {
            Logger::error("Failed to create execution context");
            return false;
        }

        // 分配缓冲区
        if (!allocateBuffers()) {
            Logger::error("Failed to allocate CUDA buffers");
            return false;
        }

        // 同步设备
        cudaDeviceSynchronize();

        Logger::info("Successfully loaded TensorRT engine");
        return true;
    }
    catch (const std::exception& e) {
        Logger::error("Error in loadEngine: " + std::string(e.what()));
        return false;
    }
}

bool TrtInference::infer(const std::vector<cv::Mat>& images, std::vector<DetectionResult>& results) {
    if (!isLoaded()) {
        Logger::error("Engine not loaded");
        return false;
    }

    try {
        // 确保CUDA上下文有效
        cudaFree(0);
        
        // 1. 预处理
        if (!preprocess(images)) {
            return false;
        }

        // 2. 设置输入形状
        if (!context_->setInputShape("images", 
            nvinfer1::Dims4(images.size(), config_.model.input_shape[0],
                           config_.model.input_shape[1], config_.model.input_shape[2]))) {
            Logger::error("Failed to set input shape");
            return false;
        }

        // 3. 设置输入和输出张量地址
        if (!context_->setTensorAddress("images", input_buffers_[0]->devicePtr())) {
            Logger::error("Failed to set input tensor address");
            return false;
        }
        
        const char* output_name = engine_->getIOTensorName(engine_->getNbIOTensors() - 1);
        if (!context_->setTensorAddress(output_name, output_buffers_[0]->devicePtr())) {
            Logger::error("Failed to set output tensor address");
            return false;
        }

        // 4. 执行推理
        if (!context_->enqueueV3(stream_->get())) {
            Logger::error("Inference failed");
            return false;
        }

        // 5. 后处理
        if (!postprocess(results)) {
            return false;
        }

        // 6. 同步流和设备
        if (!stream_->sync()) {
            Logger::error("Failed to synchronize CUDA stream");
            return false;
        }
        cudaDeviceSynchronize();

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

bool TrtInference::preprocess(const std::vector<cv::Mat>& images) {
    try {
        // 检查输入
        if (images.empty() || images.size() > static_cast<size_t>(config_.max_batch_size)) {
            Logger::error("Invalid batch size: " + std::to_string(images.size()));
            return false;
        }

        // 使用预处理器处理图像
        if (!preprocessor_->process(images, *input_buffers_[0])) {
            Logger::error("Preprocessing failed");
            return false;
        }

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

bool TrtInference::postprocess(std::vector<DetectionResult>& results) {
    try {
        // 获取输出尺寸
        const auto output_dims = engine_->getTensorShape(output_names_[0].c_str());
        const int num_outputs = output_dims.d[1];  // 每个检测的输出维度
        const int max_detections = output_dims.d[2];  // 最大检测数量
        
        // 分配主机内存接收结果
        const size_t output_size = output_buffers_[0]->size();
        std::vector<float> output_data(output_size / sizeof(float));
        
        // 从GPU拷贝结果
        if (cudaMemcpyAsync(output_data.data(), output_buffers_[0]->devicePtr(),
                           output_size, cudaMemcpyDeviceToHost, stream_->get()) != cudaSuccess) {
            Logger::error("Failed to copy output data from GPU");
            return false;
        }
        
        // 等待拷贝完成
        if (cudaStreamSynchronize(stream_->get()) != cudaSuccess) {
            Logger::error("Failed to synchronize CUDA stream");
            return false;
        }

        // 解析结果
        results.clear();
        results.resize(1);  // 每批次的结果
        auto& result = results[0];
        
        // YOLOv8输出格式：[batch_id, x, y, w, h, conf, class_conf, class_id]
        for (int i = 0; i < max_detections; ++i) {
            const float* det = output_data.data() + i * num_outputs;
            const float conf = det[5];
            
            if (conf < config_.threshold.conf) continue;
            
            BBox box;
            box.x1 = det[1] - det[3] / 2.0f;  // center_x - width/2
            box.y1 = det[2] - det[4] / 2.0f;  // center_y - height/2
            box.x2 = det[1] + det[3] / 2.0f;  // center_x + width/2
            box.y2 = det[2] + det[4] / 2.0f;  // center_y + height/2
            box.score = conf;                  // 使用 score 而不是 conf
            box.class_id = static_cast<int>(det[7]);
            
            result.boxes.push_back(box);
        }

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

void TrtInference::doNMS(std::vector<BBox>& boxes, float nms_thresh) {
    // 按置信度排序，修改 lambda 表达式确保返回 bool 值
    std::sort(boxes.begin(), boxes.end(), 
              [](const BBox& a, const BBox& b) -> bool { 
                  return a.score > b.score;    // 使用 score 而不是 conf
              });
    
    std::vector<bool> keep(boxes.size(), true);
    
    // NMS
    for (size_t i = 0; i < boxes.size(); ++i) {
        if (!keep[i]) continue;
        
        for (size_t j = i + 1; j < boxes.size(); ++j) {
            if (!keep[j]) continue;
            
            // 如果是同一类别且IOU大于阈值，则抑制
            if (boxes[i].class_id == boxes[j].class_id && 
                calculateIOU(boxes[i], boxes[j]) > nms_thresh) {
                keep[j] = false;
            }
        }
    }
    
    // 移除被抑制的框
    std::vector<BBox> nms_boxes;
    for (size_t i = 0; i < boxes.size(); ++i) {
        if (keep[i]) {
            nms_boxes.push_back(boxes[i]);
        }
    }
    
    boxes.swap(nms_boxes);
}

float TrtInference::calculateIOU(const BBox& box1, const BBox& box2) {
    float x1 = std::max(box1.x1, box2.x1);
    float y1 = std::max(box1.y1, box2.y1);
    float x2 = std::min(box1.x2, box2.x2);
    float y2 = std::min(box1.y2, box2.y2);

    float intersection = std::max(0.0f, x2 - x1) * std::max(0.0f, y2 - y1);
    float area1 = (box1.x2 - box1.x1) * (box1.y2 - box1.y1);
    float area2 = (box2.x2 - box2.x1) * (box2.y2 - box2.y1);
    
    return intersection / (area1 + area2 - intersection);
}

bool TrtInference::createContext() {
    if (!engine_) {
        return false;
    }

    context_ = engine_->createExecutionContext();
    if (!context_) {
        return false;
    }

    // 设置优化的推理维度
    if (config_.max_batch_size > 0) {
        // TensorRT 8.x 不再使用 setOptimizationProfile
        context_->setOptimizationProfileAsync(0, stream_->get());
    }

    return true;
}

bool TrtInference::allocateBuffers() {
    if (!engine_ || !context_) {
        return false;
    }

    try {
        // 先释放旧的缓冲区
        input_buffers_.clear();
        output_buffers_.clear();
        bindings_.clear();
        output_names_.clear();

        // 获取网络的输入输出数量
        int numInputs = 0;
        int numOutputs = 0;
        int maxBindings = engine_->getNbIOTensors();
        bindings_.resize(maxBindings);
        
        // 分配输入输出缓冲区
        for (int i = 0; i < maxBindings; i++) {
            const char* name = engine_->getIOTensorName(i);
            nvinfer1::TensorIOMode mode = engine_->getTensorIOMode(name);
            if (mode == nvinfer1::TensorIOMode::kINPUT) {
                numInputs++;
                nvinfer1::Dims dims = engine_->getTensorShape(name);
                nvinfer1::DataType dtype = engine_->getTensorDataType(name);
                
                // 计算缓冲区大小
                size_t size = 1;
                for (int j = 0; j < dims.nbDims; j++) {
                    if (dims.d[j] < 0) {
                        dims.d[j] = (j == 0) ? 1 : config_.model.input_shape[j-1];
                    }
                    size *= dims.d[j];
                }
                size *= (dtype == nvinfer1::DataType::kFLOAT ? 4 : 2);  // FP32或FP16
                
                // 使用固定内存分配
                void* d_ptr = nullptr;
                cudaMalloc(&d_ptr, size);
                if (!d_ptr) {
                    Logger::error("Failed to allocate input buffer");
                    return false;
                }
                
                input_buffers_.emplace_back(std::make_unique<detail::CudaBuffer>(size));
                bindings_[i] = d_ptr;
                
            } else if (mode == nvinfer1::TensorIOMode::kOUTPUT) {
                numOutputs++;
                nvinfer1::Dims dims = engine_->getTensorShape(name);
                nvinfer1::DataType dtype = engine_->getTensorDataType(name);
                
                // 计算缓冲区大小
                size_t size = 1;
                for (int j = 0; j < dims.nbDims; j++) {
                    if (dims.d[j] < 0) {
                        dims.d[j] = (j == 0) ? 1 : config_.model.input_shape[j-1];
                    }
                    size *= dims.d[j];
                }
                size *= (dtype == nvinfer1::DataType::kFLOAT ? 4 : 2);  // FP32或FP16
                
                // 使用固定内存分配
                void* d_ptr = nullptr;
                cudaMalloc(&d_ptr, size);
                if (!d_ptr) {
                    Logger::error("Failed to allocate output buffer");
                    return false;
                }
                
                output_buffers_.emplace_back(std::make_unique<detail::CudaBuffer>(size));
                bindings_[i] = d_ptr;
                output_names_.push_back(name);
            }
        }

        // 同步设备
        cudaDeviceSynchronize();
        
        return (numInputs > 0 && numOutputs > 0);
    }
    catch (const std::exception& e) {
        Logger::error("Error in allocateBuffers: " + std::string(e.what()));
        return false;
    }
}

void TrtInference::destroy() {
    try {
        // 先同步CUDA流和设备
        if (stream_) {
            stream_->sync();
        }
        cudaDeviceSynchronize();
        
        // 释放绑定的设备内存
        for (void* ptr : bindings_) {
            if (ptr) {
                cudaFree(ptr);
            }
        }
        
        // 释放缓冲区
        input_buffers_.clear();
        output_buffers_.clear();
        bindings_.clear();
        output_names_.clear();
        
        // 释放TensorRT资源
        if (context_) {
            delete context_;  // 使用delete而不是destroy()
            context_ = nullptr;
        }
        if (engine_) {
            delete engine_;   // 使用delete而不是destroy()
            engine_ = nullptr;
        }
        if (runtime_) {
            delete runtime_;  // 使用delete而不是destroy()
            runtime_ = nullptr;
        }
        
        // 最后再次同步设备
        cudaDeviceSynchronize();
    }
    catch (const std::exception& e) {
        Logger::error("Error in destroy: " + std::string(e.what()));
    }
}

} // namespace pipeline