// 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 <algorithm> // For std::count
#include <cstring>   // For std::memcpy
#include <regex>     // For std::regex
#include <stdexcept> // For std::stof exceptions
#include <string>
#include <vector>
#include <cmath>

#define STB_IMAGE_IMPLEMENTATION
#define STBI_NO_LINEAR // Don't apply gamma correction, we want raw bytes
#define STBI_ONLY_PNG
#define STBI_ONLY_JPEG
#define STBI_ONLY_TGA
#define STBI_ONLY_BMP
#include "stb_image.h"

#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},
};

static bool HasImageExtension(const std::string& filename) {
  std::string ext = filename.substr(filename.find_last_of(".") + 1);
  // simple case-insensitive check (assuming lowercase for simplicity or just
  // basic checks)
  if (ext == "png" || ext == "jpg" || ext == "jpeg" || ext == "tga" ||
      ext == "bmp")
    return true;
  return false;
}

static bool HasMeshExtension(const std::string& filename) {
  std::string ext = filename.substr(filename.find_last_of(".") + 1);
  return ext == "obj";
}

// 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
};

struct Vec3 {
  float x, y, z;
  Vec3 operator+(const Vec3& o) const { return {x + o.x, y + o.y, z + o.z}; }
  Vec3 operator+=(const Vec3& o) { x += o.x; y += o.y; z += o.z; return *this; }
  Vec3 operator-(const Vec3& o) const { return {x - o.x, y - o.y, z - o.z}; }
  Vec3 operator*(float s) const { return {x * s, y * s, z * s}; }
  static Vec3 cross(const Vec3& a, const Vec3& b) {
    return {a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x};
  }
  Vec3 normalize() const {
    float len = std::sqrt(x * x + y * y + z * z);
    if (len > 1e-6f) return {x / len, y / len, z / len};
    return {0, 0, 0};
  }
};

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",
          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,
                  "Warning: Unknown procedural function: %s for asset: %s "
                  "(Runtime error will occur)\n",
                  info.proc_func_name.c_str(), info.name.c_str());
          // return 1; // Allow unknown functions for testing runtime handling
        }
      }

      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
  fprintf(assets_h_file,
          "\n// Accessors to avoid static initialization order issues\n");
  fprintf(assets_h_file, "const struct AssetRecord* GetAssetRecordTable();\n");
  fprintf(assets_h_file, "size_t GetAssetCount();\n");

  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::string full_path = base_dir + info.filename;

      std::vector<uint8_t> buffer;
      bool is_image = HasImageExtension(info.filename);
      bool is_mesh = HasMeshExtension(info.filename);

      if (is_image) {
        int w, h, channels;
        unsigned char* img_data = stbi_load(
            full_path.c_str(), &w, &h, &channels, 4); // Force 4 channels (RGBA)
        if (!img_data) {
          fprintf(stderr, "Error: Could not load image file: %s (Reason: %s)\n",
                  full_path.c_str(), stbi_failure_reason());
          return 1;
        }

        // Format: [Width(4)][Height(4)][Pixels...]
        buffer.resize(sizeof(uint32_t) * 2 + w * h * 4);
        uint32_t* header = reinterpret_cast<uint32_t*>(buffer.data());
        header[0] = (uint32_t)w;
        header[1] = (uint32_t)h;
        std::memcpy(buffer.data() + sizeof(uint32_t) * 2, img_data, w * h * 4);

        stbi_image_free(img_data);
        printf("Processed image asset %s: %dx%d RGBA\n", info.name.c_str(), w,
               h);
      } else if (is_mesh) {
        std::ifstream obj_file(full_path);
        if (!obj_file.is_open()) {
          fprintf(stderr, "Error: Could not open mesh file: %s\n",
                  full_path.c_str());
          return 1;
        }

        std::vector<float> v_pos;
        std::vector<float> v_norm;
        std::vector<float> v_uv;

        struct RawFace {
          int v[3];
          int vt[3];
          int vn[3];
        };
        std::vector<RawFace> raw_faces;

        struct Vertex {
          float p[3], n[3], u[2];
        };
        std::vector<Vertex> final_vertices;
        std::vector<uint32_t> final_indices;
        std::map<std::string, uint32_t> vertex_map;

        std::string obj_line;
        while (std::getline(obj_file, obj_line)) {
          if (obj_line.compare(0, 2, "v ") == 0) {
            float x, y, z;
            std::sscanf(obj_line.c_str(), "v %f %f %f", &x, &y, &z);
            v_pos.push_back(x);
            v_pos.push_back(y);
            v_pos.push_back(z);
          } else if (obj_line.compare(0, 3, "vn ") == 0) {
            float x, y, z;
            std::sscanf(obj_line.c_str(), "vn %f %f %f", &x, &y, &z);
            v_norm.push_back(x);
            v_norm.push_back(y);
            v_norm.push_back(z);
          } else if (obj_line.compare(0, 3, "vt ") == 0) {
            float u, v;
            std::sscanf(obj_line.c_str(), "vt %f %f", &u, &v);
            v_uv.push_back(u);
            v_uv.push_back(v);
          } else if (obj_line.compare(0, 2, "f ") == 0) {
            char s1[64], s2[64], s3[64];
            if (std::sscanf(obj_line.c_str(), "f %s %s %s", s1, s2, s3) == 3) {
              std::string parts[3] = {s1, s2, s3};
              RawFace face = {};
              for (int i = 0; i < 3; ++i) {
                int v_idx = 0, vt_idx = 0, vn_idx = 0;
                if (parts[i].find("//") != std::string::npos) {
                  std::sscanf(parts[i].c_str(), "%d//%d", &v_idx, &vn_idx);
                } else if (std::count(parts[i].begin(), parts[i].end(), '/') ==
                           2) {
                  std::sscanf(parts[i].c_str(), "%d/%d/%d", &v_idx, &vt_idx,
                              &vn_idx);
                } else if (std::count(parts[i].begin(), parts[i].end(), '/') ==
                           1) {
                  std::sscanf(parts[i].c_str(), "%d/%d", &v_idx, &vt_idx);
                } else {
                  std::sscanf(parts[i].c_str(), "%d", &v_idx);
                }
                face.v[i] = v_idx;
                face.vt[i] = vt_idx;
                face.vn[i] = vn_idx;
              }
              raw_faces.push_back(face);
            }
          }
        }

        // Generate normals if missing
        if (v_norm.empty() && !v_pos.empty()) {
          printf("Generating normals for %s...\n", info.name.c_str());
          std::vector<Vec3> temp_normals(v_pos.size() / 3, {0, 0, 0});
          for (auto& face : raw_faces) {
            // Indices are 1-based in OBJ
            int idx0 = face.v[0] - 1;
            int idx1 = face.v[1] - 1;
            int idx2 = face.v[2] - 1;
            
            if (idx0 >= 0 && idx1 >= 0 && idx2 >= 0) {
              Vec3 p0 = {v_pos[idx0 * 3], v_pos[idx0 * 3 + 1], v_pos[idx0 * 3 + 2]};
              Vec3 p1 = {v_pos[idx1 * 3], v_pos[idx1 * 3 + 1], v_pos[idx1 * 3 + 2]};
              Vec3 p2 = {v_pos[idx2 * 3], v_pos[idx2 * 3 + 1], v_pos[idx2 * 3 + 2]};
              
              Vec3 normal = Vec3::cross(p1 - p0, p2 - p0).normalize();
              temp_normals[idx0] += normal;
              temp_normals[idx1] += normal;
              temp_normals[idx2] += normal;
            }
          }
          
          for (const auto& n : temp_normals) {
            Vec3 normalized = n.normalize();
            v_norm.push_back(normalized.x);
            v_norm.push_back(normalized.y);
            v_norm.push_back(normalized.z);
          }
          
          // Assign generated normals to faces
          for (auto& face : raw_faces) {
            face.vn[0] = face.v[0];
            face.vn[1] = face.v[1];
            face.vn[2] = face.v[2];
          }
        }

        // Build final vertices
        for (const auto& face : raw_faces) {
          for (int i = 0; i < 3; ++i) {
            // Reconstruct key string for uniqueness check
            char key_buf[128];
            std::sprintf(key_buf, "%d/%d/%d", face.v[i], face.vt[i], face.vn[i]);
            std::string key = key_buf;

            if (vertex_map.find(key) == vertex_map.end()) {
              vertex_map[key] = (uint32_t)final_vertices.size();
              
              Vertex v = {};
              if (face.v[i] > 0) {
                v.p[0] = v_pos[(face.v[i] - 1) * 3];
                v.p[1] = v_pos[(face.v[i] - 1) * 3 + 1];
                v.p[2] = v_pos[(face.v[i] - 1) * 3 + 2];
              }
              if (face.vn[i] > 0) {
                v.n[0] = v_norm[(face.vn[i] - 1) * 3];
                v.n[1] = v_norm[(face.vn[i] - 1) * 3 + 1];
                v.n[2] = v_norm[(face.vn[i] - 1) * 3 + 2];
              }
              if (face.vt[i] > 0) {
                v.u[0] = v_uv[(face.vt[i] - 1) * 2];
                v.u[1] = v_uv[(face.vt[i] - 1) * 2 + 1];
              }
              final_vertices.push_back(v);
            }
            final_indices.push_back(vertex_map[key]);
          }
        }

        // Format: [num_vertices(u32)][Vertex * num_vertices][num_indices(u32)][uint32_t
        // * num_indices]
        buffer.resize(sizeof(uint32_t) + final_vertices.size() * sizeof(Vertex) +
                      sizeof(uint32_t) +
                      final_indices.size() * sizeof(uint32_t));
        uint8_t* out_ptr = buffer.data();
        *reinterpret_cast<uint32_t*>(out_ptr) = (uint32_t)final_vertices.size();
        out_ptr += sizeof(uint32_t);
        std::memcpy(out_ptr, final_vertices.data(),
                    final_vertices.size() * sizeof(Vertex));
        out_ptr += final_vertices.size() * sizeof(Vertex);
        *reinterpret_cast<uint32_t*>(out_ptr) = (uint32_t)final_indices.size();
        out_ptr += sizeof(uint32_t);
        std::memcpy(out_ptr, final_indices.data(),
                    final_indices.size() * sizeof(uint32_t));

        printf("Processed mesh asset %s: %zu vertices, %zu indices\n",
               info.name.c_str(), final_vertices.size(), final_indices.size());
      } else {
        std::ifstream asset_file(full_path, std::ios::binary);
        if (!asset_file.is_open()) {
          fprintf(stderr, "Error: Could not open asset file: %s\n",
                  full_path.c_str());
          return 1;
        }
        buffer.assign((std::istreambuf_iterator<char>(asset_file)),
                      std::istreambuf_iterator<char>());
      }

      size_t original_size = buffer.size();
      buffer.push_back(0); // Null terminator for safety

      fprintf(assets_data_cc_file, "const size_t ASSET_SIZE_%s = %zu;\n",
              info.name.c_str(), original_size);
      fprintf(assets_data_cc_file,
              "alignas(16) 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, "const AssetRecord* GetAssetRecordTable() {\n");
  fprintf(assets_data_cc_file, "  static const AssetRecord 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, ASSET_SIZE_%s, false, nullptr, nullptr, 0",
              info.data_array_name.c_str(), info.name.c_str());
    }
    fprintf(assets_data_cc_file, " },\n");
  }
  fprintf(assets_data_cc_file, "  };\n");
  fprintf(assets_data_cc_file, "  return assets;\n");
  fprintf(assets_data_cc_file, "}\n\n");

  fprintf(assets_data_cc_file, "size_t GetAssetCount() {\n");
  fprintf(assets_data_cc_file, "  return %zu;\n", asset_build_infos.size());
  fprintf(assets_data_cc_file, "}\n\n");

  std::fclose(assets_data_cc_file);

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

  return 0;
}