// This file is part of the 64k demo project.
// It implements the generic asset retrieval logic with runtime caching.

#include "util/asset_manager.h"
#if defined(USE_TEST_ASSETS)
#include "test_assets.h"
#else
#include "generated/assets.h"
#endif /* defined(USE_TEST_ASSETS) */

#include <cstdlib>  // For free
#include <iostream> // For std::cerr
#include <map>      // For kAssetManagerProcGenFuncMap
#include <new>      // For placement new
#include <string>   // For std::string in map
#include <vector>   // For potential dynamic allocation for procedural assets

#include "procedural/generator.h" // For ProcGenFunc and procedural functions

// Map of procedural function names to their pointers (for runtime dispatch)
static const std::map<std::string, ProcGenFunc> kAssetManagerProcGenFuncMap = {
    {"gen_noise", procedural::gen_noise},
    {"gen_grid", procedural::gen_grid},
    {"make_periodic", procedural::make_periodic},
};

// Array-based cache for assets.
// Initialized to all zeros (nullptr data, 0 size, false is_procedural)
// The size is derived from the generated ASSET_LAST_ID enum value.
static AssetRecord g_asset_cache[(size_t)AssetId::ASSET_LAST_ID] = {};

const uint8_t* GetAsset(AssetId asset_id, size_t* out_size) {
  uint16_t index = (uint16_t)asset_id;

  // Assert that ASSET_LAST_ID is not used for retrieval.
  // This ensures the size of the cache is correctly used and not accessed out
  // of bounds. If this assert fails, it means assets.txt has an ID larger than
  // expected or ASSET_LAST_ID is not correctly generated/updated. This is a
  // design decision: ASSET_LAST_ID is purely for sizing and range checking, not
  // for a valid asset retrieval itself.
  if (index >= (uint16_t)AssetId::ASSET_LAST_ID) {
    if (out_size)
      *out_size = 0;
    return nullptr; // Invalid asset_id
  }

  // Check cache first
  if (g_asset_cache[index].data != nullptr) {
    if (out_size)
      *out_size = g_asset_cache[index].size;
    return g_asset_cache[index].data;
  }

  const AssetRecord* assets = GetAssetRecordTable();
  size_t count = GetAssetCount();

  // Not in cache, retrieve from static data (packed in binary) or generate
  // procedurally
  if (index >= count) {
    if (out_size)
      *out_size = 0;
    return nullptr; // Invalid asset_id or asset not in static packed data.
  }

  AssetRecord source_record = assets[index];
  
  AssetRecord cached_record = source_record;

  if (source_record.is_procedural) {
    // Dynamically generate the asset
    auto it =
        kAssetManagerProcGenFuncMap.find(source_record.proc_func_name_str);
    if (it == kAssetManagerProcGenFuncMap.end()) {
      std::cerr << "Error: Unknown procedural function at runtime: "
                << source_record.proc_func_name_str << std::endl;
      if (out_size)
        *out_size = 0;
      return nullptr; // Procedural asset without a generation function.
    }
    ProcGenFunc proc_gen_func_ptr = it->second;

    // For this demo, assuming procedural textures are RGBA8 256x256 (for
    // simplicity and bump mapping). A more generic solution would pass
    // dimensions in proc_params.
    int width = 256, height = 256;
    size_t data_size = width * height * 4; // RGBA8
    uint8_t* generated_data = new (std::nothrow) uint8_t[data_size];
    if (!generated_data) {
      std::cerr << "Error: Failed to allocate memory for procedural asset."
                << std::endl;
      if (out_size)
        *out_size = 0;
      return nullptr;
    }
    proc_gen_func_ptr(generated_data, width, height, source_record.proc_params,
                      source_record.num_proc_params);

    cached_record.data = generated_data;
    cached_record.size = data_size;
    cached_record.is_procedural = true;
    // proc_gen_func, proc_params, num_proc_params already copied from
    // source_record

  } else {
    // Static asset (copy from g_assets)
    cached_record.is_procedural = false;
  }

  // Store in cache for future use
  g_asset_cache[index] = cached_record;

  if (out_size)
    *out_size = cached_record.size;
  return cached_record.data;
}

void DropAsset(AssetId asset_id, const uint8_t* asset) {
  uint16_t index = (uint16_t)asset_id;
  if (index >= (uint16_t)AssetId::ASSET_LAST_ID) {
    return; // Invalid asset_id
  }

  // Check if the asset is in cache and is procedural, and if the pointer
  // matches. This prevents accidentally freeing static data or freeing twice.
  if (g_asset_cache[index].data == asset &&
      g_asset_cache[index].is_procedural) {
    delete[] g_asset_cache[index].data;
    g_asset_cache[index] = {}; // Zero out the struct to force re-generation
  }
  // For static assets, no dynamic memory to free.
}