#include <iostream>
#include <vector>
#include <string>
#include <unordered_map>
#include <dlfcn.h>
#include <fstream>
#include <sstream>
#include <cstdlib>
#include <regex>

typedef void (*BatchFunc)(double*, double*, int);

// ──────────────────────────────────────────────
// JIT Engine with expression-level compile cache
// ──────────────────────────────────────────────
class JITEngine {
    // Cache: canonical expression → compiled function pointer
    std::unordered_map<std::string, BatchFunc> func_cache;
    // Keep handles alive so dlclose isn't called prematurely
    std::vector<void*> handles;

    BatchFunc current_func = nullptr;

    // Convert user expression → valid cmath C++ expression.
    // Handles:  x^y  →  pow(x,y)
    std::string translate_to_cpp(const std::string& raw) {
        std::string e = raw;

        // Replace x^y with pow(x,y) using a regex that handles simple tokens
        // on both sides.  For full robustness a proper parser would be needed,
        // but this covers the common cases correctly.
        std::regex pow_re(R"((\w+|\))\s*\^\s*(\w+|\())");
        std::string prev;
        while (prev != e) {
            prev = e;
            e = std::regex_replace(e, pow_re, "pow($1,$2)");
        }
        return e;
    }

    // Build a unique .so filename for this expression so cached .so files
    // on disk don't collide between different expressions.
    std::string hash_to_name(const std::string& expr) {
        size_t h = std::hash<std::string>{}(expr);
        std::ostringstream oss;
        oss << "jit_" << std::hex << h << ".so";
        return oss.str();
    }

public:
    // Returns the compiled function for `expr`, compiling only if not cached.
    BatchFunc get_or_compile(const std::string& raw_expr) {
        std::string expr = translate_to_cpp(raw_expr);

        // ── Cache hit ──────────────────────────────────────────────────────
        auto it = func_cache.find(expr);
        if (it != func_cache.end()) {
            return it->second;
        }

        // ── Cache miss: compile ────────────────────────────────────────────
        std::string so_name = hash_to_name(expr);

        std::ofstream ofs("jit_code.cpp");
        ofs << "#include <cmath>\n"
            << "using std::pow; using std::log; using std::log2; using std::log10;\n"
            << "using std::sin;  using std::cos;  using std::tan;\n"
            << "using std::asin; using std::acos; using std::atan; using std::atan2;\n"
            << "using std::sinh; using std::cosh; using std::tanh;\n"
            << "using std::sqrt; using std::exp;  using std::fabs;\n"
            << "using std::ceil; using std::floor;\n\n"
            << "extern \"C\" void batch_fn(double* x, double* res, int n) {\n"
            << "#pragma omp simd\n"
            << "  for (int i = 0; i < n; ++i) {\n"
            << "    double val = x[i];\n"   // keep 'val' as alias if user writes it
            << "    res[i] = " << expr << ";\n"
            << "  }\n"
            << "}\n";
        ofs.close();

        std::string cmd =
            "g++ -O3 -ffast-math -march=native -fPIC -shared "
            "-fopenmp -o " + so_name + " jit_code.cpp 2>/dev/null";
        int rc = std::system(cmd.c_str());

        if (rc != 0) {
            std::cerr << "[JITEngine] Compilation failed for: " << expr << "\n";
            return nullptr;
        }

        void* handle = dlopen(("./" + so_name).c_str(), RTLD_NOW);
        if (!handle) {
            std::cerr << "[JITEngine] dlopen failed: " << dlerror() << "\n";
            return nullptr;
        }

        auto fn = (BatchFunc)dlsym(handle, "batch_fn");
        if (!fn) {
            dlclose(handle);
            return nullptr;
        }

        handles.push_back(handle);   // keep alive
        func_cache[expr] = fn;
        return fn;
    }

    void run(const std::string& expr, double* x, double* res, int n) {
        BatchFunc fn = get_or_compile(expr);
        if (fn) fn(x, res, n);
    }

    ~JITEngine() {
        for (auto h : handles) dlclose(h);
    }
};

static JITEngine jit;

// ──────────────────────────────────────────────
// C API consumed by ctypes (main.py fallback)
// ──────────────────────────────────────────────
extern "C" {
    void compile(const char* expr) {
        jit.get_or_compile(std::string(expr));
    }

    void batch_evaluate_expr(const char* expr, double* x, double* res, int n) {
        jit.run(std::string(expr), x, res, n);
    }

    void bulk_run_expr(const char* expr, double* x, double* res, int n, int iterations) {
        if (!expr || !x || !res || n <= 0 || iterations <= 0) return;

        BatchFunc fn = jit.get_or_compile(std::string(expr));
        if (!fn) return;

        for (int it = 0; it < iterations; ++it) {
            fn(x, res, n);
        }
    }
}