#include "main.hh"

using pointer_t = std::byte* const;
struct pointer {
    const pointer_t address;
    const std::byte* const deref;
    std::shared_ptr<struct pointer> parent;
    const size_t start_offset;
    const size_t total_offset;
};
using pointers_t = std::vector<std::shared_ptr<struct pointer>>;

static pointers_t make_pointers(const memory::buffer_t& memory,
                                const pointers_t& targets,
                                const void* const start, const int tolerance,
                                const size_t total_offset) {

    pointers_t ret{};
    for (auto i = 0ul; i < std::size(memory) / sizeof(pointer_t); ++i) {
        const auto byte_offset = i * sizeof(pointer_t);

        const std::byte* const deref = [&]() {
            std::byte* ret;
            std::memcpy(&ret, &memory[byte_offset], sizeof(void*));
            return ret;
        }();

        // -1 is 0xff...etc
        if (deref == nullptr || deref > (pointer_t)(-1) - tolerance) {
            continue;
        }

        const auto lower = std::lower_bound(
            std::begin(targets), std::end(targets), deref,
            [](const auto& target, const std::byte* const& deref) {
                return target->address < deref;
            });
        const auto upper = std::upper_bound(
            lower, std::end(targets), deref + tolerance,
            [](const std::byte* const& deref, const auto& target) {
                return deref < target->address;
            });

        std::for_each(lower, upper, [&](const auto& target) {
            const auto address = (std::byte*)start + byte_offset;
            if (address == target->address || deref == target->deref) {
                return;
            }
            ret.emplace_back(std::make_shared<struct pointer>(
                pointer{.address = address,
                        .deref = deref,
                        .parent = target,
                        .start_offset = byte_offset,
                        .total_offset = total_offset + byte_offset}));
        });
    }
    return ret;
}

static pointers_t make_pointer_chain(const char* const binary_name,
                                     const pointer_t target, const int depth,
                                     const int tolerance) {
    const memory::maps maps{binary_name};
    const pid_t pid = std::stoi(maps.get_pid());

    pointers_t ret{std::make_shared<struct pointer>(
        pointer{.address = target, .parent = nullptr})};
    for (int i = 0; i < depth; ++i) {

        size_t total_offset = 0;
        pointers_t pointers{};
        for (const auto& map : maps) {
            const auto size = static_cast<std::size_t>((std::byte*)map.end -
                                                       (std::byte*)map.start);
            total_offset += size;
            if (!map.is_read() || !map.is_write() || !map.is_protected()) {
                continue;
            }

            const memory::buffer_t memory =
                memory::read_buffer(size, pid, map.start);

            std::ranges::move(make_pointers(memory, ret, map.start, tolerance,
                                            total_offset - size),
                              std::back_inserter(pointers));
        }

        std::cout << "depth: " << i + 1 << " = " << std::size(pointers)
                  << " chains...\n";

        if (std::size(pointers) == 0) {
            std::cout << "warning: only discovered chains of depth " << i
                      << '\n';
            break;
        }

        ret = std::move(pointers);
    }
    return ret;
}

static void do_pointer_scanner(const char* const binary_name,
                               const pointer_t target, const int depth,
                               const int tolerance) {

    const auto pointer_chains =
        make_pointer_chain(binary_name, target, depth, tolerance);

    for (const auto& chain : pointer_chains) {
        std::cout << std::hex << chain->total_offset << " = " << chain->address
                  << std::dec;
        const std::byte* prev = chain->deref;
        for (const decltype(chain->parent)* p = &chain->parent; p != nullptr;
             p = &(*p)->parent) {
            if (*p == nullptr) {
                break;
            }

            if ((*p)->address != nullptr) {
                std::cout << " ] ";
            }
            if (prev != nullptr) {
                std::cout << ((*p)->address - prev);
            }
            prev = (*p)->deref;

            if ((*p)->address == nullptr) {
                break;
            }
        }
        std::cout << '\n';
    }
}

int main(const int argc, const char* const argv[]) {
    if (argc < 3 || argc > 5) {
        std::cout << "usage: " << argv[0]
                  << " BINARY_NAME 0xADDRESS DEPTH=10 "
                     "TOLERANCE=4096bytes\n";
        return EXIT_SUCCESS;
    }

    std::cout.sync_with_stdio(false);

    try {
        const char* const binary_name = argv[1];
        const pointer_t address = [&]() {
            std::stringstream ss{argv[2]};
            void* ptr;
            ss >> std::hex >> ptr;
            return static_cast<pointer_t>(ptr);
        }();
        const int depth = (argc >= 4 ? std::stoi(argv[3]) : 10);
        const int tolerance = (argc >= 5 ? std::stoi(argv[4]) : 4096);

        do_pointer_scanner(binary_name, address, depth, tolerance);
    } catch (const std::exception& e) {
        std::cerr << "unhandled exception!\n    what(): " << e.what() << '\n';
    } catch (...) {
        std::cerr << "unhandled unknown exception!\n";
    }

    return EXIT_SUCCESS;
}