#include "ggml_graph.hpp"
#include "backend.hpp"
#include "common.hpp"
#include "ggml.h"
#include <atomic>
#include <memory>
#include <mutex>

namespace pk {

// Process-global compute-thread override. 0 == unset (use the default thread
// count). Atomic so it is safe to set from one thread and read from the compute
// path; the benchmark CLI sets it once before any inference.
static std::atomic<int> g_num_threads{0};

void set_num_threads(int n) { g_num_threads.store(n < 0 ? 0 : n, std::memory_order_relaxed); }
int  num_threads()          { return g_num_threads.load(std::memory_order_relaxed); }

namespace {
// Process-global persistent backend. Created on first use with the configured
// thread count and reused across every graph computation. Holding it here (vs.
// per-Model) keeps the migration minimal: the components call pk::run_graph,
// which routes through this single Backend, so the per-call ggml_init/ggml_free
// churn is eliminated without threading a Backend handle through every layer.
// Task 2+ will fuse graphs on top of the same Backend.
std::unique_ptr<Backend> g_backend;
int g_backend_threads = 0;
// Serializes access to the process-global Backend. The old run_graph allocated
// a fresh per-call ggml context (so concurrent transcribes were independent);
// the shared Backend (one gallocr + pending-input list, not re-entrant) is not,
// so we serialize compute() across threads. In practice inference is driven
// from a single thread per process (the parallelism is inside the graph's
// worker threads), so this lock is uncontended; it only guards against a caller
// that drives transcribe() concurrently from multiple threads.
std::mutex g_backend_mutex;
} // namespace

// Default ggml compute-thread count when --threads is unset. On this class of
// many-core desktop CPUs (e.g. the 20-core box this was tuned on, and rf-detr's
// dual-CCD finding) using ALL cores REGRESSES throughput: a single-CCD-sized
// count parallelises the now-fused encoder graph while avoiding cross-CCD cache
// traffic and threadpool contention. Sweep on the fused 110m offline path
// (LibriSpeech): t=1 25.6, t=2 41.6, t=4 60.9, t=8 75.4, t=16 70.9, t=20 43.4
// RTFx -> 8 is the clear winner (+24% over 16, +74% over 20, +24% over the old
// default of 4). The unit tests pass an explicit per-call n_threads (4) and so
// are unaffected by this default.
constexpr int kDefaultThreads = 8;

Backend& global_backend() {
    // Lazy create (reset-safe: shutdown_backend() can free it, and a later call
    // recreates it). Always reached under g_backend_mutex (run_graph holds it)
    // or before any inference thread exists, so a plain null-check is sufficient.
    if (!g_backend) {
        const int g = g_num_threads.load(std::memory_order_relaxed);
        const int n = g > 0 ? g : kDefaultThreads;
        g_backend = std::make_unique<Backend>(n);
        g_backend_threads = n;
    }
    // Honor a late --threads override: keep the live backend's thread count in
    // sync with the global override (the override is set once before inference).
    const int g = g_num_threads.load(std::memory_order_relaxed);
    if (g > 0 && g != g_backend_threads) {
        g_backend->set_n_threads(g);
        g_backend_threads = g;
    }
    return *g_backend;
}

void shutdown_backend() {
    // Free the process-global backend explicitly. Required for GPU backends: the
    // backend (and its gallocr's device buffer) must be released while the CUDA/
    // etc. driver is still alive. Relying on static destruction frees it during
    // process exit, AFTER the driver's atexit handler, which aborts with
    // "driver shutting down". Call from main() before returning (model objects,
    // which hold their own device weight buffers, must already be destroyed).
    std::lock_guard<std::mutex> lock(g_backend_mutex);
    g_backend.reset();
    g_backend_threads = 0;
}

bool run_graph(size_t /*mem_bytes*/, int n_threads,
               const std::function<ggml_tensor*(ggml_context*)>& build,
               std::vector<float>& out) {
    std::lock_guard<std::mutex> lock(g_backend_mutex);
    Backend& be = global_backend();
    // When no global override is set, honor the caller's per-call n_threads (the
    // historical behavior, used by the unit tests). A positive global override
    // already pinned the backend's thread count in global_backend().
    const int g = g_num_threads.load(std::memory_order_relaxed);
    if (g <= 0 && n_threads > 0 && n_threads != g_backend_threads) {
        be.set_n_threads(n_threads);
        g_backend_threads = n_threads;
    }

    // Backend::compute builds the graph in a no_alloc context, allocates via the
    // persistent gallocr, pushes inputs AFTER alloc, computes, and reads back.
    return be.compute(build, out);
}

} // namespace pk