#include <opencv2/opencv.hpp>
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <cuda_runtime_api.h>
#include <fstream>
#include <iostream>
#include <memory>
#include <vector>
#include <algorithm>
#include <chrono>
#include <string>

using namespace nvinfer1;

// Logger for TensorRT
class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        if (severity != Severity::kINFO) {
            std::cerr << msg << std::endl;
        }
    }
} gLogger;

// Save TensorRT engine to file
bool saveEngine(nvinfer1::ICudaEngine* engine, const std::string& engineFile) {
    nvinfer1::IHostMemory* serializedEngine = engine->serialize();
    if (!serializedEngine) {
        std::cerr << "Error: Failed to serialize engine." << std::endl;
        return false;
    }
    std::ofstream ofs(engineFile, std::ios::binary);
    if (!ofs) {
        std::cerr << "Error: Cannot open engine file for writing: " << engineFile << std::endl;
        delete serializedEngine;
        return false;
    }
    ofs.write(reinterpret_cast<const char*>(serializedEngine->data()), serializedEngine->size());
    ofs.close();
    delete serializedEngine;
    return true;
}

// Load TensorRT engine from file
nvinfer1::ICudaEngine* loadEngine(const std::string& engineFile, nvinfer1::ILogger& logger) {
    std::ifstream ifs(engineFile, std::ios::binary);
    if (!ifs) {
        std::cerr << "Error: Cannot open engine file for reading: " << engineFile << std::endl;
        return nullptr;
    }
    ifs.seekg(0, std::ifstream::end);
    size_t size = ifs.tellg();
    ifs.seekg(0, std::ifstream::beg);
    std::vector<char> engineData(size);
    ifs.read(engineData.data(), size);
    ifs.close();

    nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), size, nullptr);
    delete runtime;
    return engine;
}

// Build TensorRT engine from ONNX
ICudaEngine* buildEngine(const std::string& onnxFile, int maxBatchSize) {
    IBuilder* builder = createInferBuilder(gLogger);
    const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);

    std::ifstream file(onnxFile, std::ios::binary);
    if (!file.good()) {
        std::cerr << "Error: Cannot open ONNX file: " << onnxFile << std::endl;
        return nullptr;
    }
    file.seekg(0, file.end);
    size_t size = file.tellg();
    file.seekg(0, file.beg);
    std::vector<char> modelData(size);
    file.read(modelData.data(), size);
    file.close();

    if (!parser->parse(modelData.data(), size)) {
        std::cerr << "Error: Failed to parse ONNX model" << std::endl;
        return nullptr;
    }

    IBuilderConfig* config = builder->createBuilderConfig();
    config->setMaxWorkspaceSize(1ULL << 30); // 1GB
    if (builder->platformHasFastFp16()) {
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
    }
    builder->setMaxBatchSize(maxBatchSize);
    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);

    delete parser;
    delete network;
    delete config;
    delete builder;

    return engine;
}

// Bounding box structure
struct BBox {
    float x1, y1, x2, y2, score;
    int classId;
};

// Non-Maximum Suppression
std::vector<BBox> nms(const std::vector<BBox>& boxes, float iouThreshold, float scoreThreshold) {
    std::vector<BBox> filteredBoxes;
    std::vector<BBox> sortedBoxes = boxes;
    std::sort(sortedBoxes.begin(), sortedBoxes.end(), [](const BBox& a, const BBox& b) {
        return a.score > b.score;
    });

    while (!sortedBoxes.empty()) {
        BBox best = sortedBoxes.front();
        if (best.score >= scoreThreshold) {
            filteredBoxes.push_back(best);
        }
        sortedBoxes.erase(sortedBoxes.begin());
        sortedBoxes.erase(std::remove_if(sortedBoxes.begin(), sortedBoxes.end(), [&](const BBox& b) {
            float x1 = std::max(best.x1, b.x1);
            float y1 = std::max(best.y1, b.y1);
            float x2 = std::min(best.x2, b.x2);
            float y2 = std::min(best.y2, b.y2);
            float w = std::max(0.0f, x2 - x1);
            float h = std::max(0.0f, y2 - y1);
            float inter = w * h;
            float area1 = (best.x2 - best.x1) * (best.y2 - best.y1);
            float area2 = (b.x2 - b.x1) * (b.y2 - b.y1);
            float iou = inter / (area1 + area2 - inter);
            return iou > iouThreshold;
        }), sortedBoxes.end());
    }
    return filteredBoxes;
}

// Inference with CUDA buffers
void doInference(IExecutionContext* context, float* input, float* output, int batchSize, int inputSize, int outputSize, void** buffers) {
    auto start = std::chrono::high_resolution_clock::now();
    cudaMemcpy(buffers[0], input, batchSize * inputSize * sizeof(float), cudaMemcpyHostToDevice);
    auto memcpy_h2d_end = std::chrono::high_resolution_clock::now();
    context->executeV2(buffers);
    auto inference_end = std::chrono::high_resolution_clock::now();
    cudaMemcpy(output, buffers[1], batchSize * outputSize * sizeof(float), cudaMemcpyDeviceToHost);
    auto memcpy_d2h_end = std::chrono::high_resolution_clock::now();
    auto h2d_time = std::chrono::duration<double>(memcpy_h2d_end - start).count() * 1000;
    auto inference_time = std::chrono::duration<double>(inference_end - memcpy_h2d_end).count() * 1000;
    auto d2h_time = std::chrono::duration<double>(memcpy_d2h_end - inference_end).count() * 1000;
    std::cout << "[Inference Timing] H2D: " << h2d_time << "ms, Inference: " << inference_time << "ms, D2H: " << d2h_time << "ms" << std::endl;
}

// Preprocess input image for YOLO (resize, BGR->RGB, normalize, CHW)
void preprocessImage(const cv::Mat& img, float* input, int width, int height) {
    cv::Mat resized, rgb;
    cv::resize(img, resized, cv::Size(width, height));
    cv::cvtColor(resized, rgb, cv::COLOR_BGR2RGB);
    cv::Mat floatImg;
    rgb.convertTo(floatImg, CV_32FC3, 1.0 / 255.0);
    std::vector<cv::Mat> channels(3);
    cv::split(floatImg, channels);
    int index = 0;
    for (int c = 0; c < 3; ++c) {
        for (int h = 0; h < height; ++h) {
            for (int w = 0; w < width; ++w) {
                input[index++] = channels[c].at<float>(h, w);
            }
        }
    }
}

// Post-process YOLO output
std::vector<BBox> postProcess(float* output, int numBoxes, float confThreshold, int origWidth, int origHeight, int modelWidth, int modelHeight) {
    std::vector<BBox> boxes;
    float scaleX = (float)origWidth / (float)modelWidth;
    float scaleY = (float)origHeight / (float)modelHeight;
    for (int i = 0; i < numBoxes; ++i) {
        float x_center = output[0 * numBoxes + i];
        float y_center = output[1 * numBoxes + i];
        float w = output[2 * numBoxes + i];
        float h = output[3 * numBoxes + i];
        float conf = output[4 * numBoxes + i];
        if (conf < confThreshold) continue;
        float x1 = (x_center - w / 2.0f) * scaleX;
        float y1 = (y_center - h / 2.0f) * scaleY;
        float x2 = (x_center + w / 2.0f) * scaleX;
        float y2 = (y_center + h / 2.0f) * scaleY;
        x1 = std::max(0.0f, std::min(x1, (float)origWidth - 1));
        y1 = std::max(0.0f, std::min(y1, (float)origHeight - 1));
        x2 = std::max(0.0f, std::min(x2, (float)origWidth - 1));
        y2 = std::max(0.0f, std::min(y2, (float)origHeight - 1));
        boxes.push_back({x1, y1, x2, y2, conf, 0});
    }
    return nms(boxes, 0.45f, confThreshold);
}

int main(int argc, char** argv) {
    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " <onnx_model_path> <video_path or camera>" << std::endl;
        std::cerr << "Example for camera: " << argv[0] << " best.onnx /dev/video0" << std::endl;
        std::cerr << "Example for camera: " << argv[0] << " best.onnx 0" << std::endl;
        std::cerr << "Example for video:  " << argv[0] << " best.onnx sample.mp4" << std::endl;
        return -1;
    }

    std::string onnxFile = argv[1];
    std::string videoPath = argv[2];
    std::string engineFile = onnxFile + ".engine";

    nvinfer1::ICudaEngine* engine = nullptr;
    std::ifstream f(engineFile);
    if (f.good()) {
        std::cout << "Loading cached TensorRT engine from " << engineFile << std::endl;
        engine = loadEngine(engineFile, gLogger);
    } else {
        std::cout << "Building TensorRT engine from ONNX and caching to " << engineFile << std::endl;
        engine = buildEngine(onnxFile, 1);
        if (engine) saveEngine(engine, engineFile);
    }
    if (!engine) {
        std::cerr << "Failed to load or build TensorRT engine" << std::endl;
        return -1;
    }

    IExecutionContext* context = engine->createExecutionContext();
    if (!context) {
        std::cerr << "Failed to create execution context" << std::endl;
        delete engine;
        return -1;
    }

    int inputIdx = engine->getBindingIndex("images");
    int outputIdx = engine->getBindingIndex("output0");
    if (inputIdx == -1 || outputIdx == -1) {
        std::cerr << "Error: Could not find input/output bindings. Check tensor names." << std::endl;
        for (int i = 0; i < engine->getNbBindings(); ++i) {
            std::cout << "  [" << i << "] " << engine->getBindingName(i) << std::endl;
        }
        delete context;
        delete engine;
        return -1;
    }
    Dims inputDims = engine->getBindingDimensions(inputIdx);
    Dims outputDims = engine->getBindingDimensions(outputIdx);

    int inputSize = inputDims.d[1] * inputDims.d[2] * inputDims.d[3];
    int numBoxes = outputDims.d[2];
    int outputSize = outputDims.d[1] * outputDims.d[2];
    int width = inputDims.d[3];
    int height = inputDims.d[2];

    std::cout << "[Buffer Info] Input size: " << inputSize << ", Output size: " << outputSize << ", Num boxes: " << numBoxes << std::endl;

    void* buffers[2];
    cudaMalloc(&buffers[0], inputSize * sizeof(float));
    cudaMalloc(&buffers[1], outputSize * sizeof(float));

    // VideoCapture: support for file, /dev/video*, or camera index
    cv::VideoCapture cap;
    if (videoPath == "0") {
        cap.open(0);
    } else if (videoPath.find("/dev/video") == 0) {
        cap.open(videoPath, cv::CAP_V4L2);
    } else {
        cap.open(videoPath);
    }
    if (!cap.isOpened()) {
        std::cerr << "Error: Cannot open video or camera: " << videoPath << std::endl;
        delete context;
        delete engine;
        cudaFree(buffers[0]);
        cudaFree(buffers[1]);
        return -1;
    }

    cv::namedWindow("YOLOv11n Inference", cv::WINDOW_NORMAL);
    cv::resizeWindow("YOLOv11n Inference", 960, 540);

    std::vector<float> input(inputSize);
    std::vector<float> output(outputSize);

    cv::Mat frame;
    while (cap.read(frame)) {
        preprocessImage(frame, input.data(), width, height);
        doInference(context, input.data(), output.data(), 1, inputSize, outputSize, buffers);
        std::vector<BBox> boxes = postProcess(output.data(), numBoxes, 0.3f, frame.cols, frame.rows, width, height);

        for (const auto& box : boxes) {
            cv::rectangle(frame, cv::Point((int)box.x1, (int)box.y1), cv::Point((int)box.x2, (int)box.y2), cv::Scalar(0, 255, 0), 2);
            std::string label = "Object: " + std::to_string(box.score).substr(0, 4);
            cv::putText(frame, label, cv::Point((int)box.x1, (int)box.y1 - 10), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 2);
        }

        cv::imshow("YOLOv11n Inference", frame);
        if (cv::waitKey(1) == 27) break; // ESC to exit
    }

    cap.release();
    cv::destroyAllWindows();
    delete context;
    delete engine;
    cudaFree(buffers[0]);
    cudaFree(buffers[1]);
    return 0;
}