/**
 * chi_refresh_repo - ChiMod Autogeneration Utility
 *
 * This utility automatically generates libexec source files (.cc) and methods
 * header files (.h) for ChiMods based on their chimaera_mod.yaml configuration files.
 *
 * The libexec source files implement Container virtual API methods (Run, Del,
 * SaveTask, LoadTask, NewCopy) with switch-case dispatch.
 *
 * Usage:
 *     chi_refresh_repo <chimod_repo_path>
 */

#include <algorithm>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <string>
#include <vector>

#include <yaml-cpp/yaml.h>

namespace fs = std::filesystem;

/**
 * Structure to hold method information
 */
struct Method {
  std::string constant_name;
  std::string method_name;
  int method_id;
  bool is_inherited;
};

/**
 * ChiMod Generator class
 * Handles generation of autogenerated files for ChiMods
 */
class ChiModGenerator {
 private:
  fs::path repo_path_;
  fs::path repo_yaml_path_;
  std::string repo_namespace_;

 public:
  explicit ChiModGenerator(const std::string& repo_path)
      : repo_path_(fs::absolute(repo_path)),
        repo_yaml_path_(repo_path_ / "chimaera_repo.yaml") {
    if (!fs::exists(repo_yaml_path_)) {
      throw std::runtime_error("Repository YAML not found: " + repo_yaml_path_.string());
    }

    // Load repository namespace from chimaera_repo.yaml
    YAML::Node repo_config = LoadRepoConfig();
    repo_namespace_ = repo_config["namespace"] ? repo_config["namespace"].as<std::string>() : "chimaera";
  }

  /**
   * Load the repository configuration
   */
  YAML::Node LoadRepoConfig() {
    try {
      return YAML::LoadFile(repo_yaml_path_.string());
    } catch (const YAML::Exception& e) {
      throw std::runtime_error("Failed to parse repository YAML: " + std::string(e.what()));
    }
  }

  /**
   * Load a specific ChiMod configuration
   */
  YAML::Node LoadChiModConfig(const std::string& chimod_name) {
    fs::path chimod_yaml_path = repo_path_ / chimod_name / "chimaera_mod.yaml";
    if (!fs::exists(chimod_yaml_path)) {
      throw std::runtime_error("ChiMod YAML not found: " + chimod_yaml_path.string());
    }

    try {
      return YAML::LoadFile(chimod_yaml_path.string());
    } catch (const YAML::Exception& e) {
      throw std::runtime_error("Failed to parse ChiMod YAML: " + std::string(e.what()));
    }
  }

  /**
   * Extract methods from configuration, filtering out those with value -1
   */
  std::vector<Method> GetMethods(const YAML::Node& config) {
    std::vector<Method> methods;

    // Process all key-value pairs in the config
    for (auto it = config.begin(); it != config.end(); ++it) {
      std::string key = it->first.as<std::string>();
      
      // Skip non-method keys
      if (key.length() == 0 || key[0] != 'k' || !it->second.IsScalar()) {
        continue;
      }

      // Try to parse as integer
      int value;
      try {
        value = it->second.as<int>();
      } catch (const YAML::Exception&) {
        continue; // Skip non-integer values
      }

      // Skip unimplemented methods (value -1)
      if (value == -1) {
        continue;
      }

      std::string method_name = key.substr(1); // Remove the 'k' prefix
      Method method;
      method.constant_name = key;
      method.method_name = method_name;
      method.method_id = value;
      method.is_inherited = value < 10;  // Inherited methods have IDs < 10
      methods.push_back(method);
    }

    // Sort by method ID
    std::sort(methods.begin(), methods.end(),
              [](const Method& a, const Method& b) {
                return a.method_id < b.method_id;
              });

    return methods;
  }

  /**
   * Generate the methods header file
   */
  std::string GenerateMethodsHeader(const std::string& chimod_name, const YAML::Node& config,
                                     const std::string& repo_namespace) {
    auto methods = GetMethods(config);
    std::string namespace_name = repo_namespace;
    std::string module_name = config["module_name"] ? config["module_name"].as<std::string>() : chimod_name;

    std::ostringstream oss;
    std::string chimod_upper;
    std::transform(chimod_name.begin(), chimod_name.end(), std::back_inserter(chimod_upper),
                   [](unsigned char c) { return std::toupper(c); });

    // Build the header content
    oss << "#ifndef " << chimod_upper << "_AUTOGEN_METHODS_H_\n";
    oss << "#define " << chimod_upper << "_AUTOGEN_METHODS_H_\n";
    oss << "\n";
    oss << "#include <chimaera/chimaera.h>\n";
    oss << "\n";
    oss << "/**\n";
    oss << " * Auto-generated method definitions for " << module_name << "\n";
    oss << " */\n";
    oss << "\n";
    oss << "namespace " << namespace_name << "::" << chimod_name << " {\n";
    oss << "\n";
    oss << "namespace Method {\n";

    // Add inherited methods
    std::vector<Method> inherited_methods;
    std::copy_if(methods.begin(), methods.end(), std::back_inserter(inherited_methods),
                 [](const Method& m) { return m.is_inherited; });

    if (!inherited_methods.empty()) {
      oss << "// Inherited methods\n";
      for (const auto& method : inherited_methods) {
        oss << "GLOBAL_CONST chi::u32 " << method.constant_name << " = " << method.method_id << ";\n";
      }
      oss << "\n";
    }

    // Add custom methods
    std::vector<Method> custom_methods;
    std::copy_if(methods.begin(), methods.end(), std::back_inserter(custom_methods),
                 [](const Method& m) { return !m.is_inherited; });

    if (!custom_methods.empty()) {
      oss << "// " << module_name << "-specific methods\n";
      for (const auto& method : custom_methods) {
        oss << "GLOBAL_CONST chi::u32 " << method.constant_name << " = " << method.method_id << ";\n";
      }
    }

    oss << "}  // namespace Method\n";
    oss << "\n";
    oss << "}  // namespace " << namespace_name << "::" << chimod_name << "\n";
    oss << "\n";
    oss << "#endif  // " << chimod_upper << "_AUTOGEN_METHODS_H_\n";

    return oss.str();
  }

  /**
   * Get the task type name for a method
   */
  std::string GetTaskTypeName(const std::string& method_name, const std::string& chimod_name) {
    // Special cases for common inherited methods
    if (method_name == "Create") {
      return "CreateTask";
    } else if (method_name == "Destroy") {
      return "DestroyTask";
    } else if (method_name.length() >= 11 && method_name.substr(0, 11) == "GetOrCreate") {
      // Template-based task
      return chimod_name + "::" + method_name + "Task<" + chimod_name + "::CreateParams>";
    }
    
    // Regular task (default case)
    return method_name + "Task";
  }

  /**
   * Generate the lib_exec source file (.cc) with Container virtual API implementations
   */
  std::string GenerateLibExecSource(const std::string& chimod_name, const YAML::Node& config,
                                     const std::string& repo_namespace) {
    auto methods = GetMethods(config);
    std::string namespace_name = repo_namespace;
    std::string module_name = config["module_name"] ? config["module_name"].as<std::string>() : chimod_name;

    std::ostringstream oss;

    // Build the source content
    oss << "/**\n";
    oss << " * Auto-generated execution implementation for " << module_name << " ChiMod\n";
    oss << " * Implements Container virtual APIs (Run, Del, SaveTask, LoadTask, NewCopy, Aggregate)\n";
    oss << " * using switch-case dispatch\n";
    oss << " * \n";
    oss << " * This file is autogenerated - do not edit manually.\n";
    oss << " * Changes should be made to the autogen tool or the YAML configuration.\n";
    oss << " */\n";
    oss << "\n";
    oss << "#include \"" << namespace_name << "/" << module_name << "/" << chimod_name << "_runtime.h\"\n";
    oss << "#include \"" << namespace_name << "/" << module_name << "/autogen/" << chimod_name << "_methods.h\"\n";
    oss << "#include <chimaera/chimaera.h>\n";
    oss << "\n";
    oss << "namespace " << namespace_name << "::" << chimod_name << " {\n";
    oss << "\n";
    oss << "//==============================================================================\n";
    oss << "// Container Virtual API Implementations\n";
    oss << "//==============================================================================\n";
    oss << "\n";
    oss << "void Runtime::Init(const chi::PoolId &pool_id, const std::string &pool_name,\n";
    oss << "                   chi::u32 container_id) {\n";
    oss << "  // Call base class initialization\n";
    oss << "  chi::Container::Init(pool_id, pool_name, container_id);\n";
    oss << "\n";
    oss << "  // Initialize the client for this ChiMod\n";
    oss << "  client_ = Client(pool_id);\n";
    oss << "}\n";
    oss << "\n";
    oss << "void Runtime::Run(chi::u32 method, hipc::FullPtr<chi::Task> task_ptr, chi::RunContext& rctx) {\n";
    oss << "  switch (method) {\n";

    // Add Run switch cases for each method
    for (const auto& method : methods) {
      std::string task_type = GetTaskTypeName(method.method_name, chimod_name);
      oss << "    case Method::" << method.constant_name << ": {\n";
      oss << "      " << method.method_name << "(task_ptr.Cast<" << task_type << ">(), rctx);\n";
      oss << "      break;\n";
      oss << "    }\n";
    }

    oss << "    default: {\n";
    oss << "      // Unknown method - do nothing\n";
    oss << "      break;\n";
    oss << "    }\n";
    oss << "  }\n";
    oss << "}\n";
    oss << "\n";
    oss << "void Runtime::Del(chi::u32 method, hipc::FullPtr<chi::Task> task_ptr) {\n";
    oss << "  // Use IPC manager to deallocate task from shared memory\n";
    oss << "  auto* ipc_manager = CHI_IPC;\n";
    oss << "  \n";
    oss << "  switch (method) {\n";

    // Add Del switch cases
    for (const auto& method : methods) {
      std::string task_type = GetTaskTypeName(method.method_name, chimod_name);
      oss << "    case Method::" << method.constant_name << ": {\n";
      oss << "      ipc_manager->DelTask(task_ptr.Cast<" << task_type << ">());\n";
      oss << "      break;\n";
      oss << "    }\n";
    }

    oss << "    default: {\n";
    oss << "      // For unknown methods, still try to delete from main segment\n";
    oss << "      ipc_manager->DelTask(task_ptr);\n";
    oss << "      break;\n";
    oss << "    }\n";
    oss << "  }\n";
    oss << "}\n";
    oss << "\n";
    oss << "void Runtime::SaveTask(chi::u32 method, chi::SaveTaskArchive& archive, \n";
    oss << "                        hipc::FullPtr<chi::Task> task_ptr) {\n";
    oss << "  switch (method) {\n";

    // Add SaveTask switch cases
    for (const auto& method : methods) {
      std::string task_type = GetTaskTypeName(method.method_name, chimod_name);
      oss << "    case Method::" << method.constant_name << ": {\n";
      oss << "      auto typed_task = task_ptr.Cast<" << task_type << ">();\n";
      oss << "      archive << *typed_task;\n";
      oss << "      break;\n";
      oss << "    }\n";
    }

    oss << "    default: {\n";
    oss << "      // Unknown method - do nothing\n";
    oss << "      break;\n";
    oss << "    }\n";
    oss << "  }\n";
    oss << "}\n";
    oss << "\n";
    oss << "void Runtime::LoadTask(chi::u32 method, chi::LoadTaskArchive& archive, \n";
    oss << "                        hipc::FullPtr<chi::Task>& task_ptr) {\n";
    oss << "  auto* ipc_manager = CHI_IPC;\n";
    oss << "  \n";
    oss << "  switch (method) {\n";

    // Add LoadTask switch cases
    for (const auto& method : methods) {
      std::string task_type = GetTaskTypeName(method.method_name, chimod_name);
      oss << "    case Method::" << method.constant_name << ": {\n";
      oss << "      // Allocate task using typed NewTask if not already allocated\n";
      oss << "      if (task_ptr.IsNull()) {\n";
      oss << "        task_ptr = ipc_manager->NewTask<" << task_type << ">().template Cast<chi::Task>();\n";
      oss << "      }\n";
      oss << "      auto typed_task = task_ptr.Cast<" << task_type << ">();\n";
      oss << "      archive >> *typed_task;\n";
      oss << "      break;\n";
      oss << "    }\n";
    }

    oss << "    default: {\n";
    oss << "      // Unknown method - do nothing\n";
    oss << "      break;\n";
    oss << "    }\n";
    oss << "  }\n";
    oss << "}\n";
    oss << "\n";
    oss << "void Runtime::NewCopy(chi::u32 method, const hipc::FullPtr<chi::Task>& orig_task,\n";
    oss << "                       hipc::FullPtr<chi::Task>& dup_task, bool deep) {\n";
    oss << "  auto* ipc_manager = CHI_IPC;\n";
    oss << "  if (!ipc_manager) {\n";
    oss << "    return;\n";
    oss << "  }\n";
    oss << "  \n";
    oss << "  switch (method) {\n";

    // Add NewCopy switch cases
    for (const auto& method : methods) {
      std::string task_type = GetTaskTypeName(method.method_name, chimod_name);
      oss << "    case Method::" << method.constant_name << ": {\n";
      oss << "      // Allocate new task using SHM default constructor\n";
      oss << "      auto typed_task = ipc_manager->NewTask<" << task_type << ">();\n";
      oss << "      if (!typed_task.IsNull()) {\n";
      oss << "        // Copy base Task fields first\n";
      oss << "        typed_task.template Cast<chi::Task>()->Copy(orig_task);\n";
      oss << "        // Then copy task-specific fields\n";
      oss << "        typed_task->Copy(orig_task.Cast<" << task_type << ">());\n";
      oss << "        // Cast to base Task type for return\n";
      oss << "        dup_task = typed_task.template Cast<chi::Task>();\n";
      oss << "      }\n";
      oss << "      break;\n";
      oss << "    }\n";
    }

    oss << "    default: {\n";
    oss << "      // For unknown methods, create base Task copy\n";
    oss << "      auto typed_task = ipc_manager->NewTask<chi::Task>();\n";
    oss << "      if (!typed_task.IsNull()) {\n";
    oss << "        typed_task->Copy(orig_task);\n";
    oss << "        dup_task = typed_task;  // Already chi::Task type\n";
    oss << "      }\n";
    oss << "      break;\n";
    oss << "    }\n";
    oss << "  }\n";
    oss << "  \n";
    oss << "  (void)deep;    // Deep copy parameter reserved for future use\n";
    oss << "}\n";
    oss << "\n";
    oss << "void Runtime::Aggregate(chi::u32 method, hipc::FullPtr<chi::Task> origin_task,\n";
    oss << "                         hipc::FullPtr<chi::Task> replica_task) {\n";
    oss << "  switch (method) {\n";

    // Add Aggregate switch cases
    for (const auto& method : methods) {
      std::string task_type = GetTaskTypeName(method.method_name, chimod_name);
      oss << "    case Method::" << method.constant_name << ": {\n";
      oss << "      auto typed_origin = origin_task.Cast<" << task_type << ">();\n";
      oss << "      auto typed_replica = replica_task.Cast<" << task_type << ">();\n";
      oss << "      // Call base Task aggregate to propagate return codes\n";
      oss << "      origin_task->Aggregate(replica_task);\n";
      oss << "      // Use SFINAE-based macro to call task-specific Aggregate if available, otherwise Copy\n";
      oss << "      CHI_AGGREGATE_OR_COPY(typed_origin, typed_replica);\n";
      oss << "      break;\n";
      oss << "    }\n";
    }

    oss << "    default: {\n";
    oss << "      // For unknown methods, use base Task Aggregate (which also propagates return codes)\n";
    oss << "      origin_task->Aggregate(replica_task);\n";
    oss << "      break;\n";
    oss << "    }\n";
    oss << "  }\n";
    oss << "}\n";
    oss << "\n";
    oss << "} // namespace " << namespace_name << "::" << chimod_name << "\n";

    return oss.str();
  }

  /**
   * Generate autogen files for a single ChiMod
   */
  void GenerateChiModFiles(const std::string& chimod_name) {
    std::cout << "Generating files for ChiMod: " << chimod_name << std::endl;

    // Load ChiMod configuration
    YAML::Node config;
    try {
      config = LoadChiModConfig(chimod_name);
    } catch (const std::runtime_error& e) {
      std::cout << "Warning: " << e.what() << std::endl;
      return;
    }

    // Create include autogen directory for methods header
    // Structure: [chimod_directory]/include/[namespace]/[module_name]/autogen/
    std::string module_name = config["module_name"] ? config["module_name"].as<std::string>() : chimod_name;
    fs::path include_autogen_dir = repo_path_ / chimod_name / "include" / repo_namespace_ / module_name / "autogen";
    fs::create_directories(include_autogen_dir);

    // Create src autogen directory for lib_exec source
    fs::path src_autogen_dir = repo_path_ / chimod_name / "src" / "autogen";
    fs::create_directories(src_autogen_dir);

    // Generate methods header
    std::string methods_content = GenerateMethodsHeader(chimod_name, config, repo_namespace_);
    fs::path methods_file = include_autogen_dir / (chimod_name + "_methods.h");
    std::ofstream methods_stream(methods_file);
    if (!methods_stream) {
      throw std::runtime_error("Failed to create methods file: " + methods_file.string());
    }
    methods_stream << methods_content;
    methods_stream.close();
    std::cout << "  Generated: " << methods_file << std::endl;

    // Generate lib_exec source file
    std::string lib_exec_content = GenerateLibExecSource(chimod_name, config, repo_namespace_);
    fs::path lib_exec_file = src_autogen_dir / (chimod_name + "_lib_exec.cc");
    std::ofstream lib_exec_stream(lib_exec_file);
    if (!lib_exec_stream) {
      throw std::runtime_error("Failed to create lib_exec file: " + lib_exec_file.string());
    }
    lib_exec_stream << lib_exec_content;
    lib_exec_stream.close();
    std::cout << "  Generated: " << lib_exec_file << std::endl;
  }

  /**
   * Refresh all ChiMods in the repository
   */
  void RefreshRepo() {
    YAML::Node repo_config = LoadRepoConfig();
    
    if (!repo_config["modules"] || !repo_config["modules"].IsSequence()) {
      throw std::runtime_error("No modules found in repository configuration");
    }

    std::vector<std::string> modules;
    for (const auto& module : repo_config["modules"]) {
      modules.push_back(module.as<std::string>());
    }

    std::cout << "Refreshing ChiMod repository: " << repo_path_ << std::endl;
    std::cout << "Found " << modules.size() << " modules: ";
    for (size_t i = 0; i < modules.size(); ++i) {
      if (i > 0) std::cout << ", ";
      std::cout << modules[i];
    }
    std::cout << std::endl << std::endl;

    for (const auto& chimod_name : modules) {
      GenerateChiModFiles(chimod_name);
      std::cout << std::endl;
    }

    std::cout << "Repository refresh complete!" << std::endl;
  }
};

void PrintUsage(const char* program_name) {
  std::cout << "Usage: " << program_name << " <chimod_repo_path>" << std::endl;
  std::cout << std::endl;
  std::cout << "Autogenerate ChiMod libexec and methods files" << std::endl;
  std::cout << std::endl;
  std::cout << "Examples:" << std::endl;
  std::cout << "  " << program_name << " /path/to/chimods" << std::endl;
  std::cout << "  " << program_name << " ." << std::endl;
}

int main(int argc, char* argv[]) {
  if (argc != 2) {
    PrintUsage(argv[0]);
    return 1;
  }

  std::string repo_path = argv[1];

  try {
    ChiModGenerator generator(repo_path);
    generator.RefreshRepo();
    return 0;
  } catch (const std::exception& e) {
    std::cerr << "Error: " << e.what() << std::endl;
    return 1;
  }
}