// Author: Simon-Pierre Boucher — contact@spboucher.ai
#include "core/fmodel.h"

#include "core/device.h"

#include <Metal/Metal.hpp>

#include <nlohmann/json.hpp>

#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

#include <cinttypes>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <unordered_map>
#include <unordered_set>

namespace forge::fmodel {

namespace {

namespace fs = std::filesystem;
using nlohmann::json;

// Apple Silicon VM page — the alignment contract of the whole format.
constexpr size_t kPage = 16384;
constexpr uint32_t kShardMagic = 0x53475246; // "FRGS"
constexpr uint32_t kVersion = 1;

[[noreturn]] void die(const std::string& msg) {
    std::fprintf(stderr, "forge/fmodel: %s\n", msg.c_str());
    std::abort();
}

size_t page_align(size_t n) { return (n + kPage - 1) & ~(kPage - 1); }

uint64_t fnv1a64(const void* data, size_t n) {
    const uint8_t* p = static_cast<const uint8_t*>(data);
    uint64_t h = 1469598103934665603ull;
    for (size_t i = 0; i < n; ++i) {
        h ^= p[i];
        h *= 1099511628211ull;
    }
    return h;
}

std::string hex64(uint64_t v) {
    char buf[17];
    std::snprintf(buf, sizeof(buf), "%016" PRIx64, v);
    return buf;
}

DType dtype_from_name(const std::string& s) {
    if (s == "f32") return DType::F32;
    if (s == "f16") return DType::F16;
    if (s == "bf16") return DType::BF16;
    if (s == "u16") return DType::U16;
    if (s == "i32") return DType::I32;
    die("unknown dtype in manifest: " + s);
}

// Convert one f32 tensor into `dtype` bytes (identity for f32).
std::vector<uint8_t> convert_out(const Tensor& t, DType dtype) {
    const int64_t n = t.numel();
    const float* src = t.data<float>();
    std::vector<uint8_t> out(size_t(n) * dtype_size(dtype));
    if (dtype == DType::F32) {
        std::memcpy(out.data(), src, out.size());
    } else if (dtype == DType::F16) {
        f16_t* d = reinterpret_cast<f16_t*>(out.data());
        for (int64_t i = 0; i < n; ++i) d[i] = f16_t(src[i]);
    } else if (dtype == DType::BF16) {
        uint16_t* d = reinterpret_cast<uint16_t*>(out.data());
        for (int64_t i = 0; i < n; ++i) d[i] = float_to_bf16(src[i]);
    } else {
        die("save: unsupported storage dtype");
    }
    return out;
}

std::vector<std::pair<std::string, Var>> dedupe(
    const std::vector<std::pair<std::string, Var>>& named) {
    std::vector<std::pair<std::string, Var>> out;
    std::unordered_set<const void*> seen;
    for (const auto& [name, p] : named) {
        if (!p.defined()) continue;
        if (seen.insert(p.id()).second) out.emplace_back(name, p);
    }
    return out;
}

json load_json(const fs::path& p) {
    std::ifstream in(p);
    if (!in) die("cannot open " + p.string());
    return json::parse(in);
}

} // namespace

std::string save(const std::string& repo_dir, const std::string& config_json,
                 const std::vector<std::pair<std::string, Var>>& named_params,
                 const SaveOptions& opts) {
    const fs::path repo(repo_dir);
    fs::create_directories(repo / "manifests");
    fs::create_directories(repo / "objects");

    // Parent manifest (if any): hash -> existing location, for delta reuse.
    json parent;
    std::string parent_name;
    const fs::path head = repo / "manifest-latest.json";
    if (fs::exists(head)) {
        parent = load_json(head);
        parent_name = parent.value("self", "");
    }
    std::unordered_map<std::string, json> known; // content hash -> tensor entry
    if (parent.contains("tensors"))
        for (const auto& [name, e] : parent.at("tensors").items())
            known.emplace(e.at("hash").get<std::string>(), e);

    // Hash every (converted) tensor; split into reused vs new.
    struct Pending {
        std::string name;
        std::vector<uint8_t> bytes;
        std::string hash;
        std::vector<int64_t> shape;
    };
    json tensors = json::object();
    std::vector<Pending> fresh;
    size_t reused = 0, reused_bytes = 0;
    for (const auto& [name, p] : dedupe(named_params)) {
        Pending pd;
        pd.name = name;
        pd.shape = p.value().shape();
        pd.bytes = convert_out(p.value(), opts.dtype);
        pd.hash = hex64(fnv1a64(pd.bytes.data(), pd.bytes.size()));
        auto it = known.find(pd.hash);
        if (it != known.end()) {
            tensors[name] = it->second; // unchanged since parent: reference it
            ++reused;
            reused_bytes += pd.bytes.size();
        } else {
            fresh.push_back(std::move(pd));
        }
    }

    // Pack new tensors into page-aligned shards capped at shard_mb.
    const size_t cap = size_t(opts.shard_mb) * 1024 * 1024;
    size_t si = 0;
    size_t written_bytes = 0;
    while (si < fresh.size()) {
        std::vector<uint8_t> blob(kPage, 0); // shard header page
        {
            uint32_t* h = reinterpret_cast<uint32_t*>(blob.data());
            h[0] = kShardMagic;
            h[1] = kVersion;
        }
        json members = json::array();
        uint64_t shard_hash = 1469598103934665603ull;
        size_t end = si;
        while (end < fresh.size()) {
            const size_t start = page_align(blob.size());
            const size_t next = page_align(start + fresh[end].bytes.size());
            if (end > si && next > cap + kPage) break; // shard full
            blob.resize(start, 0);
            blob.insert(blob.end(), fresh[end].bytes.begin(), fresh[end].bytes.end());
            members.push_back(json{{"i", end}, {"offset", start}});
            shard_hash ^= fnv1a64(fresh[end].hash.data(), fresh[end].hash.size());
            shard_hash *= 1099511628211ull;
            ++end;
        }
        blob.resize(page_align(blob.size()), 0);

        const std::string shard_rel = "objects/" + hex64(shard_hash) + ".fshard";
        const fs::path shard_path = repo / shard_rel;
        if (!fs::exists(shard_path)) { // content-addressed: identical = skip
            const fs::path tmp = shard_path.string() + ".tmp";
            std::ofstream out(tmp, std::ios::binary);
            out.write(reinterpret_cast<const char*>(blob.data()),
                      std::streamsize(blob.size()));
            out.close();
            fs::rename(tmp, shard_path);
        }
        written_bytes += blob.size();

        for (const auto& m : members) {
            const Pending& pd = fresh[m.at("i").get<size_t>()];
            tensors[pd.name] = json{
                {"dtype", dtype_name(opts.dtype)}, {"shape", pd.shape},
                {"hash", pd.hash},                 {"shard", shard_rel},
                {"offset", m.at("offset")},        {"nbytes", pd.bytes.size()}};
        }
        si = end;
    }

    // Commit: numbered manifest + refreshed HEAD.
    int64_t seq = 0;
    for (const auto& e : fs::directory_iterator(repo / "manifests")) {
        (void)e;
        ++seq;
    }
    char mname[64];
    std::snprintf(mname, sizeof(mname), "manifest-%06lld.json",
                  static_cast<long long>(seq));
    json manifest{{"fmodel", kVersion},
                  {"self", mname},
                  {"parent", parent_name},
                  {"step", opts.step},
                  {"tag", opts.tag},
                  {"dtype", dtype_name(opts.dtype)},
                  {"config", json::parse(config_json.empty() ? "{}" : config_json)},
                  {"tensors", tensors}};
    const fs::path mpath = repo / "manifests" / mname;
    {
        std::ofstream out(mpath);
        out << manifest.dump(1);
    }
    fs::copy_file(mpath, head, fs::copy_options::overwrite_existing);

    std::printf("fmodel: %s — %zu tensors reused (%.1f MB), %zu written (%.1f MB)\n",
                (repo / mname).c_str(), reused, double(reused_bytes) / 1e6,
                fresh.size(), double(written_bytes) / 1e6);
    return mpath.string();
}

// ---- Snapshot ---------------------------------------------------------------

struct Snapshot::Impl {
    struct Mapped {
        void* base = nullptr;
        size_t len = 0;
        MTL::Buffer* buffer = nullptr;
    };
    struct Entry {
        DType dtype;
        std::vector<int64_t> shape;
        std::string shard;
        size_t offset, nbytes;
    };

    fs::path repo;
    std::string config;
    int64_t step = -1;
    std::vector<std::string> order;
    std::unordered_map<std::string, Entry> index;
    std::unordered_map<std::string, Mapped> maps; // shard rel-path -> mapping

    ~Impl() {
        for (auto& [rel, m] : maps) {
            if (m.buffer) m.buffer->release();
            if (m.base) munmap(m.base, m.len);
        }
    }

    const Mapped& mapped(const std::string& rel) {
        auto it = maps.find(rel);
        if (it != maps.end()) return it->second;

        const fs::path p = repo / rel;
        const int fd = ::open(p.c_str(), O_RDONLY);
        if (fd < 0) die("cannot open shard " + p.string());
        struct stat st;
        if (fstat(fd, &st) != 0) die("fstat failed: " + p.string());
        const size_t len = size_t(st.st_size);
        if (len % kPage != 0) die("shard not page-padded: " + p.string());

        // MAP_PRIVATE + RW: pages stay clean (backed by the file cache)
        // unless someone writes — and the read-only contract means nobody
        // does. newBuffer(bytesNoCopy) needs page alignment + page-multiple
        // length, which mmap and the writer guarantee.
        void* base = mmap(nullptr, len, PROT_READ | PROT_WRITE, MAP_PRIVATE, fd, 0);
        ::close(fd);
        if (base == MAP_FAILED) die("mmap failed: " + p.string());
        if (reinterpret_cast<const uint32_t*>(base)[0] != kShardMagic)
            die("bad shard magic: " + p.string());

        MTL::Buffer* buf = Device::get().mtl()->newBuffer(
            base, len, MTL::ResourceStorageModeShared, nullptr);
        if (!buf) die("newBuffer(bytesNoCopy) failed: " + p.string());
        return maps.emplace(rel, Mapped{base, len, buf}).first->second;
    }
};

Snapshot Snapshot::open(const std::string& path) {
    fs::path manifest(path);
    fs::path repo;
    if (fs::is_directory(manifest)) {
        repo = manifest;
        manifest = repo / "manifest-latest.json";
    } else {
        repo = manifest.parent_path();
        if (repo.filename() == "manifests") repo = repo.parent_path();
    }
    if (!fs::exists(manifest)) die("no manifest at " + manifest.string());

    auto impl = std::make_shared<Impl>();
    impl->repo = repo;
    json j = load_json(manifest);
    impl->config = j.value("config", json::object()).dump();
    impl->step = j.value("step", int64_t(-1));
    for (const auto& [name, e] : j.at("tensors").items()) {
        Impl::Entry ent;
        ent.dtype = dtype_from_name(e.at("dtype").get<std::string>());
        ent.shape = e.at("shape").get<std::vector<int64_t>>();
        ent.shard = e.at("shard").get<std::string>();
        ent.offset = e.at("offset").get<size_t>();
        ent.nbytes = e.at("nbytes").get<size_t>();
        impl->order.push_back(name);
        impl->index.emplace(name, std::move(ent));
    }

    Snapshot s;
    s.impl_ = std::move(impl);
    return s;
}

const std::string& Snapshot::config_json() const { return impl_->config; }
int64_t Snapshot::step() const { return impl_->step; }
std::vector<std::string> Snapshot::names() const { return impl_->order; }
bool Snapshot::has(const std::string& name) const { return impl_->index.count(name); }

const std::vector<int64_t>& Snapshot::shape(const std::string& name) const {
    auto it = impl_->index.find(name);
    if (it == impl_->index.end()) die("no tensor " + name);
    return it->second.shape;
}

Tensor Snapshot::tensor_f32(const std::string& name) const {
    auto it = impl_->index.find(name);
    if (it == impl_->index.end()) die("no tensor " + name);
    const Impl::Entry& e = it->second;
    const Impl::Mapped& m = impl_->mapped(e.shard);

    if (e.dtype == DType::F32)
        return Tensor::from_buffer(m.buffer, e.offset, e.shape); // zero-copy

    // f16/bf16 storage: convert into a fresh f32 tensor.
    Tensor out = Tensor::empty(e.shape);
    const int64_t n = out.numel();
    float* dst = out.data<float>();
    const uint8_t* src = static_cast<const uint8_t*>(m.base) + e.offset;
    if (e.dtype == DType::F16) {
        const f16_t* s = reinterpret_cast<const f16_t*>(src);
        for (int64_t i = 0; i < n; ++i) dst[i] = float(s[i]);
    } else if (e.dtype == DType::BF16) {
        const uint16_t* s = reinterpret_cast<const uint16_t*>(src);
        for (int64_t i = 0; i < n; ++i) dst[i] = bf16_to_float(s[i]);
    } else {
        die("tensor_f32: unsupported stored dtype for " + name);
    }
    return out;
}

} // namespace forge::fmodel