// opencl_cache.cpp
#include <CL/opencl.hpp>        // OpenCL C++ bindings
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>
#include <filesystem>       // C++17
#include <stdexcept>

// -------------------------------------------------------------------
// Helpers
// -------------------------------------------------------------------

// Read a text file into a string
std::string read_file(const std::string &path)
{
    std::ifstream f(path, std::ios::in | std::ios::binary);
    if (!f) throw std::runtime_error("Cannot open file: " + path);
    std::ostringstream ss;
    ss << f.rdbuf();
    return ss.str();
}

// Read a binary file into a std::vector<uint8_t>
std::vector<uint8_t> read_binary(const std::string &path)
{
    std::ifstream f(path, std::ios::binary);
    if (!f) throw std::runtime_error("Cannot open binary file: " + path);
    return std::vector<uint8_t>(std::istreambuf_iterator<char>(f),
                                std::istreambuf_iterator<char>());
}

// Write a binary buffer to file
void write_binary(const std::string &path, const std::vector<uint8_t> &data)
{
    std::ofstream f(path, std::ios::binary);
    if (!f) throw std::runtime_error("Cannot write binary file: " + path);
    f.write(reinterpret_cast<const char*>(data.data()), data.size());
}

// Create a unique key for the device (vendor + name + global mem size)
std::string device_key(const cl::Device &dev)
{
    std::ostringstream ss;
    ss << dev.getInfo<CL_DEVICE_VENDOR>()
       << "-" << dev.getInfo<CL_DEVICE_NAME>()
       << "-" << dev.getInfo<CL_DEVICE_GLOBAL_MEM_SIZE>();
    return ss.str();
}

// -------------------------------------------------------------------
// Main
// -------------------------------------------------------------------
int main()
{
    try
    {
        // 1) Pick the first GPU device on the first platform
        std::vector<cl::Platform> platforms;
        cl::Platform::get(&platforms);
        if (platforms.empty())
            throw std::runtime_error("No OpenCL platforms found");

        cl::Platform platform = platforms[0];
        std::vector<cl::Device> devices;
        platform.getDevices(CL_DEVICE_TYPE_GPU, &devices);
        if (devices.empty())
            throw std::runtime_error("No GPU devices on the platform");

        cl::Device device = devices[0];
        std::cout << "Using platform: "
                  << platform.getInfo<CL_PLATFORM_NAME>() << '\n';
        std::cout << "Using device  : "
                  << device.getInfo<CL_DEVICE_NAME>() << '\n';

        // 2) Context & cache file name
        cl::Context ctx(device);
        std::string bin_file = device_key(device) + ".bin";
        std::filesystem::path bin_path = std::filesystem::current_path() / bin_file;

        cl::Program program;

        // 3) If we already have a binary, load it
        if (std::filesystem::exists(bin_path))
        {
            std::cout << "Loading binary from " << bin_path << '\n';
            std::vector<uint8_t> bin = read_binary(bin_path.string());

            // OpenCL expects an array of *unsigned char* pointers, one per
            // device.  We are using only one device here.
            std::vector<const unsigned char*> binaries(1, bin.data());
            std::vector<size_t> lengths(1, bin.size());

            std::vector<cl::Program> progs(1);
            cl_int err = ctx.createProgramWithBinary(devices,
                                                    lengths,
                                                    binaries.data(),
                                                    nullptr,
                                                    &progs[0]);

            if (err != CL_SUCCESS)
                throw std::runtime_error("clCreateProgramWithBinary failed");

            program = progs[0];
        }
        else
        {
            // 4) Compile from source, extract the binary, save it

            std::cout << "Compiling source (first run)...\n";

            std::string src = read_file("example.cl");
            cl::Program::Sources srcs(1, { src.c_str(), src.length() + 1 });

            program = cl::Program(ctx, srcs);
            cl_int err = program.build(devices);
            if (err != CL_SUCCESS)
            {
                std::string log = program.getBuildInfo<CL_PROGRAM_BUILD_LOG>(device);
                throw std::runtime_error("Program build failed: " + log);
            }

            // Retrieve the binary
            std::vector<unsigned char> bin;
            std::vector<size_t> lengths;
            program.getInfo(CL_PROGRAM_BINARY_SIZES, &lengths);
            program.getInfo(CL_PROGRAM_BINARIES, &bin);

            // Save it for next run
            write_binary(bin_path.string(), std::vector<uint8_t>(bin.begin(), bin.end()));
            std::cout << "Saved binary to " << bin_path << '\n';
        }

        // -------------------------------------------------------------------
        // 5) Demo: add two float vectors
        // -------------------------------------------------------------------
        size_t N = 1024;
        std::vector<float> a(N, 1.0f), b(N, 2.0f), c(N, 0.0f);

        // Allocate device buffers
        cl::Buffer bufA(ctx, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                        sizeof(float) * N, a.data());
        cl::Buffer bufB(ctx, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR,
                        sizeof(float) * N, b.data());
        cl::Buffer bufC(ctx, CL_MEM_WRITE_ONLY,
                        sizeof(float) * N);

        // Create kernel
        cl::Kernel kernel(program, "add_vectors");

        // Set arguments
        kernel.setArg(0, bufA);
        kernel.setArg(1, bufB);
        kernel.setArg(2, bufC);

        // Launch
        cl::CommandQueue queue(ctx, device);
        queue.enqueueNDRangeKernel(kernel, cl::NullRange,
                                   cl::NDRange(N), cl::NullRange);
        queue.finish();

        // Read back
        queue.enqueueReadBuffer(bufC, CL_TRUE, 0,
                                sizeof(float) * N, c.data());

        // Verify a few elements
        bool ok = true;
        for (size_t i = 0; i < N; ++i)
        {
            if (c[i] != a[i] + b[i]) { ok = false; break; }
        }
        std::cout << "Demo " << (ok ? "succeeded" : "failed") << '\n';
    }
    catch (const std::exception &e)
    {
        std::cerr << "ERROR: " << e.what() << '\n';
        return 1;
    }

    return 0;
}