#include "tool.hpp"
#include <CLI11.hpp>
#include <tree_sitter/api.h>

// Forward declarations
extern "C" TSLanguage *tree_sitter_cpp();

static std::string filepath;

// Helper function to dump Tree-sitter AST
void dumpTreeSitterAST(TSNode node, const char *source_code, int depth) {
  std::string indent(depth * 2, ' ');
  const char *type = ts_node_type(node);

  uint32_t start = ts_node_start_byte(node);
  uint32_t end = ts_node_end_byte(node);

  // Get the text content for leaf nodes or small nodes
  std::string content;
  if (end - start < 100) { // Only show content for small nodes
    content = extractNodeText(node, source_code);
    // Replace newlines with \n for better readability
    std::regex newline_regex("\n");
    content = std::regex_replace(content, newline_regex, "\\n");
    // Truncate if still too long
    if (content.length() > 50) {
      content = content.substr(0, 47) + "...";
    }
  }

  if (!content.empty()) {
    spdlog::info("{}{}[{}:{}] \"{}\"", indent, type, start, end, content);
  } else {
    spdlog::info("{}{}[{}:{}]", indent, type, start, end);
  }

  // Recursively dump children
  uint32_t child_count = ts_node_child_count(node);
  for (uint32_t i = 0; i < child_count; i++) {
    TSNode child = ts_node_child(node, i);
    dumpTreeSitterAST(child, source_code, depth + 1);
  }
}

bool dumpTreeFile(const std::string &filepath) {
  std::ifstream file(filepath);
  if (!file.is_open()) {
    spdlog::error("Could not open file {}", filepath);
    return false;
  }

  std::string file_content((std::istreambuf_iterator<char>(file)),
                           std::istreambuf_iterator<char>());

  TSParser *parser = ts_parser_new();
  ts_parser_set_language(parser, tree_sitter_cpp());

  TSTree *tree = ts_parser_parse_string(parser, nullptr, file_content.c_str(),
                                        file_content.length());
  TSNode root_node = ts_tree_root_node(tree);

  if (ts_node_is_null(root_node)) {
    spdlog::error("Failed to parse file {}", filepath);
    ts_tree_delete(tree);
    ts_parser_delete(parser);
    return false;
  }

  spdlog::info("=== Tree-sitter AST for {} ===", filepath);
  dumpTreeSitterAST(root_node, file_content.c_str());
  spdlog::info("=== End of AST dump ===");

  ts_tree_delete(tree);
  ts_parser_delete(parser);
  return true;
}

void register_cmd_dump(CLI::App &app) {
  auto cmd =
      app.add_subcommand("dump-tree", "Dump the tree-sitter AST for a file");
  cmd->add_option("file", filepath,
                 "Input C++ file to parse")->required();
  cmd->final_callback([]() {
    spdlog::info("=== Processing: {} ===", filepath);
    dumpTreeFile(filepath);
  });
}