/// TRACES: GR-004 | SR-001
#include "embedder_stamp.hpp"
#include "types.hpp"

#include <algorithm>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <map>
#include <mutex>
#include <sstream>
#include <stdexcept>
#include <vector>

namespace fs = std::filesystem;

// ── SHA-256 (FIPS 180-4) ──────────────────────────────────────────────────────
// Self-contained rather than pulled from OpenSSL: the gallery library already
// links OpenCV, HDF5, FFmpeg and a GPU backend, and the unit tests deliberately
// link none of those crypto stacks. ~80 lines of table-driven code is cheaper
// than another find_package that CI has to satisfy on an Intel N100.
namespace {

struct Sha256 {
    uint32_t h[8] = {0x6a09e667u, 0xbb67ae85u, 0x3c6ef372u, 0xa54ff53au,
                     0x510e527fu, 0x9b05688cu, 0x1f83d9abu, 0x5be0cd19u};
    uint64_t len = 0;
    uint8_t  buf[64]{};
    size_t   buf_n = 0;

    static uint32_t ror(uint32_t x, int n) { return (x >> n) | (x << (32 - n)); }

    void block(const uint8_t* p) {
        static const uint32_t k[64] = {
            0x428a2f98u,0x71374491u,0xb5c0fbcfu,0xe9b5dba5u,0x3956c25bu,0x59f111f1u,
            0x923f82a4u,0xab1c5ed5u,0xd807aa98u,0x12835b01u,0x243185beu,0x550c7dc3u,
            0x72be5d74u,0x80deb1feu,0x9bdc06a7u,0xc19bf174u,0xe49b69c1u,0xefbe4786u,
            0x0fc19dc6u,0x240ca1ccu,0x2de92c6fu,0x4a7484aau,0x5cb0a9dcu,0x76f988dau,
            0x983e5152u,0xa831c66du,0xb00327c8u,0xbf597fc7u,0xc6e00bf3u,0xd5a79147u,
            0x06ca6351u,0x14292967u,0x27b70a85u,0x2e1b2138u,0x4d2c6dfcu,0x53380d13u,
            0x650a7354u,0x766a0abbu,0x81c2c92eu,0x92722c85u,0xa2bfe8a1u,0xa81a664bu,
            0xc24b8b70u,0xc76c51a3u,0xd192e819u,0xd6990624u,0xf40e3585u,0x106aa070u,
            0x19a4c116u,0x1e376c08u,0x2748774cu,0x34b0bcb5u,0x391c0cb3u,0x4ed8aa4au,
            0x5b9cca4fu,0x682e6ff3u,0x748f82eeu,0x78a5636fu,0x84c87814u,0x8cc70208u,
            0x90befffau,0xa4506cebu,0xbef9a3f7u,0xc67178f2u};
        uint32_t w[64];
        for (int i = 0; i < 16; ++i)
            w[i] = (uint32_t(p[i * 4]) << 24) | (uint32_t(p[i * 4 + 1]) << 16) |
                   (uint32_t(p[i * 4 + 2]) << 8) | uint32_t(p[i * 4 + 3]);
        for (int i = 16; i < 64; ++i) {
            uint32_t s0 = ror(w[i - 15], 7) ^ ror(w[i - 15], 18) ^ (w[i - 15] >> 3);
            uint32_t s1 = ror(w[i - 2], 17) ^ ror(w[i - 2], 19) ^ (w[i - 2] >> 10);
            w[i] = w[i - 16] + s0 + w[i - 7] + s1;
        }
        uint32_t a = h[0], b = h[1], c = h[2], d = h[3];
        uint32_t e = h[4], f = h[5], g = h[6], hh = h[7];
        for (int i = 0; i < 64; ++i) {
            uint32_t S1 = ror(e, 6) ^ ror(e, 11) ^ ror(e, 25);
            uint32_t ch = (e & f) ^ (~e & g);
            uint32_t t1 = hh + S1 + ch + k[i] + w[i];
            uint32_t S0 = ror(a, 2) ^ ror(a, 13) ^ ror(a, 22);
            uint32_t mj = (a & b) ^ (a & c) ^ (b & c);
            uint32_t t2 = S0 + mj;
            hh = g; g = f; f = e; e = d + t1;
            d = c; c = b; b = a; a = t1 + t2;
        }
        h[0] += a; h[1] += b; h[2] += c; h[3] += d;
        h[4] += e; h[5] += f; h[6] += g; h[7] += hh;
    }

    void update(const uint8_t* p, size_t n) {
        len += n;
        while (n) {
            size_t take = std::min(n, size_t(64) - buf_n);
            std::memcpy(buf + buf_n, p, take);
            buf_n += take; p += take; n -= take;
            if (buf_n == 64) { block(buf); buf_n = 0; }
        }
    }

    std::string hex() {
        uint64_t bits = len * 8;
        uint8_t pad = 0x80;
        update(&pad, 1);
        uint8_t zero = 0;
        while (buf_n != 56) update(&zero, 1);
        uint8_t tail[8];
        for (int i = 0; i < 8; ++i) tail[i] = uint8_t(bits >> (56 - i * 8));
        // update() would re-count these into len, but len is already frozen in bits.
        std::memcpy(buf + buf_n, tail, 8);
        block(buf);
        buf_n = 0;

        static const char* d = "0123456789abcdef";
        std::string out;
        out.reserve(64);
        for (int i = 0; i < 8; ++i)
            for (int s = 28; s >= 0; s -= 4)
                out += d[(h[i] >> s) & 0xF];
        return out;
    }
};

// (path, mtime, size) → digest. Hashing a 250 MB ONNX is cheap but not free, and
// the optimizer constructs many networks in one process against the same model.
std::mutex                                  g_hash_mu;
std::map<std::string, std::string>          g_hash_cache;

std::string short_hash(const std::string& hex) {
    return hex.size() > 12 ? hex.substr(0, 12) + "…" : hex;
}

}  // namespace

std::string sha256_hex(const std::string& bytes) {
    Sha256 s;
    s.update(reinterpret_cast<const uint8_t*>(bytes.data()), bytes.size());
    return s.hex();
}

std::string sha256_file_hex(const std::string& path) {
    if (path.empty()) return "";

    std::error_code ec;
    auto size  = fs::file_size(path, ec);
    if (ec) return "";
    auto mtime = fs::last_write_time(path, ec);
    if (ec) return "";

    std::ostringstream key;
    key << path << '|' << size << '|'
        << mtime.time_since_epoch().count();

    {
        std::lock_guard<std::mutex> lk(g_hash_mu);
        auto it = g_hash_cache.find(key.str());
        if (it != g_hash_cache.end()) return it->second;
    }

    std::ifstream f(path, std::ios::binary);
    if (!f) return "";
    Sha256 s;
    std::vector<char> chunk(1 << 20);
    while (f) {
        f.read(chunk.data(), static_cast<std::streamsize>(chunk.size()));
        std::streamsize got = f.gcount();
        if (got > 0) s.update(reinterpret_cast<const uint8_t*>(chunk.data()),
                              static_cast<size_t>(got));
    }
    std::string hex = s.hex();

    std::lock_guard<std::mutex> lk(g_hash_mu);
    g_hash_cache[key.str()] = hex;
    return hex;
}

// ── EmbedderStamp ─────────────────────────────────────────────────────────────

std::string EmbedderStamp::describe() const {
    std::string name = model_name.empty() ? "<unnamed model>" : model_name;
    if (model_sha256.empty())
        return name + " (sha256 unavailable)";
    return name + " (sha256 " + short_hash(model_sha256) + ")";
}

EmbedderStamp make_embedder_stamp(const std::string& model_path) {
    EmbedderStamp s;
    if (model_path.empty()) return s;
    s.model_name   = fs::path(model_path).filename().string();
    s.model_sha256 = sha256_file_hex(model_path);
    if (s.model_sha256.empty())
        std::cerr << "[gallery] cannot hash embedder model " << model_path
                  << " — model binding falls back to filename only (GR-004)\n";
    return s;
}

bool require_gallery_stamp_from_env() {
    const char* v = std::getenv("SAE_REQUIRE_GALLERY_STAMP");
    return v && *v && std::strcmp(v, "0") != 0;
}

// ── Comparison ────────────────────────────────────────────────────────────────

StampCheck compare_embedder_stamps(const EmbedderStamp& built_with,
                                   const EmbedderStamp& loading_with,
                                   const std::string&   gallery_desc,
                                   const std::string&   embedder_desc) {
    StampCheck out;
    std::ostringstream m;

    // The gallery predates GR-004 (or was written by a tool that does not stamp).
    if (built_with.empty()) {
        out.verdict = StampVerdict::unstamped;
        m << "gallery '" << gallery_desc << "' carries no embedder stamp (GR-004).\n"
          << "  gallery was built with : UNKNOWN — this file predates model binding\n"
          << "  embedder now loaded    : " << loading_with.describe()
          << "  [" << embedder_desc << "]\n"
          << "  If these are not the same model every similarity from this run is\n"
          << "  meaningless but will look plausible. Rebuild or re-stamp the gallery\n"
          << "  (scripts/stamp_gallery.py), or run with SAE_REQUIRE_GALLERY_STAMP=1 to\n"
          << "  make this a hard error.";
        out.message = m.str();
        return out;
    }

    // Gallery is stamped but we cannot say what is about to embed.
    if (loading_with.empty()) {
        out.verdict = StampVerdict::unknown_embedder;
        m << "cannot identify the embedder being used against gallery '"
          << gallery_desc << "' (GR-004).\n"
          << "  gallery was built with : " << built_with.describe() << "\n"
          << "  embedder now loaded    : UNKNOWN [" << embedder_desc << "]\n"
          << "  The binding cannot be checked, so it is not being checked.";
        out.message = m.str();
        return out;
    }

    const bool have_both_hashes =
        !built_with.model_sha256.empty() && !loading_with.model_sha256.empty();

    // Embedding width disagreeing is a mismatch on its own terms — different
    // spaces entirely, and it will not even be caught by a cosine that "looks fine".
    if (built_with.embed_dim != loading_with.embed_dim) {
        out.verdict = StampVerdict::mismatch;
        m << "gallery/embedder MODEL MISMATCH — refusing to run (GR-004).\n"
          << "  gallery was built with : " << built_with.describe()
          << ", dim=" << built_with.embed_dim << "  [" << gallery_desc << "]\n"
          << "  embedder now loaded    : " << loading_with.describe()
          << ", dim=" << loading_with.embed_dim << "  [" << embedder_desc << "]\n"
          << "  Embedding dimensions differ; these are not the same space.";
        out.message = m.str();
        return out;
    }

    if (have_both_hashes) {
        if (built_with.model_sha256 == loading_with.model_sha256) {
            out.verdict = StampVerdict::match;
            m << "embedder binding verified: " << built_with.describe();
            if (built_with.model_name != loading_with.model_name)
                m << " (gallery recorded it as '" << built_with.model_name
                  << "', loaded from '" << loading_with.model_name
                  << "' — same bytes, renamed file)";
            out.message = m.str();
            return out;
        }
        out.verdict = StampVerdict::mismatch;
        m << "gallery/embedder MODEL MISMATCH — refusing to run (GR-004).\n"
          << "  gallery was built with : " << built_with.model_name
          << "  sha256=" << built_with.model_sha256 << "\n"
          << "                           [" << gallery_desc << "]\n"
          << "  embedder now loaded    : " << loading_with.model_name
          << "  sha256=" << loading_with.model_sha256 << "\n"
          << "                           [" << embedder_desc << "]\n"
          << "  Cosine similarities between embeddings from different models are\n"
          << "  meaningless but look plausible. Rebuild the gallery with the loaded\n"
          << "  model, or point the embedder at the model the gallery was built with.";
        out.message = m.str();
        return out;
    }

    // One side has no hash (e.g. a TRT deployment with the .onnx absent). Names
    // are all we have; agreeing on them is evidence, not proof.
    if (!built_with.model_name.empty() &&
        built_with.model_name == loading_with.model_name) {
        out.verdict = StampVerdict::weak_match;
        m << "embedder binding UNPROVEN for gallery '" << gallery_desc << "' (GR-004).\n"
          << "  gallery was built with : " << built_with.describe() << "\n"
          << "  embedder now loaded    : " << loading_with.describe()
          << "  [" << embedder_desc << "]\n"
          << "  Filenames agree but at least one SHA-256 is unavailable, so an\n"
          << "  in-place re-export under the same name would not be detected.";
        out.message = m.str();
        return out;
    }

    out.verdict = StampVerdict::mismatch;
    m << "gallery/embedder MODEL MISMATCH — refusing to run (GR-004).\n"
      << "  gallery was built with : " << built_with.describe()
      << "  [" << gallery_desc << "]\n"
      << "  embedder now loaded    : " << loading_with.describe()
      << "  [" << embedder_desc << "]\n"
      << "  Cosine similarities between embeddings from different models are\n"
      << "  meaningless but look plausible. Rebuild the gallery with the loaded\n"
      << "  model, or point the embedder at the model the gallery was built with.";
    out.message = m.str();
    return out;
}

void enforce_embedder_stamp(const EmbedderStamp& built_with,
                            const EmbedderStamp& loading_with,
                            const std::string&   gallery_desc,
                            const std::string&   embedder_desc,
                            bool                 require_stamp) {
    const bool strict = require_stamp || require_gallery_stamp_from_env();
    StampCheck chk = compare_embedder_stamps(built_with, loading_with,
                                             gallery_desc, embedder_desc);

    if (chk.fatal(strict)) {
        if (chk.verdict != StampVerdict::mismatch)
            throw std::runtime_error(chk.message +
                "\n  (fatal because SAE_REQUIRE_GALLERY_STAMP / --require-gallery-stamp is set)");
        throw std::runtime_error(chk.message);
    }

    if (chk.verdict == StampVerdict::match) {
        std::cerr << "[gallery] " << chk.message << "\n";
    } else {
        std::cerr << "\n[gallery] ***** WARNING (GR-004) *****\n"
                  << chk.message << "\n"
                  << "[gallery] ****************************\n\n";
    }
}

void verify_gallery_embedder(const ActorGallery& gallery,
                             const std::string&  gallery_path,
                             const std::string&  arcface_model_path,
                             bool                require_stamp) {
    enforce_embedder_stamp(gallery.embedder,
                           make_embedder_stamp(arcface_model_path),
                           gallery_path,
                           arcface_model_path.empty() ? "no --arcface given"
                                                      : arcface_model_path,
                           require_stamp);
}