// This file is part of the 64k demo project.
// It implements the asset packer tool for demoscene resource management.
// Converts external files into embedded C++ byte arrays and look-up records.

#include <cstdio> // for simplicity, use fprintf() for output generation
#include <fstream>
#include <map>
#include <regex>     // For std::regex
#include <stdexcept> // For std::stof exceptions
#include <string>
#include <vector>

#include "procedural/generator.h" // For ProcGenFunc and procedural functions
#include "util/asset_manager.h"   // For AssetRecord and AssetId

// Map of procedural function names to their pointers (used only internally by
// asset_packer here, not generated)
static const std::map<std::string, ProcGenFunc> kAssetPackerProcGenFuncMap = {
    {"gen_noise", procedural::gen_noise},
    {"gen_grid", procedural::gen_grid},
    {"make_periodic", procedural::make_periodic},
};

// Helper struct to hold all information about an asset during parsing
struct AssetBuildInfo {
  std::string name;
  std::string filename; // Original filename for static assets
  bool is_procedural;
  std::string proc_func_name;     // Function name string
  std::vector<float> proc_params; // Parameters for procedural function

  // For generated C++ code
  std::string data_array_name;    // ASSET_DATA_xxx for static
  std::string params_array_name;  // ASSET_PROC_PARAMS_xxx for procedural
  std::string func_name_str_name; // ASSET_PROC_FUNC_STR_xxx for procedural
};

int main(int argc, char* argv[]) {
  if (argc != 4) {
    fprintf(stderr,
            "Usage: %s <assets.txt_path> <output_assets_h_path> "
            "<output_assets_data_cc_path>\n",
            argv[0]);
    return 1;
  }

  std::string assets_txt_path = argv[1];
  std::string output_assets_h_path = argv[2];
  std::string output_assets_data_cc_path = argv[3];

  std::ifstream assets_txt_file(assets_txt_path);
  if (!assets_txt_file.is_open()) {
    fprintf(stderr, "Error: Could not open assets.txt at %s\n",
            assets_txt_path.c_str());
    return 1;
  }

  FILE* assets_h_file = std::fopen(output_assets_h_path.c_str(), "w");
  if (!assets_h_file) {
    fprintf(stderr, "Error: Could not open output assets.h at %s\n",
            output_assets_h_path.c_str());
    return 1;
  }

  FILE* assets_data_cc_file =
      std::fopen(output_assets_data_cc_path.c_str(), "w");
  if (!assets_data_cc_file) {
    fprintf(stderr, "Error: Could not open output assets_data.cc at %s\n",
            output_assets_data_cc_path.c_str());
    return 1;
  }

  // Generate assets.h header
  fprintf(
      assets_h_file,
      "// This file is auto-generated by asset_packer.cc. Do not edit.\n\n");
  fprintf(assets_h_file, "#pragma once\n");
  fprintf(assets_h_file, "#include <cstdint>\n\n");
  fprintf(assets_h_file, "enum class AssetId : uint16_t {\n");

  std::string generated_header_name =
      output_assets_h_path.substr(output_assets_h_path.find_last_of("/\\") + 1);

  // Generate assets_data.cc header
  fprintf(
      assets_data_cc_file,
      "// This file is auto-generated by asset_packer.cc. Do not edit.\n\n");
  fprintf(assets_data_cc_file, "#include \"util/asset_manager.h\"\n");
  fprintf(assets_data_cc_file, "#include \"%s\"\n\n",
          generated_header_name.c_str());

  // Forward declare procedural functions for AssetRecord initialization
  fprintf(assets_data_cc_file,
          "namespace procedural { void gen_noise(uint8_t*, int, int, const "
          "float*, int); }\n");
  fprintf(assets_data_cc_file,
          "namespace procedural { void gen_grid(uint8_t*, int, int, const "
          "float*, int); }\n");
  fprintf(assets_data_cc_file,
          "namespace procedural { void make_periodic(uint8_t*, int, int, const "
          "float*, int); }\n\n");

  std::vector<AssetBuildInfo> asset_build_infos;
  int asset_id_counter = 0;
  std::string line;
  // Updated regex pattern for new asset list format (Name, CompressionType,
  // Filename, Description)
  std::regex asset_line_regex(
      R"(^\s*([A-Z0-9_]+)\s*,\s*(PROC\([^)]*\)|[^,]+)\s*(?:,\s*([^,]*))?\s*(?:,\s*\"(.*)\")?\s*$)");

  while (std::getline(assets_txt_file, line)) {
    if (line.empty() || line[0] == '#')
      continue;

    std::smatch matches;
    if (std::regex_search(line, matches, asset_line_regex) &&
        matches.size() >= 3) {
      AssetBuildInfo info;
      info.name = matches[1].str();
      std::string compression_type_str = matches[2].str();
      // Filename is now matches[3]
      info.filename =
          (matches.size() >= 4 && matches[3].matched) ? matches[3].str() : "_";

      info.data_array_name = "ASSET_DATA_" + info.name;
      info.params_array_name = "ASSET_PROC_PARAMS_" + info.name;
      info.func_name_str_name = "ASSET_PROC_FUNC_STR_" + info.name;
      info.is_procedural = false;

      if (compression_type_str.rfind("PROC(", 0) == 0) {
        info.is_procedural = true;
        size_t open_paren = compression_type_str.find('(');
        size_t close_paren = compression_type_str.rfind(')');
        if (open_paren == std::string::npos ||
            close_paren == std::string::npos) {
          fprintf(stderr,
                  "Error: Invalid PROC() syntax for asset: %s, string: [%s]\n",
                  info.name.c_str(), compression_type_str.c_str());
          return 1;
        }
        std::string func_and_params_str = compression_type_str.substr(
            open_paren + 1, close_paren - open_paren - 1);

        size_t params_start = func_and_params_str.find(',');
        if (params_start != std::string::npos) {
          std::string params_str = func_and_params_str.substr(params_start + 1);
          info.proc_func_name = func_and_params_str.substr(0, params_start);

          size_t current_pos = 0;
          while (current_pos < params_str.length()) {
            size_t comma_pos = params_str.find(',', current_pos);
            std::string param_val_str =
                (comma_pos == std::string::npos)
                    ? params_str.substr(current_pos)
                    : params_str.substr(current_pos, comma_pos - current_pos);
            param_val_str.erase(0, param_val_str.find_first_not_of(" \t\r\n"));
            param_val_str.erase(param_val_str.find_last_not_of(" \t\r\n") + 1);
            try {
              info.proc_params.push_back(std::stof(param_val_str));
            } catch (...) {
              fprintf(stderr, "Error: Invalid proc param for %s: %s\n",
                      info.name.c_str(), param_val_str.c_str());
              return 1;
            }
            if (comma_pos == std::string::npos)
              break;
            current_pos = comma_pos + 1;
          }
        } else {
          info.proc_func_name = func_and_params_str;
        }

        // Validate procedural function name
        // kAssetPackerProcGenFuncMap is defined globally for validation
        if (kAssetPackerProcGenFuncMap.find(info.proc_func_name) ==
            kAssetPackerProcGenFuncMap.end()) {
          fprintf(stderr,
                  "Error: Unknown procedural function: %s for asset: %s\n",
                  info.proc_func_name.c_str(), info.name.c_str());
          return 1;
        }
      }

      asset_build_infos.push_back(info);
      fprintf(assets_h_file, "  ASSET_%s = %d,\n", info.name.c_str(),
              asset_id_counter++);
    } else {
      fprintf(stderr, "Warning: Skipping malformed line in assets.txt: %s\n",
              line.c_str());
    }
  }

  fprintf(assets_h_file, "  ASSET_LAST_ID = %d,\n", asset_id_counter);
  fprintf(assets_h_file, "};\n");

  fprintf(assets_h_file,
          "#include \"util/asset_manager.h\"\n"); // Include here AFTER enum
                                                  // definition
  std::fclose(assets_h_file);

  for (const auto& info : asset_build_infos) {
    if (!info.is_procedural) {
      std::string base_dir =
          assets_txt_path.substr(0, assets_txt_path.find_last_of("/\\") + 1);
      std::ifstream asset_file(base_dir + info.filename, std::ios::binary);
      if (!asset_file.is_open()) {
        fprintf(stderr, "Error: Could not open asset file: %s\n",
                (base_dir + info.filename).c_str());
        return 1;
      }
      std::vector<uint8_t> buffer((std::istreambuf_iterator<char>(asset_file)),
                                  std::istreambuf_iterator<char>());
      fprintf(assets_data_cc_file, "static const uint8_t %s[] = {\n    ",
              info.data_array_name.c_str());
      for (size_t i = 0; i < buffer.size(); ++i) {
        if (i > 0 && i % 12 == 0)
          fprintf(assets_data_cc_file, "\n    ");
        fprintf(assets_data_cc_file, "0x%02x%s", buffer[i],
                (i == buffer.size() - 1 ? "" : ", "));
      }
      fprintf(assets_data_cc_file, "\n};\n");
    } else {
      fprintf(assets_data_cc_file, "static const float %s[] = {",
              info.params_array_name.c_str());
      for (size_t i = 0; i < info.proc_params.size(); ++i) {
        if (i > 0)
          fprintf(assets_data_cc_file, ", ");
        fprintf(assets_data_cc_file, "%f", info.proc_params[i]);
      }
      fprintf(assets_data_cc_file, "};\n\n");
      fprintf(assets_data_cc_file, "static const char* %s = \"%s\";\n\n",
              info.func_name_str_name.c_str(), info.proc_func_name.c_str());
    }
  }

  fprintf(assets_data_cc_file, "extern const AssetRecord g_assets[] = {\n");
  for (const auto& info : asset_build_infos) {
    fprintf(assets_data_cc_file, "    { ");
    if (info.is_procedural) {
      fprintf(assets_data_cc_file, "nullptr, 0, true, %s, %s, %zu",
              info.func_name_str_name.c_str(), info.params_array_name.c_str(),
              info.proc_params.size());
    } else {
      fprintf(assets_data_cc_file, "%s, sizeof(%s), false, nullptr, nullptr, 0",
              info.data_array_name.c_str(), info.data_array_name.c_str());
    }
    fprintf(assets_data_cc_file, " },\n");
  }
  fprintf(assets_data_cc_file, "};\n");
  fprintf(assets_data_cc_file,
          "extern const size_t g_assets_count = sizeof(g_assets) / "
          "sizeof(g_assets[0]);\n");

  std::fclose(assets_data_cc_file);

  printf("Asset packer successfully generated records for %zu assets.\n",
         asset_build_infos.size());

  return 0;
}