#include "lua/persistence.h"

extern "C" {
#include <lua.h>
#include <lauxlib.h>
}

#include <wx/log.h>
#include <fstream>
#include <sstream>
#include <set>
#include <cstdio>

namespace wherigo {

// List of Wherigo game object classes to save
static const std::set<std::string> WHERIGO_CLASSES = {
  "ZCartridge", "Zone", "ZItem", "ZCharacter",
  "ZTask", "ZTimer", "ZInput", "ZMedia"
};

std::vector<std::string> LuaPersistence::getGameGlobals(lua_State* L) {
  std::vector<std::string> globals;

  lua_getglobal(L, "_G");
  lua_pushnil(L);

  while (lua_next(L, -2) != 0) {
    if (lua_isstring(L, -2)) {
      const char* key = lua_tostring(L, -2);

      // Check if it's a Wherigo game object
      if (lua_istable(L, -1)) {
        lua_getfield(L, -1, "ClassName");
        if (lua_isstring(L, -1)) {
          std::string className = lua_tostring(L, -1);
          if (WHERIGO_CLASSES.count(className) > 0) {
            globals.push_back(key);
          }
        }
        lua_pop(L, 1); // pop ClassName
      }
      // Also save simple global variables (strings, numbers, booleans)
      else if (lua_isnumber(L, -1) || lua_isstring(L, -1) || lua_isboolean(L, -1)) {
        std::string keyStr = key;
        // Skip internal Lua variables
        if (keyStr != "_G" && keyStr != "_VERSION" && keyStr.find("_") != 0) {
          globals.push_back(key);
        }
      }
    }
    lua_pop(L, 1); // pop value
  }
  lua_pop(L, 1); // pop _G

  // Always include Player object
  globals.push_back("Player");

  return globals;
}

void LuaPersistence::serializeValue(lua_State* L, int index, std::string& output, int depth) {
  if (depth > 10) {
    output += "nil"; // Prevent infinite recursion
    return;
  }

  int type = lua_type(L, index);

  switch (type) {
    case LUA_TNIL:
      output += "nil";
      break;

    case LUA_TBOOLEAN:
      output += lua_toboolean(L, index) ? "true" : "false";
      break;

    case LUA_TNUMBER:
      output += std::to_string(lua_tonumber(L, index));
      break;

    case LUA_TSTRING: {
      const char* str = lua_tostring(L, index);
      output += "\"";
      // Properly escape all special characters for Lua strings
      for (const char* p = str; *p; ++p) {
        unsigned char c = (unsigned char)*p;
        switch (c) {
          case '"':
            output += "\\\"";
            break;
          case '\\':
            output += "\\\\";
            break;
          case '\n':
            output += "\\n";
            break;
          case '\r':
            output += "\\r";
            break;
          case '\t':
            output += "\\t";
            break;
          case '\b':
            output += "\\b";
            break;
          case '\f':
            output += "\\f";
            break;
          default:
            // For control characters and high bytes, use decimal escape
            if (c < 32 || c >= 127) {
              char hex[6];
              snprintf(hex, sizeof(hex), "\\%d", c);
              output += hex;
            } else {
              output += c;
            }
            break;
        }
      }
      output += "\"";
      break;
    }

    case LUA_TTABLE:
      serializeTable(L, index, output, depth);
      break;

    default:
      // Functions, userdata, threads -> save as nil
      output += "nil";
      break;
  }
}

void LuaPersistence::serializeTable(lua_State* L, int index, std::string& output, int depth) {
  output += "{";

  // Normalize stack index
  if (index < 0) index = lua_gettop(L) + index + 1;

  bool first = true;
  lua_pushnil(L);

  while (lua_next(L, index) != 0) {
    if (!first) output += ",";
    first = false;

    // Serialize key
    output += "[";
    serializeValue(L, -2, output, depth + 1);
    output += "]=";

    // Serialize value
    serializeValue(L, -1, output, depth + 1);

    lua_pop(L, 1); // pop value, keep key for next iteration
  }

  output += "}";
}

bool LuaPersistence::saveGlobals(lua_State* L, const std::string& filePath,
                                 const std::vector<std::string>& globals) {
  std::string output = "-- Wherigo Save State\nreturn {\n";

  for (const auto& globalName : globals) {
    lua_getglobal(L, globalName.c_str());

    if (!lua_isnil(L, -1)) {
      output += "  [\"" + globalName + "\"] = ";
      serializeValue(L, -1, output, 0);
      output += ",\n";
    }

    lua_pop(L, 1);
  }

  output += "}\n";

  // Write to file
  std::ofstream file(filePath);
  if (!file.is_open()) {
    wxLogError("Failed to open save file: %s", filePath);
    return false;
  }

  file << output;
  file.close();

  wxLogDebug("Saved %zu globals to %s", globals.size(), filePath);
  return true;
}

bool LuaPersistence::saveState(lua_State* L, const std::string& filePath) {
  auto globals = getGameGlobals(L);
  return saveGlobals(L, filePath, globals);
}

bool LuaPersistence::loadGlobals(lua_State* L, const std::string& filePath) {
  // Load the file
  std::ifstream file(filePath);
  if (!file.is_open()) {
    wxLogError("Failed to open load file: %s", filePath);
    return false;
  }

  std::stringstream buffer;
  buffer << file.rdbuf();
  std::string content = buffer.str();
  file.close();

  // Execute the Lua file
  if (luaL_loadstring(L, content.c_str()) != 0) {
    wxLogError("Failed to parse save file: %s", lua_tostring(L, -1));
    lua_pop(L, 1);
    return false;
  }

  if (lua_pcall(L, 0, 1, 0) != 0) {
    wxLogError("Failed to execute save file: %s", lua_tostring(L, -1));
    lua_pop(L, 1);
    return false;
  }

  // Result should be a table
  if (!lua_istable(L, -1)) {
    wxLogError("Save file did not return a table");
    lua_pop(L, 1);
    return false;
  }

  // Restore globals - but for tables (Wherigo objects), merge instead of replace
  // Strategy: For each saved object, find the matching existing object and merge
  int count = 0;
  lua_pushnil(L);
  while (lua_next(L, -2) != 0) {
    if (lua_isstring(L, -2) && lua_istable(L, -1)) {
      const char* savedKey = lua_tostring(L, -2);
      wxLogDebug("Restoring global: %s", savedKey);
      int stackBefore = lua_gettop(L);

      // Get the Name field of the saved object (if it exists)
      lua_getfield(L, -1, "Name");
      std::string savedName;
      if (lua_isstring(L, -1)) {
        savedName = lua_tostring(L, -1);
      }
      lua_pop(L, 1);

      // Try to find matching existing object
      bool merged = false;

      if (!savedName.empty()) {
        // Search for existing object with same Name
        lua_getglobal(L, "_G");
        lua_pushnil(L);

        while (lua_next(L, -2) != 0) {
          if (lua_istable(L, -1)) {
            lua_getfield(L, -1, "Name");
            if (lua_isstring(L, -1)) {
              if (savedName == lua_tostring(L, -1)) {
                lua_pop(L, 1); // pop Name
                // Found matching object! Merge data
                lua_pushnil(L);
                while (lua_next(L, -3) != 0) {
                  if (lua_isstring(L, -2)) {
                    const char* field = lua_tostring(L, -2);
                    if (!lua_isfunction(L, -1)) {
                      lua_pushvalue(L, -1);
                      lua_setfield(L, -4, field);
                    }
                  }
                  lua_pop(L, 1);
                }
                merged = true;
                lua_pop(L, 2); // pop iterator and object
                break;
              }
            }
            lua_pop(L, 1);
          }
          lua_pop(L, 1);
        }
        lua_pop(L, 1); // pop _G
      }

      if (!merged) {
        // No match found or no Name field - use key matching
        lua_getglobal(L, savedKey);
        if (lua_istable(L, -1)) {
          // Object exists by key - merge it
          lua_pushnil(L);
          while (lua_next(L, -3) != 0) {
            if (lua_isstring(L, -2)) {
              const char* field = lua_tostring(L, -2);
              if (!lua_isfunction(L, -1)) {
                lua_getfield(L, -3, field); // get existing value
                if (lua_istable(L, -1) && lua_istable(L, -2)) {
                  // Both are tables: merge recursively
                  lua_pushnil(L);
                  while (lua_next(L, -3) != 0) {
                    if (lua_isstring(L, -2)) {
                      const char* subfield = lua_tostring(L, -2);
                      if (!lua_isfunction(L, -1)) {
                        lua_pushvalue(L, -1);
                        lua_setfield(L, -5, subfield); // merge into existing table
                      }
                    }
                    lua_pop(L, 1);
                  }
                  lua_pop(L, 1); // pop existing value
                } else {
                  lua_pop(L, 1); // pop existing value
                  lua_pushvalue(L, -1);
                  lua_setfield(L, -3, field);
                }
              }
            }
            lua_pop(L, 1);
          }
          lua_pop(L, 1); // pop existing object
          merged = true;
        } else {
          lua_pop(L, 1); // pop nil
        }
      }
      if (!merged) {
        // No existing object found - create new global
        lua_pushvalue(L, -1);
        lua_setglobal(L, savedKey);
      }
      int stackAfter = lua_gettop(L);
      if (stackBefore != stackAfter) {
        wxLogError("Stack imbalance detected for global %s: before=%d after=%d", savedKey, stackBefore, stackAfter);
      }
      count++;
    }
    lua_pop(L, 1); // pop value from lua_next
  }

  lua_pop(L, 1); // pop table

  lua_pop(L, 1); // pop table

  wxLogDebug("Restored %d globals from %s", count, filePath);
  return true;
}

bool LuaPersistence::loadState(lua_State* L, const std::string& filePath) {
  return loadGlobals(L, filePath);
}

} // namespace wherigo