// ── Gallery GEMM backend ──────────────────────────────────────────────────────
// Similarity engine for the identity matcher. Uploads the reference gallery
// once and computes the per-frame similarity matrix with a single SGEMM. The
// math backend is selected at compile time by CMake (SAE_GEMM_BACKEND):
// cuBLAS/CUDA, rocBLAS/HIP, or a portable CPU reference (SAE_GEMM_CPU). The two
// GPU paths are the ONLY translation units that include cublas/cuda or
// rocblas/hip headers; the CPU path pulls in no GPU headers at all and exists so
// the pipeline can build and be tested on a machine without a GPU (CI).

#include "inference/similarity.hpp"

#if defined(SAE_GEMM_CUDA)
#include <cublas_v2.h>
#include <cuda_runtime_api.h>
#elif defined(SAE_GEMM_ROCM)
#include <hip/hip_runtime.h>
#include <rocblas/rocblas.h>
#elif defined(SAE_GEMM_CPU)
// no external headers — portable reference implementation below
#else
#error "gemm_backend.cpp requires SAE_GEMM_CUDA, SAE_GEMM_ROCM or SAE_GEMM_CPU to be defined"
#endif

#include <cstring>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

namespace {

constexpr int kDim = 512;

#if defined(SAE_GEMM_CPU)

#if defined(SAE_GEMM_CBLAS)
#include <cblas.h>
#endif

// ── CPU reference engine ──────────────────────────────────────────────────────
// Used for CI and as the correctness oracle for the GPU backends.
//
// TRACES: AR-026, AR-027 | SR-001
// Backed by CBLAS (OpenBLAS), which CMake now REQUIRES for this backend. The
// scalar loop below is portable but scales badly: scoring one face against a
// 5000-embedding gallery is 2.6 MFLOP, and a crowded frame multiplies that by
// the face count. Since AR-003 removed the per-frame face cap and CI has no GPU,
// the CPU path is the one that has to hold up under a library-scale gallery
// (AR-027) rather than merely be correct — so falling back to it silently would
// mean measuring AR-027 on a path no release runs.
//
// The fallback is kept as the correctness oracle the two BLAS backends are
// diffed against when a similarity looks wrong, and is reachable only via
// -DSAE_ALLOW_SCALAR_GEMM=ON. The gallery is L2-normalised (as are the queries),
// so each similarity is a plain dot product. S is stored column-major to match
// the GPU backends: the gallery similarities for face fi start at
// result + fi*n_gallery().
class SimilarityEngine final : public ISimilarityEngine {
public:
    SimilarityEngine(const float* gallery_row_major, int n_gallery, int max_faces)
        : n_gallery_(n_gallery), max_faces_(max_faces),
          gallery_(gallery_row_major,
                   gallery_row_major + static_cast<size_t>(n_gallery) * kDim)
    {
        host_sims_.resize(static_cast<size_t>(max_faces_) * n_gallery_);
        std::cerr << "[similarity] CPU engine ("
#if defined(SAE_GEMM_CBLAS)
                  << "CBLAS"
#else
                  << "scalar fallback — no CBLAS; expect poor scaling on a large gallery"
#endif
                  << "): gallery resident in host RAM ("
                  << (gallery_.size() * sizeof(float)) / (1024 * 1024) << " MiB)\n";
    }

    int max_faces() const override { return max_faces_; }
    int n_gallery() const override { return n_gallery_; }

    /// TRACES: AR-026 | SR-001
    /// Promotions join the resident matrix, so the annex is scored by the same
    /// SGEMM as the baked references. std::vector already grows geometrically,
    /// so this is amortised O(1) per row.
    void append_rows(const float* rows_row_major, int n_rows) override {
        if (n_rows <= 0) return;
        gallery_.insert(gallery_.end(), rows_row_major,
                        rows_row_major + static_cast<size_t>(n_rows) * kDim);
        n_gallery_ += n_rows;
        host_sims_.resize(static_cast<size_t>(max_faces_) * n_gallery_);
    }

    const float* compute(const float* query_row_major, int n_faces) override {
        if (n_faces <= 0) return host_sims_.data();
        if (n_faces > max_faces_)
            throw std::runtime_error("SimilarityEngine: n_faces exceeds max_faces");

        // S(g, f) col-major = dot(gallery[g], query[f]). Viewed as row-major
        // [n_faces x n_gallery] that is exactly query * gallery^T, so it is one
        // GEMM rather than a loop nest.
#if defined(SAE_GEMM_CBLAS)
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
                    /*M=*/n_faces, /*N=*/n_gallery_, /*K=*/kDim,
                    /*alpha=*/1.0f,
                    query_row_major, /*lda=*/kDim,
                    gallery_.data(), /*ldb=*/kDim,
                    /*beta=*/0.0f,
                    host_sims_.data(), /*ldc=*/n_gallery_);
#else
        for (int f = 0; f < n_faces; ++f) {
            const float* q = query_row_major + static_cast<size_t>(f) * kDim;
            float* out = host_sims_.data() + static_cast<size_t>(f) * n_gallery_;
            for (int g = 0; g < n_gallery_; ++g) {
                const float* row = gallery_.data() + static_cast<size_t>(g) * kDim;
                float acc = 0.f;
                for (int d = 0; d < kDim; ++d) acc += row[d] * q[d];
                out[g] = acc;
            }
        }
#endif
        return host_sims_.data();
    }

private:
    int n_gallery_{0};
    int max_faces_{0};
    std::vector<float> gallery_;    // n_gallery × 512, row-major
    std::vector<float> host_sims_;  // max_faces × n_gallery, column-major
};

}  // namespace

std::unique_ptr<ISimilarityEngine> make_similarity_engine(
    const float* gallery_row_major, int n_gallery, int max_faces) {
    return std::make_unique<SimilarityEngine>(gallery_row_major, n_gallery, max_faces);
}

#else  // GPU backends (CUDA / ROCM)

struct GpuError : std::runtime_error {
    using std::runtime_error::runtime_error;
};

#if defined(SAE_GEMM_CUDA)

using stream_t      = cudaStream_t;
using blas_handle_t = cublasHandle_t;

inline void check_gpu(cudaError_t e, const char* what) {
    if (e != cudaSuccess)
        throw GpuError(std::string(what) + ": " + cudaGetErrorString(e));
}
inline void check_blas(cublasStatus_t s, const char* what) {
    if (s != CUBLAS_STATUS_SUCCESS)
        throw GpuError(std::string(what) + ": cublas error " + std::to_string(s));
}

inline void gpu_malloc(void** p, size_t bytes)                     { check_gpu(cudaMalloc(p, bytes), "cudaMalloc"); }
inline void gpu_free(void* p)                                      { cudaFree(p); }
inline void gpu_memcpy_h2d(void* dst, const void* src, size_t n, stream_t s) { check_gpu(cudaMemcpyAsync(dst, src, n, cudaMemcpyHostToDevice, s), "H2D"); }
inline void gpu_memcpy_d2h(void* dst, const void* src, size_t n, stream_t s) { check_gpu(cudaMemcpyAsync(dst, src, n, cudaMemcpyDeviceToHost, s), "D2H"); }
inline void gpu_memcpy_h2d_sync(void* dst, const void* src, size_t n)        { check_gpu(cudaMemcpy(dst, src, n, cudaMemcpyHostToDevice), "H2D_sync"); }
inline void gpu_memcpy_d2d_sync(void* dst, const void* src, size_t n)        { check_gpu(cudaMemcpy(dst, src, n, cudaMemcpyDeviceToDevice), "D2D_sync"); }
inline void stream_create(stream_t* s)                             { check_gpu(cudaStreamCreate(s), "cudaStreamCreate"); }
inline void stream_destroy(stream_t s)                             { cudaStreamDestroy(s); }
inline void stream_sync(stream_t s)                                { check_gpu(cudaStreamSynchronize(s), "cudaStreamSync"); }
inline void blas_create(blas_handle_t* h)                          { check_blas(cublasCreate(h), "cublasCreate"); }
inline void blas_destroy(blas_handle_t h)                          { cublasDestroy(h); }
inline void blas_set_stream(blas_handle_t h, stream_t s)           { check_blas(cublasSetStream(h, s), "cublasSetStream"); }
inline void blas_sgemm(blas_handle_t h, int m, int n, int k,
                       const float* A, const float* B, float* C) {
    const float alpha = 1.f, beta = 0.f;
    check_blas(cublasSgemm(h, CUBLAS_OP_T, CUBLAS_OP_N,
                            m, n, k, &alpha, A, k, B, k, &beta, C, m),
               "cublasSgemm");
}
inline const char* backend_name() { return "cuBLAS/CUDA"; }

#else  // SAE_GEMM_ROCM

using stream_t      = hipStream_t;
using blas_handle_t = rocblas_handle;

inline void check_gpu(hipError_t e, const char* what) {
    if (e != hipSuccess)
        throw GpuError(std::string(what) + ": " + hipGetErrorString(e));
}
inline void check_blas(rocblas_status s, const char* what) {
    if (s != rocblas_status_success)
        throw GpuError(std::string(what) + ": rocblas error " + std::to_string(s));
}

inline void gpu_malloc(void** p, size_t bytes)                     { check_gpu(hipMalloc(p, bytes), "hipMalloc"); }
inline void gpu_free(void* p)                                      { (void)hipFree(p); }
inline void gpu_memcpy_h2d(void* dst, const void* src, size_t n, stream_t s) { check_gpu(hipMemcpyAsync(dst, src, n, hipMemcpyHostToDevice, s), "H2D"); }
inline void gpu_memcpy_d2h(void* dst, const void* src, size_t n, stream_t s) { check_gpu(hipMemcpyAsync(dst, src, n, hipMemcpyDeviceToHost, s), "D2H"); }
inline void gpu_memcpy_h2d_sync(void* dst, const void* src, size_t n)        { check_gpu(hipMemcpy(dst, src, n, hipMemcpyHostToDevice), "H2D_sync"); }
inline void gpu_memcpy_d2d_sync(void* dst, const void* src, size_t n)        { check_gpu(hipMemcpy(dst, src, n, hipMemcpyDeviceToDevice), "D2D_sync"); }
inline void stream_create(stream_t* s)                             { check_gpu(hipStreamCreate(s), "hipStreamCreate"); }
inline void stream_destroy(stream_t s)                             { (void)hipStreamDestroy(s); }
inline void stream_sync(stream_t s)                                { check_gpu(hipStreamSynchronize(s), "hipStreamSync"); }
inline void blas_create(blas_handle_t* h)                          { check_blas(rocblas_create_handle(h), "rocblas_create_handle"); }
inline void blas_destroy(blas_handle_t h)                          { rocblas_destroy_handle(h); }
inline void blas_set_stream(blas_handle_t h, stream_t s)           { check_blas(rocblas_set_stream(h, s), "rocblas_set_stream"); }
inline void blas_sgemm(blas_handle_t h, int m, int n, int k,
                       const float* A, const float* B, float* C) {
    const float alpha = 1.f, beta = 0.f;
    // rocblas_sgemm is column-major; same transposition trick as cuBLAS:
    //   C(m×n) = A(k×m)^T * B(k×n)  →  S(N_gallery × n_faces) = G^T * Q
    check_blas(rocblas_sgemm(h, rocblas_operation_transpose, rocblas_operation_none,
                              m, n, k, &alpha, A, k, B, k, &beta, C, m),
               "rocblas_sgemm");
}
inline const char* backend_name() { return "rocBLAS/HIP"; }

#endif

class SimilarityEngine final : public ISimilarityEngine {
public:
    SimilarityEngine(const float* gallery_row_major, int n_gallery, int max_faces)
        : n_gallery_(n_gallery), max_faces_(max_faces)
    {
        gpu_malloc(reinterpret_cast<void**>(&d_query_),
                   static_cast<size_t>(max_faces_) * kDim * sizeof(float));

        // Allocates d_gallery_/d_sims_ at the initial row count; append_rows()
        // grows them geometrically from here.
        reserve_rows(std::max(n_gallery_, 1));
        const size_t gallery_floats = static_cast<size_t>(n_gallery_) * kDim;
        if (gallery_floats)
            gpu_memcpy_h2d_sync(d_gallery_, gallery_row_major, gallery_floats * sizeof(float));

        stream_create(&stream_);
        blas_create(&handle_);
        blas_set_stream(handle_, stream_);

        host_sims_.resize(static_cast<size_t>(max_faces_) * n_gallery_);

        std::cerr << "[similarity] " << backend_name() << " engine: gallery resident on GPU ("
                  << (gallery_floats * sizeof(float)) / (1024 * 1024) << " MiB)\n";
    }

    ~SimilarityEngine() override {
        if (d_gallery_) gpu_free(d_gallery_);
        if (d_query_)   gpu_free(d_query_);
        if (d_sims_)    gpu_free(d_sims_);
        if (handle_)    blas_destroy(handle_);
        if (stream_)    stream_destroy(stream_);
    }

    SimilarityEngine(const SimilarityEngine&) = delete;
    SimilarityEngine& operator=(const SimilarityEngine&) = delete;

    int max_faces() const override { return max_faces_; }
    int n_gallery() const override { return n_gallery_; }

    /// TRACES: AR-026 | SR-001
    /// Promotions join the GPU-resident matrix, so the annex is scored by the
    /// same SGEMM as the baked references rather than by a host-side loop.
    /// Capacity doubles on overflow, so the gallery is re-uploaded O(log n)
    /// times over a film rather than once per promotion.
    void append_rows(const float* rows_row_major, int n_rows) override {
        if (n_rows <= 0) return;
        const int want = n_gallery_ + n_rows;
        if (want > capacity_) reserve_rows(std::max(want, capacity_ * 2));

        gpu_memcpy_h2d_sync(d_gallery_ + static_cast<size_t>(n_gallery_) * kDim,
                            rows_row_major,
                            static_cast<size_t>(n_rows) * kDim * sizeof(float));
        n_gallery_ = want;
        host_sims_.resize(static_cast<size_t>(max_faces_) * n_gallery_);
    }

    const float* compute(const float* query_row_major, int n_faces) override {
        if (n_faces <= 0) return host_sims_.data();
        if (n_faces > max_faces_)
            throw std::runtime_error("SimilarityEngine: n_faces exceeds max_faces");

        gpu_memcpy_h2d(d_query_, query_row_major,
                       static_cast<size_t>(n_faces) * kDim * sizeof(float), stream_);

        // S (N_gallery × n_faces) col-major = G(512 × N_gallery)^T * Q(512 × n_faces)
        blas_sgemm(handle_, n_gallery_, n_faces, kDim, d_gallery_, d_query_, d_sims_);

        gpu_memcpy_d2h(host_sims_.data(), d_sims_,
                       static_cast<size_t>(n_gallery_) * n_faces * sizeof(float), stream_);
        stream_sync(stream_);
        return host_sims_.data();
    }

private:
    // Grow the resident gallery (and the similarity output sized against it) to
    // `rows` capacity, preserving the n_gallery_ rows already there. The copy is
    // device-to-device, so a promotion never re-uploads the baked gallery across
    // the bus.
    void reserve_rows(int rows) {
        if (rows <= capacity_) return;

        float* d_new_gallery = nullptr;
        gpu_malloc(reinterpret_cast<void**>(&d_new_gallery),
                   static_cast<size_t>(rows) * kDim * sizeof(float));
        if (d_gallery_ && n_gallery_ > 0)
            gpu_memcpy_d2d_sync(d_new_gallery, d_gallery_,
                                static_cast<size_t>(n_gallery_) * kDim * sizeof(float));
        if (d_gallery_) gpu_free(d_gallery_);
        d_gallery_ = d_new_gallery;

        // S is (capacity × n_faces); its contents are rewritten by every
        // compute(), so this one is a plain reallocation with nothing to keep.
        float* d_new_sims = nullptr;
        gpu_malloc(reinterpret_cast<void**>(&d_new_sims),
                   static_cast<size_t>(max_faces_) * rows * sizeof(float));
        if (d_sims_) gpu_free(d_sims_);
        d_sims_ = d_new_sims;

        capacity_ = rows;
    }

    int     n_gallery_{0};
    int     capacity_{0};
    int     max_faces_{0};
    float*  d_gallery_{nullptr};
    float*  d_query_{nullptr};
    float*  d_sims_{nullptr};
    std::vector<float> host_sims_;
    stream_t      stream_{};
    blas_handle_t handle_{};
};

}  // namespace

std::unique_ptr<ISimilarityEngine> make_similarity_engine(
    const float* gallery_row_major, int n_gallery, int max_faces) {
    return std::make_unique<SimilarityEngine>(gallery_row_major, n_gallery, max_faces);
}

#endif  // SAE_GEMM_CPU / GPU backends