#include <iostream>
#include <fstream>
#include <string>
#include <regex>
#include <filesystem>
#include <set>
#include <map>
#include <vector>
#include <algorithm>

namespace fs = std::filesystem;

int main(int argc, char *argv[]) {
    if (argc < 3) {
        std::cerr << "Usage: gen_owners <wolfram-include-dir> <output.hpp>\n";
        return 1;
    }

    fs::path inc_dir(argv[1]);
    fs::path out_path(argv[2]);

    // Match `void wf_<T>_free(wf_<T> *name)`. Backreference ensures type
    // matches function prefix. Excludes `wf_*_json_free(char*)` (strings).
    std::regex header_re(
        R"(void\s+wf_([A-Za-z0-9_]+)_free\s*\(\s*wf_\1\s*\*\s*\w+\s*\))");

    std::map<std::string, std::string> types; // type -> include
    std::set<std::string> includes;

    // Recursively scan all .h files (including oauth/, repo/ subdirectories)
    // in sorted order, mirroring the Python generator's sorted() glob.
    std::vector<fs::path> headers;
    for (const auto &entry : fs::recursive_directory_iterator(inc_dir)) {
        if (entry.path().extension() == ".h" &&
            entry.path().string().find("json-enhanced") == std::string::npos)
            headers.push_back(entry.path());
    }
    std::sort(headers.begin(), headers.end());

    for (const auto &path : headers) {
        std::ifstream file(path);
        std::string text((std::istreambuf_iterator<char>(file)),
                         std::istreambuf_iterator<char>());

        std::smatch match;
        std::string::const_iterator search_start(text.cbegin());
        while (std::regex_search(search_start, text.cend(), match, header_re)) {
            std::string t = match[1].str();
            if (types.find(t) == types.end()) {
                fs::path rel = fs::relative(path, inc_dir);
                std::string inc = "wolfram/" + rel.generic_string();
                types[t] = inc;
                includes.insert(inc);
            }
            search_start = match.suffix().first;
        }
    }

    std::vector<std::string> lines;
    lines.emplace_back("/* Generated by cpp/wolfram-cpp/tools/gen_owners.cpp; "
                       "do not edit. */");
    lines.emplace_back("#ifndef WOLFRAM_CPP_GENERATED_OWNERS_HPP");
    lines.emplace_back("#define WOLFRAM_CPP_GENERATED_OWNERS_HPP");
    lines.emplace_back("");

    for (const auto &inc : includes) {
        lines.emplace_back("#include <" + inc + ">");
    }
    lines.emplace_back("");
    lines.emplace_back("#include \"wolfram/unique_handle.hpp\"");
    lines.emplace_back("");
    lines.emplace_back("namespace wolfram {");
    lines.emplace_back("");

    std::vector<std::string> sorted_types(types.size());
    std::transform(types.begin(), types.end(), sorted_types.begin(),
                   [](const auto &p) { return p.first; });
    std::sort(sorted_types.begin(), sorted_types.end());

    for (const auto &t : sorted_types) {
        lines.emplace_back("using wf_" + t + "_handle = unique_handle<wf_" + t +
                           ", wf_" + t + "_free>;");
    }
    lines.emplace_back("");
    lines.emplace_back("} // namespace wolfram");
    lines.emplace_back("");
    lines.emplace_back("#endif");

    std::ofstream out(out_path);
    for (const auto &line : lines) {
        out << line << "\n";
    }

    std::cout << "wrote " << out_path << ": " << types.size()
              << " owned handles from " << includes.size() << " headers\n";
    return 0;
}