#include "gallery_store.hpp"

#include <nlohmann/json.hpp>
#include <H5Cpp.h>
#include <chrono>
#include <fstream>
#include <iostream>
#include <stdexcept>

using json = nlohmann::json;

// ── HDF5 (see gallery_store.hpp for the full layout) ─────────────────────────
// A 170MB gallery JSON parses in ~18s (nlohmann). The same data as HDF5 loads in
// ~1s — a big win for the optimizer, which reloads the gallery per replay subprocess.
static bool ends_with(const std::string& s, const std::string& suf) {
    return s.size() >= suf.size() &&
           s.compare(s.size() - suf.size(), suf.size(), suf) == 0;
}

static std::vector<std::string> read_str_dataset(H5::H5File& file, const char* name, hsize_t a) {
    H5::DataSet ds = file.openDataSet(name);
    H5::StrType st = ds.getStrType();
    std::vector<std::string> out(a);
    if (st.isVariableStr()) {
        std::vector<char*> raw(a);
        ds.read(raw.data(), st);
        for (hsize_t i = 0; i < a; ++i) { out[i] = raw[i] ? raw[i] : ""; }
        H5::DataSpace sp = ds.getSpace();
        H5Dvlen_reclaim(st.getId(), sp.getId(), H5P_DEFAULT, raw.data());
    }
    return out;
}

static ActorGallery load_gallery_hdf5(const std::string& path) {
    std::cerr << "[gallery] loading " << path << " (HDF5)..." << std::flush;
    auto t0 = std::chrono::steady_clock::now();

    H5::H5File file(path, H5F_ACC_RDONLY);
    H5::DataSet emb_ds = file.openDataSet("embeddings");
    hsize_t dims[2];
    emb_ds.getSpace().getSimpleExtentDims(dims);   // [N, 512]
    const hsize_t N = dims[0];
    if (dims[1] != 512) throw std::runtime_error("gallery HDF5: embedding dim != 512");

    std::vector<float> flat(N * 512);
    emb_ds.read(flat.data(), H5::PredType::NATIVE_FLOAT);

    H5::DataSet off_ds = file.openDataSet("offset");
    hsize_t adim[1];
    off_ds.getSpace().getSimpleExtentDims(adim);
    const hsize_t A = adim[0];
    std::vector<int64_t> offset(A);
    off_ds.read(offset.data(), H5::PredType::NATIVE_INT64);
    std::vector<int32_t> count(A);
    file.openDataSet("count").read(count.data(), H5::PredType::NATIVE_INT32);

    auto imdb = read_str_dataset(file, "imdb_id", A);
    auto tmdb = read_str_dataset(file, "tmdb_id", A);
    auto jf   = read_str_dataset(file, "jellyfin_id", A);
    auto name = read_str_dataset(file, "name", A);
    std::vector<std::string> src_images;
    if (file.nameExists("source_images"))
        src_images = read_str_dataset(file, "source_images", N);

    ActorGallery gallery;
    gallery.actors.reserve(A);
    for (hsize_t a = 0; a < A; ++a) {
        ActorGallery::Actor actor;
        actor.imdb_id = imdb[a]; actor.tmdb_id = tmdb[a];
        actor.jellyfin_id = jf[a]; actor.name = name[a];
        for (int32_t e = 0; e < count[a]; ++e) {
            hsize_t row = offset[a] + e;
            Embedding emb;
            std::copy_n(flat.data() + row * 512, 512, emb.begin());
            actor.embeddings.push_back(emb);
            if (!src_images.empty())
                actor.source_images.push_back(src_images[row]);
        }
        gallery.actors.push_back(std::move(actor));
    }

    /// TRACES: GR-004 | SR-001
    // Absent /embedder group == a gallery written before model binding existed.
    // It stays readable; verify_gallery_embedder() decides what that means.
    if (file.nameExists("embedder")) {
        H5::Group eg = file.openGroup("embedder");
        H5::StrType str(H5::PredType::C_S1, H5T_VARIABLE);
        if (eg.attrExists("model_name"))
            eg.openAttribute("model_name").read(str, gallery.embedder.model_name);
        if (eg.attrExists("model_sha256"))
            eg.openAttribute("model_sha256").read(str, gallery.embedder.model_sha256);
        if (eg.attrExists("embed_dim"))
            eg.openAttribute("embed_dim").read(H5::PredType::NATIVE_INT32,
                                               &gallery.embedder.embed_dim);
    }

    if (file.nameExists("calibration")) {
        H5::Group cal = file.openGroup("calibration");
        cal.openAttribute("a").read(H5::PredType::NATIVE_FLOAT, &gallery.calib_a);
        cal.openAttribute("b").read(H5::PredType::NATIVE_FLOAT, &gallery.calib_b);
        int8_t valid = 0;
        cal.openAttribute("valid").read(H5::PredType::NATIVE_INT8, &valid);
        gallery.calib_valid = valid != 0;
        cal.openAttribute("hash").read(H5::PredType::NATIVE_UINT64, &gallery.calib_hash);
    }

    auto t1 = std::chrono::steady_clock::now();
    std::cerr << " built " << A << " actors / " << N << " embeddings in "
              << std::chrono::duration<double>(t1 - t0).count() << "s";
    if (gallery.calib_hash != 0)
        std::cerr << " (calibration cached: a=" << gallery.calib_a
                  << " b=" << gallery.calib_b << " valid=" << gallery.calib_valid << ")";
    std::cerr << "\n";
    return gallery;
}

static void write_str_dataset(H5::H5File& file, const char* name,
                              const std::vector<std::string>& values) {
    H5::StrType str_t(H5::PredType::C_S1, H5T_VARIABLE);
    hsize_t n = values.size();
    H5::DataSpace space(1, &n);
    H5::DataSet ds = file.createDataSet(name, str_t, space);
    std::vector<const char*> raw(n);
    for (hsize_t i = 0; i < n; ++i) raw[i] = values[i].c_str();
    ds.write(raw.data(), str_t);
}

static void save_gallery_hdf5(const std::string& path, const ActorGallery& gallery) {
    H5::H5File file(path, H5F_ACC_TRUNC);

    std::vector<float> flat;
    std::vector<int64_t> offset;
    std::vector<int32_t> count;
    std::vector<std::string> imdb, tmdb, jf, name, src_images;
    int64_t row = 0;
    for (const auto& a : gallery.actors) {
        offset.push_back(row);
        count.push_back(static_cast<int32_t>(a.embeddings.size()));
        row += static_cast<int64_t>(a.embeddings.size());
        for (size_t i = 0; i < a.embeddings.size(); ++i) {
            flat.insert(flat.end(), a.embeddings[i].begin(), a.embeddings[i].end());
            src_images.push_back(i < a.source_images.size() ? a.source_images[i] : "");
        }
        imdb.push_back(a.imdb_id); tmdb.push_back(a.tmdb_id);
        jf.push_back(a.jellyfin_id); name.push_back(a.name);
    }

    hsize_t N = flat.size() / 512;
    hsize_t emb_dims[2] = {N, 512};
    H5::DataSpace emb_space(2, emb_dims);
    file.createDataSet("embeddings", H5::PredType::NATIVE_FLOAT, emb_space)
        .write(flat.data(), H5::PredType::NATIVE_FLOAT);

    hsize_t A = gallery.actors.size();
    H5::DataSpace a_space(1, &A);
    file.createDataSet("offset", H5::PredType::NATIVE_INT64, a_space)
        .write(offset.data(), H5::PredType::NATIVE_INT64);
    file.createDataSet("count", H5::PredType::NATIVE_INT32, a_space)
        .write(count.data(), H5::PredType::NATIVE_INT32);

    write_str_dataset(file, "imdb_id", imdb);
    write_str_dataset(file, "tmdb_id", tmdb);
    write_str_dataset(file, "jellyfin_id", jf);
    write_str_dataset(file, "name", name);
    write_str_dataset(file, "source_images", src_images);

    /// TRACES: GR-004 | SR-001
    // Bind the file to the embedder that produced its vectors. Written only when
    // known — an empty stamp must round-trip as "unstamped", not as a stamp
    // claiming an unnamed model.
    if (!gallery.embedder.empty()) {
        H5::Group eg = file.createGroup("embedder");
        H5::DataSpace scalar(H5S_SCALAR);
        H5::StrType str(H5::PredType::C_S1, H5T_VARIABLE);
        eg.createAttribute("model_name", str, scalar).write(str, gallery.embedder.model_name);
        eg.createAttribute("model_sha256", str, scalar).write(str, gallery.embedder.model_sha256);
        eg.createAttribute("embed_dim", H5::PredType::NATIVE_INT32, scalar)
          .write(H5::PredType::NATIVE_INT32, &gallery.embedder.embed_dim);
    }

    if (gallery.calib_hash != 0) {
        H5::Group cal = file.createGroup("calibration");
        H5::DataSpace scalar(H5S_SCALAR);
        cal.createAttribute("a", H5::PredType::NATIVE_FLOAT, scalar)
           .write(H5::PredType::NATIVE_FLOAT, &gallery.calib_a);
        cal.createAttribute("b", H5::PredType::NATIVE_FLOAT, scalar)
           .write(H5::PredType::NATIVE_FLOAT, &gallery.calib_b);
        int8_t valid = gallery.calib_valid ? 1 : 0;
        cal.createAttribute("valid", H5::PredType::NATIVE_INT8, scalar)
           .write(H5::PredType::NATIVE_INT8, &valid);
        cal.createAttribute("hash", H5::PredType::NATIVE_UINT64, scalar)
           .write(H5::PredType::NATIVE_UINT64, &gallery.calib_hash);
    }

    std::cerr << "[gallery] saved " << A << " actors / " << N
              << " embeddings to " << path << " (HDF5)\n";
}

ActorGallery load_gallery(const std::string& path) {
    if (ends_with(path, ".h5") || ends_with(path, ".hdf5"))
        return load_gallery_hdf5(path);

    std::ifstream f(path);
    if (!f.is_open())
        throw std::runtime_error("load_gallery: cannot open " + path);

    std::cerr << "[gallery] loading " << path << "..." << std::flush;
    auto t0 = std::chrono::steady_clock::now();

    json j;
    f >> j;

    auto t1 = std::chrono::steady_clock::now();
    std::cerr << " parsed JSON in "
              << std::chrono::duration<double>(t1 - t0).count() << "s\n";

    ActorGallery gallery;

    /// TRACES: GR-004 | SR-001
    // Optional top-level "embedder" object, matching the HDF5 /embedder group.
    // Written by the JSON-era helper scripts; absent in anything older.
    if (j.contains("embedder") && j.at("embedder").is_object()) {
        const auto& je = j.at("embedder");
        gallery.embedder.model_name   = je.value("model_name", "");
        gallery.embedder.model_sha256 = je.value("model_sha256", "");
        gallery.embedder.embed_dim    = je.value("embed_dim", 512);
    }

    for (const auto& ja : j.at("actors")) {
        ActorGallery::Actor actor;
        actor.imdb_id     = ja.value("imdb_id", "");
        actor.tmdb_id     = ja.value("tmdb_id", "");
        // older make_jellyfin_gallery.py galleries used "jellyfin_person_id"
        actor.jellyfin_id = ja.value("jellyfin_id", ja.value("jellyfin_person_id", ""));
        actor.name        = ja.at("name").get<std::string>();

        if (ja.contains("source_images"))
            actor.source_images = ja.at("source_images").get<std::vector<std::string>>();

        for (const auto& je : ja.at("embeddings")) {
            Embedding emb = je.get<Embedding>();
            actor.embeddings.push_back(emb);
        }

        gallery.actors.push_back(std::move(actor));
    }

    size_t n_emb = 0;
    for (const auto& actor : gallery.actors) n_emb += actor.embeddings.size();
    auto t2 = std::chrono::steady_clock::now();
    std::cerr << "[gallery] built " << gallery.actors.size() << " actors / "
              << n_emb << " embeddings in "
              << std::chrono::duration<double>(t2 - t1).count() << "s\n";

    return gallery;
}

// Always writes HDF5. If `path` doesn't already end in .h5/.hdf5, the
// extension is replaced (galleries are never written as JSON anymore).
void save_gallery(const std::string& path, const ActorGallery& gallery) {
    std::string out_path = path;
    if (!ends_with(out_path, ".h5") && !ends_with(out_path, ".hdf5")) {
        auto dot = out_path.find_last_of('.');
        out_path = (dot == std::string::npos ? out_path : out_path.substr(0, dot)) + ".h5";
        std::cerr << "[gallery] save_gallery: writing HDF5 to " << out_path
                  << " (galleries are no longer written as JSON)\n";
    }
    save_gallery_hdf5(out_path, gallery);
}