#include "cuda_helper.hpp"
#include <iostream>

namespace pipeline {
namespace detail {

CudaStream::CudaStream() {
    cudaError_t err = cudaStreamCreate(&stream_);
    if (err != cudaSuccess) {
        std::cerr << "Failed to create CUDA stream: " << cudaGetErrorString(err) << std::endl;
        stream_ = nullptr;
    }
}

CudaStream::~CudaStream() {
    if (stream_) {
        cudaStreamDestroy(stream_);
    }
}

bool CudaStream::sync() const {
    if (!stream_) return false;
    CUDA_CHECK(cudaStreamSynchronize(stream_));
    return true;
}

CudaBuffer::CudaBuffer(size_t size) : size_(size) {
    // 分配设备内存
    cudaError_t err = cudaMalloc(&d_ptr_, size);
    if (err != cudaSuccess) {
        std::cerr << "Failed to allocate device memory: " << cudaGetErrorString(err) << std::endl;
        d_ptr_ = nullptr;
        return;
    }

    // 分配主机内存（页锁定）
    err = cudaMallocHost(&h_ptr_, size);
    if (err != cudaSuccess) {
        std::cerr << "Failed to allocate host memory: " << cudaGetErrorString(err) << std::endl;
        cudaFree(d_ptr_);
        d_ptr_ = nullptr;
        h_ptr_ = nullptr;
        return;
    }
}

CudaBuffer::~CudaBuffer() {
    if (d_ptr_) {
        cudaFree(d_ptr_);
        d_ptr_ = nullptr;
    }
    if (h_ptr_) {
        cudaFreeHost(h_ptr_);
        h_ptr_ = nullptr;
    }
}

bool CudaBuffer::copyH2D(cudaStream_t stream) {
    if (!d_ptr_ || !h_ptr_) return false;
    
    // 执行拷贝
    cudaError_t err = cudaMemcpyAsync(d_ptr_, h_ptr_, size_, 
        cudaMemcpyHostToDevice, stream);
    if (err != cudaSuccess) {
        std::cerr << "H2D copy failed: " << cudaGetErrorString(err) << std::endl;
        return false;
    }
    
    // 如果是同步模式，等待拷贝完成
    if (stream == nullptr) {
        err = cudaDeviceSynchronize();
        if (err != cudaSuccess) {
            std::cerr << "H2D sync failed: " << cudaGetErrorString(err) << std::endl;
            return false;
        }
    }
    
    return true;
}

bool CudaBuffer::copyD2H(cudaStream_t stream) {
    if (!d_ptr_ || !h_ptr_) return false;
    
    // 执行拷贝
    cudaError_t err = cudaMemcpyAsync(h_ptr_, d_ptr_, size_, 
        cudaMemcpyDeviceToHost, stream);
    if (err != cudaSuccess) {
        std::cerr << "D2H copy failed: " << cudaGetErrorString(err) << std::endl;
        return false;
    }
    
    // 如果是同步模式，等待拷贝完成
    if (stream == nullptr) {
        err = cudaDeviceSynchronize();
        if (err != cudaSuccess) {
            std::cerr << "D2H sync failed: " << cudaGetErrorString(err) << std::endl;
            return false;
        }
    }
    
    return true;
}

} // namespace detail
} // namespace pipeline