// 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 <fstream>
#include <iostream>
#include <map>
#include <string>
#include <vector>
#include <stdexcept> // For std::stof exceptions
#include <regex>     // For std::regex

#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) {
    std::cerr << "Usage: " << argv[0]
              << " <assets.txt_path> <output_assets_h_path> "
                 "<output_assets_data_cc_path>\n";
    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()) {
    std::cerr << "Error: Could not open assets.txt at " << assets_txt_path
              << "\n";
    return 1;
  }

  std::ofstream assets_h_file(output_assets_h_path);
  if (!assets_h_file.is_open()) {
    std::cerr << "Error: Could not open output assets.h at "
              << output_assets_h_path << "\n";
    return 1;
  }

  std::ofstream assets_data_cc_file(output_assets_data_cc_path);
  if (!assets_data_cc_file.is_open()) {
    std::cerr << "Error: Could not open output assets_data.cc at "
              << output_assets_data_cc_path << "\n";
    return 1;
  }

  // Generate assets.h header
  assets_h_file
      << "// This file is auto-generated by asset_packer.cc. Do not edit.\n\n"
      << "#pragma once\n"
      << "#include <cstdint>\n\n"
      << "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
  assets_data_cc_file
      << "// This file is auto-generated by asset_packer.cc. Do not edit.\n\n"
      << "#include \"util/asset_manager.h\"\n"
      << "#include \"generated/" << generated_header_name << "\"\n\n";

  // Forward declare procedural functions for AssetRecord initialization
  assets_data_cc_file << "namespace procedural { void gen_noise(uint8_t*, int, int, const float*, int); }\n"
                      << "namespace procedural { void gen_grid(uint8_t*, int, int, const float*, int); }\n"
                      << "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) {
          std::cerr << "Error: Invalid PROC() syntax for asset: " << info.name << ", string: [" << compression_type_str << "]\n";
          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 (...) {
              std::cerr << "Error: Invalid proc param for " << info.name << ": " << param_val_str << "\n";
              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()) {
          std::cerr << "Error: Unknown procedural function: " << info.proc_func_name << " for asset: " << info.name << "\n";
          return 1;
        }
      }
      
      asset_build_infos.push_back(info);
      assets_h_file << "  ASSET_" << info.name << " = " << asset_id_counter++ << ",\n";
    } else {
        std::cerr << "Warning: Skipping malformed line in assets.txt: " << line << "\n";
    }
  }

  assets_h_file << "  ASSET_LAST_ID = " << asset_id_counter << ",\n";
  assets_h_file << "};\n";

  assets_h_file << "#include \"util/asset_manager.h\"\n"; // Include here AFTER enum definition
  assets_h_file.close();

  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()) {
              std::cerr << "Error: Could not open asset file: " << base_dir + info.filename << "\n";
              return 1;
          }
          std::vector<uint8_t> buffer((std::istreambuf_iterator<char>(asset_file)), std::istreambuf_iterator<char>());
          assets_data_cc_file << "static const uint8_t " << info.data_array_name << "[] = {\n    ";
          for (size_t i = 0; i < buffer.size(); ++i) {
              if (i > 0 && i % 12 == 0) assets_data_cc_file << "\n    ";
              assets_data_cc_file << "0x" << std::hex << (int)buffer[i] << std::dec << (i == buffer.size() - 1 ? "" : ", ");
          }
          assets_data_cc_file << "\n};\n";
      } else {
          assets_data_cc_file << "static const float " << info.params_array_name << "[] = {";
          for (size_t i = 0; i < info.proc_params.size(); ++i) {
              if (i > 0) assets_data_cc_file << ", ";
              assets_data_cc_file << info.proc_params[i];
          }
          assets_data_cc_file << "};\n\n";
          assets_data_cc_file << "static const char* " << info.func_name_str_name << " = \"" << info.proc_func_name << "\";\n\n";
      }
  }

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

  assets_data_cc_file.close();
  
  std::cout << "Asset packer successfully generated records for "
            << asset_build_infos.size() << " assets.\n";

  return 0;
}