#include <fstream>
#include <iostream>
#include <map>
#include <string>
#include <vector>

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
  assets_h_file << "#pragma once\n";
  assets_h_file << "#include <cstdint>\n";
  assets_h_file << "#include <cstddef>\n\n";
  assets_h_file << "enum class AssetId : uint16_t {\n";

  // Generate assets_data.cc header
  assets_data_cc_file << "#include \"assets.h\"\n\n";
  assets_data_cc_file << "#include <vector>\n\n";

  std::string line;
  int asset_id_counter = 0;
  std::vector<std::string> asset_names;

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

    size_t first_comma = line.find(',');
    if (first_comma == std::string::npos) {
      std::cerr << "Warning: Skipping malformed line in assets.txt: " << line
                << "\n";
      continue;
    }

    std::string asset_name = line.substr(0, first_comma);
    asset_name.erase(0, asset_name.find_first_not_of(" \t\r\n"));
    asset_name.erase(asset_name.find_last_not_of(" \t\r\n") + 1);

    size_t second_comma = line.find(',', first_comma + 1);
    if (second_comma == std::string::npos) {
      std::cerr << "Warning: Skipping malformed line in assets.txt (missing "
                   "filename): "
                << line << "\n";
      continue;
    }

    std::string filename =
        line.substr(first_comma + 1, second_comma - first_comma - 1);
    filename.erase(0, filename.find_first_not_of(" \t\r\n"));
    filename.erase(filename.find_last_not_of(" \t\r\n") + 1);

    // Read the actual file
    std::string base_dir =
        assets_txt_path.substr(0, assets_txt_path.find_last_of("/\\") + 1);
    std::ifstream asset_file(base_dir + filename, std::ios::binary);
    if (!asset_file.is_open()) {
      std::cerr << "Error: Could not open asset file: " << base_dir + filename
                << "\n";
      return 1;
    }

    std::vector<uint8_t> buffer((std::istreambuf_iterator<char>(asset_file)),
                                std::istreambuf_iterator<char>());
    asset_file.close();

    asset_names.push_back(asset_name);

    // Add to assets.h enum
    assets_h_file << "  ASSET_" << asset_name << " = " << asset_id_counter
                  << ",\n";

    // Write data to assets_data.cc
    assets_data_cc_file << "const uint8_t ASSET_DATA_" << asset_name
                        << "[] = {";
    for (size_t i = 0; i < buffer.size(); ++i) {
      if (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";
    assets_data_cc_file << "const size_t ASSET_SIZE_" << asset_name << " = "
                        << buffer.size() << ";\n\n";

    asset_id_counter++;
  }

  assets_h_file << "};\n\n";

  // Generate GetAsset function declaration in assets.h
  assets_h_file << "const uint8_t *GetAsset(AssetId asset_id, size_t *out_size "
                   "= nullptr);\n";
  assets_h_file
      << "void DropAsset(AssetId asset_id, const uint8_t *asset); // For lazy "
         "decompression scaffolding\n";
  assets_h_file.close();

  // Generate GetAsset function implementation in assets_data.cc
  assets_data_cc_file
      << "const uint8_t *GetAsset(AssetId asset_id, size_t *out_size) {\n";
  assets_data_cc_file << "  if (out_size) *out_size = 0;\n"; // Default for now
  assets_data_cc_file << "  switch (asset_id) {\n";
  for (const std::string &name : asset_names) {
    assets_data_cc_file << "    case AssetId::ASSET_" << name << ":\n";
    assets_data_cc_file << "      if (out_size) *out_size = ASSET_SIZE_" << name
                        << ";\n";
    assets_data_cc_file << "      return ASSET_DATA_" << name << ";\n";
  }
  assets_data_cc_file << "    default:\n";
  assets_data_cc_file << "      return nullptr;\n";
  assets_data_cc_file << "  }\n";
  assets_data_cc_file << "}\n\n";

  // Dummy DropAsset implementation
  assets_data_cc_file
      << "void DropAsset(AssetId asset_id, const uint8_t *asset) {\n";
  assets_data_cc_file << "  (void)asset_id;\n";
  assets_data_cc_file << "  (void)asset;\n";
  assets_data_cc_file << "  // No-op for now, actual implementation for lazy "
                         "decompression goes here\n";
  assets_data_cc_file << "}\n";

  assets_data_cc_file.close();

  std::cout << "Asset packer successfully generated: " << output_assets_h_path
            << " and " << output_assets_data_cc_path << "\n";

  return 0;
}