#pragma once

#include <filesystem>
#include <memory>
#include <unordered_map>

#include "kernel_runtime.hpp"
#include "../utils/system.hpp"

namespace deep_gemm {

class KernelRuntimeCache {
    std::unordered_map<std::string, std::shared_ptr<KernelRuntime>> cache;

public:
    // TODO: consider cache capacity
    KernelRuntimeCache() = default;

    std::shared_ptr<KernelRuntime> get(const std::filesystem::path& dir_path) {
        // Hit the runtime cache
        if (const auto& iterator = cache.find(dir_path); iterator != cache.end())
            return iterator->second;

        if (KernelRuntime::check_validity(dir_path))
            return cache[dir_path] = std::make_shared<KernelRuntime>(dir_path);
        return nullptr;
    }

    // Preload all cached kernels from disk into memory
    void preload_all(const std::filesystem::path& cache_dir) {
        const auto kernel_cache_dir = cache_dir / "cache";
        
        // If cache directory doesn't exist, return early
        if (!std::filesystem::exists(kernel_cache_dir))
            return;
            
        int loaded_count = 0;
        int total_count = 0;
        
        // Iterate through all kernel subdirectories in the cache
        for (const auto& entry : std::filesystem::directory_iterator(kernel_cache_dir)) {
            if (!entry.is_directory())
                continue;
                
            total_count++;
            const auto& dir_path = entry.path();
            
            // Skip if already in memory cache
            if (cache.find(dir_path) != cache.end())
                continue;
            
            // Check validity and load the kernel
            if (KernelRuntime::check_validity(dir_path)) {
                try {
                    cache[dir_path] = std::make_shared<KernelRuntime>(dir_path);
                    loaded_count++;
                } catch (const std::exception& e) {
                    // Skip failed loads
                    if (get_env<int>("DG_JIT_DEBUG"))
                        printf("Warning: Failed to preload kernel from %s: %s\n", 
                               dir_path.c_str(), e.what());
                }
            }
        }
        
        if (get_env<int>("DG_JIT_DEBUG") || loaded_count > 0)
            printf("DeepGEMM: Preloaded %d/%d cached kernels from %s\n", 
                   loaded_count, total_count, kernel_cache_dir.c_str());
    }
};

static auto kernel_runtime_cache = std::make_shared<KernelRuntimeCache>();

} // namespace deep_gemm