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

#include <Foundation/Foundation.hpp>
#include <Metal/Metal.hpp>

#include <cstdio>
#include <cstdlib>
#include <filesystem>

#include <mach-o/dyld.h>

namespace forge {

namespace {

std::string executable_dir() {
    char buf[4096];
    uint32_t size = sizeof(buf);
    if (_NSGetExecutablePath(buf, &size) != 0) return {};
    std::error_code ec;
    auto canonical = std::filesystem::canonical(buf, ec);
    if (ec) return {};
    return canonical.parent_path().string();
}

[[noreturn]] void die(const char* msg, NS::Error* err) {
    std::fprintf(stderr, "forge: %s", msg);
    if (err && err->localizedDescription()) {
        std::fprintf(stderr, ": %s", err->localizedDescription()->utf8String());
    }
    std::fprintf(stderr, "\n");
    std::abort();
}

} // namespace

Device& Device::get() {
    static Device instance;
    return instance;
}

Device::Device() {
    device_ = MTL::CreateSystemDefaultDevice();
    if (!device_) die("no Metal device found", nullptr);
    queue_ = device_->newCommandQueue();
    if (!queue_) die("failed to create command queue", nullptr);
    allocator_ = std::make_unique<Allocator>(device_);
    load_library();
}

Device::~Device() {
    for (auto& [key, pso] : pipelines_) pso->release();
    if (library_) library_->release();
    allocator_.reset(); // must drop pooled buffers before the device
    if (queue_) queue_->release();
    if (device_) device_->release();
}

void Device::load_library() {
    std::vector<std::string> candidates;
    if (const char* env = std::getenv("FORGE_METALLIB")) candidates.push_back(env);
    if (auto dir = executable_dir(); !dir.empty()) candidates.push_back(dir + "/forge.metallib");
    candidates.push_back("forge.metallib");

    for (const auto& path : candidates) {
        if (!std::filesystem::exists(path)) continue;
        NS::Error* err = nullptr;
        NS::String* nspath = NS::String::string(path.c_str(), NS::UTF8StringEncoding);
        library_ = device_->newLibrary(nspath, &err);
        if (library_) return;
        die("failed to load forge.metallib", err);
    }
    die("forge.metallib not found (looked next to the executable, in cwd, and at $FORGE_METALLIB)",
        nullptr);
}

MTL::ComputePipelineState* Device::pipeline(const std::string& kernel_name) {
    return pipeline(kernel_name, nullptr, "");
}

MTL::ComputePipelineState* Device::pipeline(const std::string& kernel_name,
                                            const MTL::FunctionConstantValues* constants,
                                            const std::string& constants_key) {
    const std::string key =
        constants_key.empty() ? kernel_name : kernel_name + "/" + constants_key;

    std::lock_guard<std::mutex> lock(pipeline_mutex_);
    if (auto it = pipelines_.find(key); it != pipelines_.end()) return it->second;

    NS::Error* err = nullptr;
    NS::String* nsname = NS::String::string(kernel_name.c_str(), NS::UTF8StringEncoding);
    MTL::Function* fn = constants
        ? library_->newFunction(nsname, constants, &err)
        : library_->newFunction(nsname);
    if (!fn) die(("kernel not found: " + kernel_name).c_str(), err);

    MTL::ComputePipelineState* pso = device_->newComputePipelineState(fn, &err);
    fn->release();
    if (!pso) die(("pipeline creation failed: " + key).c_str(), err);

    pipelines_.emplace(key, pso);
    return pso;
}

std::string Device::name() const {
    return device_->name()->utf8String();
}

size_t Device::recommended_working_set() const {
    return device_->recommendedMaxWorkingSetSize();
}

} // namespace forge