// Author: Simon-Pierre Boucher — contact@spboucher.ai
#pragma once

#include "core/allocator.h"

#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>

namespace MTL {
class Device;
class CommandQueue;
class Library;
class ComputePipelineState;
class FunctionConstantValues;
}

namespace forge {

// Owns the Metal device, the single command queue, the kernel library and
// the pipeline cache. One instance per process (Device::get()).
//
// All pipelines are created lazily on first use and cached by
// "kernel_name" or "kernel_name/constants_key" so steady-state training
// never touches pipeline creation.
class Device {
public:
    static Device& get();

    MTL::Device* mtl() const { return device_; }
    MTL::CommandQueue* queue() const { return queue_; }
    Allocator& allocator() { return *allocator_; }

    // Plain kernel, no function constants.
    MTL::ComputePipelineState* pipeline(const std::string& kernel_name);

    // Specialized kernel. constants_key must uniquely identify the constant
    // values (e.g. "am1_an0_ak1"); it is only used as a cache key.
    MTL::ComputePipelineState* pipeline(const std::string& kernel_name,
                                        const MTL::FunctionConstantValues* constants,
                                        const std::string& constants_key);

    std::string name() const;
    size_t recommended_working_set() const;

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

private:
    Device();
    ~Device();

    // Looks for forge.metallib next to the executable, in the current
    // directory, or at $FORGE_METALLIB.
    void load_library();

    MTL::Device* device_ = nullptr;
    MTL::CommandQueue* queue_ = nullptr;
    MTL::Library* library_ = nullptr;
    std::unique_ptr<Allocator> allocator_;
    std::mutex pipeline_mutex_;
    std::unordered_map<std::string, MTL::ComputePipelineState*> pipelines_;
};

} // namespace forge