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

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

#include <cstdio>  // For fprintf
#include <cstdlib> // For free
#include <cstring> // For strcmp
#include <new>     // For placement new

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

// Map of procedural function names to their pointers (for runtime dispatch)
struct ProcGenEntry {
  const char* name;
  ProcGenFunc func;
};

static const ProcGenEntry kAssetManagerProcGenFuncs[] = {
    {"gen_noise", procedural::gen_noise},
    {"gen_perlin", procedural::gen_perlin},
    {"gen_grid", procedural::gen_grid},
    {"make_periodic", procedural::make_periodic},
    {"gen_plasma", procedural::gen_plasma},
    {"gen_voronoi", procedural::gen_voronoi},
    {"gen_normalmap", procedural::gen_normalmap},
#if !defined(STRIP_ALL)
    {"gen_fail", procedural::gen_fail},
#endif
};
static const size_t kNumProcGenFuncs =
    sizeof(kAssetManagerProcGenFuncs) / sizeof(kAssetManagerProcGenFuncs[0]);

// Array-based cache for assets.
// Initialized to all zeros (nullptr data, 0 size, AssetType::STATIC)
// 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];

#if !defined(STRIP_ALL)
  if (source_record.type == AssetType::SPEC ||
      source_record.type == AssetType::MP3 ||
      source_record.type == AssetType::WGSL) {
    const char* path = (const char*)source_record.data;
    FILE* f = fopen(path, "rb");
    if (!f) {
      if (out_size)
        *out_size = 0;
      return nullptr;
    }
    fseek(f, 0, SEEK_END);
    long size = ftell(f);
    fseek(f, 0, SEEK_SET);
    uint8_t* buffer = new (std::nothrow) uint8_t[size + 1];
    if (!buffer) {
      fclose(f);
      if (out_size)
        *out_size = 0;
      return nullptr;
    }
    fread(buffer, 1, size, f);
    fclose(f);
    buffer[size] = 0; // Null-terminate

    g_asset_cache[index].data = buffer;
    g_asset_cache[index].size = size;
    g_asset_cache[index].type = source_record.type;
    if (out_size)
      *out_size = size;
    return buffer;
  }
#endif

  AssetRecord cached_record = source_record;

  if (source_record.compression == AssetCompression::ANS_ASCII) {
    // Decompress on first access. Buffer is null-terminated for C-string use.
    const size_t uncomp = source_record.uncompressed_size;
    uint8_t* buffer = new (std::nothrow) uint8_t[uncomp + 1];
    CHECK_RETURN_BEGIN(!buffer, nullptr)
    if (out_size)
      *out_size = 0;
    ERROR_MSG("Failed to allocate buffer for ANS-compressed asset id=%u",
              index);
    return nullptr;
    CHECK_RETURN_END
    size_t got = 0;
    if (!ans::Decode(source_record.data, source_record.size, buffer, uncomp,
                     &got, GetAnsAsciiHistogram()) ||
        got != uncomp) {
      delete[] buffer;
      if (out_size)
        *out_size = 0;
      ERROR_MSG("ANS decode failed for asset id=%u", index);
      return nullptr;
    }
    buffer[uncomp] = 0;
    cached_record.data = buffer;
    cached_record.size = uncomp;
    cached_record.uncompressed_size = uncomp;
    g_asset_cache[index] = cached_record;
    if (out_size)
      *out_size = uncomp;
    return buffer;
  }

  if (source_record.type == AssetType::PROC ||
      source_record.type == AssetType::PROC_GPU) {
    // Dynamically generate the asset
    ProcGenFunc proc_gen_func_ptr = nullptr;
    for (size_t i = 0; i < kNumProcGenFuncs; ++i) {
      if (strcmp(source_record.proc_func_name_str,
                 kAssetManagerProcGenFuncs[i].name) == 0) {
        proc_gen_func_ptr = kAssetManagerProcGenFuncs[i].func;
        break;
      }
    }

    CHECK_RETURN_BEGIN(proc_gen_func_ptr == nullptr, nullptr)
    if (out_size)
      *out_size = 0;
    ERROR_MSG("Unknown procedural function at runtime: %s",
              source_record.proc_func_name_str);
    return nullptr;
    CHECK_RETURN_END

    // 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 header_size = sizeof(uint32_t) * 2;
    size_t data_size = header_size + (size_t)width * height * 4; // RGBA8
    uint8_t* generated_data = new (std::nothrow) uint8_t[data_size];
    CHECK_RETURN_BEGIN(!generated_data, nullptr)
    if (out_size)
      *out_size = 0;
    ERROR_MSG("Failed to allocate memory for procedural asset");
    return nullptr;
    CHECK_RETURN_END

    // Write header
    uint32_t* header = (uint32_t*)(generated_data);
    header[0] = (uint32_t)width;
    header[1] = (uint32_t)height;

    // Generate data after header
    CHECK_RETURN_BEGIN(!proc_gen_func_ptr(generated_data + header_size, width,
                                          height, source_record.proc_params,
                                          source_record.num_proc_params),
                       nullptr)
    delete[] generated_data;
    if (out_size)
      *out_size = 0;
    ERROR_MSG("Procedural generation failed for asset: %s",
              source_record.proc_func_name_str);
    return nullptr;
    CHECK_RETURN_END

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

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

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

TextureAsset GetTextureAsset(AssetId asset_id) {
  size_t size = 0;
  const uint8_t* data = GetAsset(asset_id, &size);
  if (!data || size < 8) {
    return {0, 0, nullptr};
  }

  const uint32_t* header = (const uint32_t*)(data);
  return {(int)header[0], (int)header[1], data + 8};
}

MeshAsset GetMeshAsset(AssetId asset_id) {
  size_t size = 0;
  const uint8_t* data = GetAsset(asset_id, &size);
  if (!data || size < 8) {
    return {0, nullptr, 0, nullptr};
  }

  const uint8_t* ptr = data;
  uint32_t num_vertices = *(const uint32_t*)(ptr);
  ptr += sizeof(uint32_t);
  const MeshVertex* vertices = (const MeshVertex*)(ptr);
  ptr += num_vertices * sizeof(MeshVertex);
  uint32_t num_indices = *(const uint32_t*)(ptr);
  ptr += sizeof(uint32_t);
  const uint32_t* indices = (const uint32_t*)(ptr);

  return {num_vertices, vertices, num_indices, indices};
}

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 the pointer matches.
  // This prevents accidentally freeing static data or freeing twice.
  if (g_asset_cache[index].data != asset) {
    return; // Pointer mismatch — not our allocation
  }

  // Heap-owned: procedural, compressed, or disk-loaded (debug)
  if (g_asset_cache[index].type == AssetType::PROC ||
      g_asset_cache[index].type == AssetType::PROC_GPU ||
      g_asset_cache[index].compression != AssetCompression::NONE
#if !defined(STRIP_ALL)
      || g_asset_cache[index].type == AssetType::SPEC ||
      g_asset_cache[index].type == AssetType::MP3 ||
      g_asset_cache[index].type == AssetType::WGSL
#endif
  ) {
    delete[] g_asset_cache[index].data;
    g_asset_cache[index] = {};
  }
  // For static assets, no dynamic memory to free.
}

#if !defined(STRIP_ALL)
// Hot-reload: Clear asset cache to force reload from disk
// Note: This only works for assets that read from disk at runtime.
// Compiled-in assets cannot be hot-reloaded.
bool ReloadAssetsFromFile(const char* config_path) {
  (void)config_path; // Unused - just for API consistency

  // Clear cache to force reload
  for (size_t i = 0; i < (size_t)AssetId::ASSET_LAST_ID; ++i) {
    const AssetRecord& e = g_asset_cache[i];
    if (e.data &&
        (e.type == AssetType::PROC || e.type == AssetType::PROC_GPU ||
         e.compression != AssetCompression::NONE || e.type == AssetType::SPEC ||
         e.type == AssetType::MP3 || e.type == AssetType::WGSL)) {
      delete[] e.data;
    }
    g_asset_cache[i] = {};
  }

  fprintf(stderr, "[ReloadAssets] Cache cleared\n");
  return true;
}
#endif // !defined(STRIP_ALL)

AssetType GetAssetType(AssetId asset_id) {
  uint16_t index = (uint16_t)asset_id;
  if (index >= (uint16_t)AssetId::ASSET_LAST_ID)
    return AssetType::BINARY;
  return GetAssetRecordTable()[index].type;
}