---

// Build: g++ -std=c++17 -O2 -Wall -Wextra callgraph.cpp -o callgraph
// Run:   sudo ./proc_callgraph <pid> [samples=30] [delay_ms=100] [out=callgraph.dot]
// Render: dot -Tpng callgraph.dot -o callgraph.png
//
// Target program should preferably be built with:
//   -O0 -g -fno-omit-frame-pointer -rdynamic
// Otherwise frame-pointer unwinding and symbol names may be incomplete.

#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <unistd.h>

#include <cerrno>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>

#include <algorithm>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#ifndef __linux__
#error This example is intended for Linux.
#endif

#ifndef __x86_64__
#error This example supports x86_64 Linux only.
#endif

struct Frame {
    std::uint64_t previousRbp;
    std::uint64_t returnAddress;
};

struct MapEntry {
    std::uint64_t start = 0;
    std::uint64_t end = 0;
    std::uint64_t offset = 0;
    std::string perms;
    std::string path;
};

static std::string hex64(std::uint64_t value) {
    std::ostringstream out;
    out << "0x" << std::hex << value;
    return out.str();
}

static std::string trim(std::string s) {
    while (!s.empty() && (s.front() == ' ' || s.front() == '\t')) s.erase(s.begin());
    while (!s.empty() && (s.back() == ' ' || s.back() == '\t' || s.back() == '\n' || s.back() == '\r')) s.pop_back();
    return s;
}

static std::string baseName(const std::string& path) {
    std::size_t pos = path.find_last_of('/');
    if (pos == std::string::npos) return path;
    return path.substr(pos + 1);
}

static std::string shellQuote(const std::string& s) {
    std::string r = "'";
    for (char c : s) {
        if (c == '\'') r += "'\\''";
        else r += c;
    }
    r += "'";
    return r;
}

static std::string dotEscape(const std::string& s) {
    std::string r;
    for (char c : s) {
        if (c == '"') r += "\\\"";
        else if (c == '\\') r += "\\\\";
        else if (c == '\n') r += "\\n";
        else r += c;
    }
    return r;
}

static bool isRealFilePath(const std::string& path) {
    return !path.empty() && path[0] == '/';
}

class ProcessMaps {
public:
    explicit ProcessMaps(pid_t pid) {
        std::ifstream in("/proc/" + std::to_string(pid) + "/maps");
        if (!in) {
            throw std::runtime_error("cannot open /proc/" + std::to_string(pid) + "/maps");
        }

        std::string line;
        while (std::getline(in, line)) {
            MapEntry e;
            std::string range, dev, inode;
            std::istringstream ss(line);
            ss >> range >> e.perms >> std::hex >> e.offset >> dev >> inode;
            std::string path;
            std::getline(ss, path);
            e.path = trim(path);

            std::size_t dash = range.find('-');
            if (dash == std::string::npos) continue;
            e.start = std::stoull(range.substr(0, dash), nullptr, 16);
            e.end = std::stoull(range.substr(dash + 1), nullptr, 16);
            entries_.push_back(e);
        }
    }

    const MapEntry* find(std::uint64_t address) const {
        for (const auto& e : entries_) {
            if (address >= e.start && address < e.end) return &e;
        }
        return nullptr;
    }

private:
    std::vector<MapEntry> entries_;
};

class Symbolizer {
public:
    explicit Symbolizer(pid_t pid) : maps_(pid) {}

    std::string name(std::uint64_t address) {
        const MapEntry* m = maps_.find(address);
        if (!m) return hex64(address);

        std::string module = m->path.empty() ? "unknown" : baseName(m->path);
        std::uint64_t adjusted = (address - m->start) + m->offset;

        std::ostringstream key;
        key << m->path << ":" << std::hex << address << ":" << adjusted;
        auto it = cache_.find(key.str());
        if (it != cache_.end()) return it->second;

        std::string result;
        if (isRealFilePath(m->path)) {
            // For PIE/shared objects adjusted address usually works.
            // For non-PIE executables absolute runtime address may work better.
            result = addr2line(m->path, adjusted);
            if (result.empty()) result = addr2line(m->path, address);
        }

        if (result.empty()) {
            result = module + "+" + hex64(adjusted);
        } else {
            result += "\n" + module + "+" + hex64(adjusted);
        }

        cache_[key.str()] = result;
        return result;
    }

private:
    ProcessMaps maps_;
    std::unordered_map<std::string, std::string> cache_;

    static std::string addr2line(const std::string& file, std::uint64_t address) {
        std::ostringstream cmd;
        cmd << "addr2line -f -C -e " << shellQuote(file) << " " << hex64(address) << " 2>/dev/null";

        FILE* pipe = popen(cmd.str().c_str(), "r");
        if (!pipe) return "";

        char buffer[4096];
        std::string functionName;
        std::string location;

        if (fgets(buffer, sizeof(buffer), pipe)) functionName = trim(buffer);
        if (fgets(buffer, sizeof(buffer), pipe)) location = trim(buffer);

        int rc = pclose(pipe);
        (void)rc;

        if (functionName.empty() || functionName == "??") return "";
        if (!location.empty() && location != "??:0") {
            return functionName + "\n" + location;
        }
        return functionName;
    }
};

class CallGraph {
public:
    void addStack(const std::vector<std::string>& rootToLeaf) {
        if (rootToLeaf.empty()) return;

        for (const auto& n : rootToLeaf) nodeHits_[n]++;
        for (std::size_t i = 1; i < rootToLeaf.size(); ++i) {
            edges_[{rootToLeaf[i - 1], rootToLeaf[i]}]++;
        }
    }

    void writeDot(const std::string& fileName) const {
        std::ofstream out(fileName);
        if (!out) throw std::runtime_error("cannot write " + fileName);

        out << "digraph ProcessCallGraph {\n";
        out << "  rankdir=TB;\n";
        out << "  graph [fontname=\"Arial\"];\n";
        out << "  node [shape=box, style=rounded, fontname=\"Arial\"];\n";
        out << "  edge [fontname=\"Arial\"];\n\n";

        for (const auto& [name, hits] : nodeHits_) {
            out << "  \"" << dotEscape(name) << "\" [label=\"" << dotEscape(name)
                << "\\nhits: " << hits << "\"];\n";
        }

        out << "\n";
        for (const auto& [edge, count] : edges_) {
            out << "  \"" << dotEscape(edge.first) << "\" -> \""
                << dotEscape(edge.second) << "\" [label=\"" << count << "\"];\n";
        }

        out << "}\n";
    }

    bool empty() const {
        return nodeHits_.empty();
    }

private:
    std::map<std::string, int> nodeHits_;
    std::map<std::pair<std::string, std::string>, int> edges_;
};

static bool readProcessMemory(pid_t pid, std::uint64_t remoteAddress, void* localBuffer, std::size_t size) {
    iovec local{};
    local.iov_base = localBuffer;
    local.iov_len = size;

    iovec remote{};
    remote.iov_base = reinterpret_cast<void*>(remoteAddress);
    remote.iov_len = size;

    ssize_t n = process_vm_readv(pid, &local, 1, &remote, 1, 0);
    if (n == static_cast<ssize_t>(size)) return true;

    // Fallback: read word-by-word via ptrace.
    std::uint8_t* dst = static_cast<std::uint8_t*>(localBuffer);
    std::size_t copied = 0;

    while (copied < size) {
        errno = 0;
        long word = ptrace(PTRACE_PEEKDATA, pid, reinterpret_cast<void*>(remoteAddress + copied), nullptr);
        if (errno != 0) return false;

        std::size_t chunk = std::min(sizeof(word), size - copied);
        std::memcpy(dst + copied, &word, chunk);
        copied += chunk;
    }

    return true;
}

class TraceAttach {
public:
    explicit TraceAttach(pid_t pid) : pid_(pid) {
        if (ptrace(PTRACE_ATTACH, pid_, nullptr, nullptr) == -1) {
            throw std::runtime_error("PTRACE_ATTACH failed: " + std::string(std::strerror(errno)));
        }

        int status = 0;
        if (waitpid(pid_, &status, 0) == -1) {
            throw std::runtime_error("waitpid failed after attach: " + std::string(std::strerror(errno)));
        }

        attached_ = true;
    }

    ~TraceAttach() {
        if (attached_) {
            ptrace(PTRACE_DETACH, pid_, nullptr, nullptr);
        }
    }

    TraceAttach(const TraceAttach&) = delete;
    TraceAttach& operator=(const TraceAttach&) = delete;

private:
    pid_t pid_;
    bool attached_ = false;
};

static std::vector<std::uint64_t> captureAddresses(pid_t pid) {
    user_regs_struct regs{};
    if (ptrace(PTRACE_GETREGS, pid, nullptr, &regs) == -1) {
        throw std::runtime_error("PTRACE_GETREGS failed: " + std::string(std::strerror(errno)));
    }

    std::vector<std::uint64_t> addresses;
    addresses.push_back(static_cast<std::uint64_t>(regs.rip));

    std::uint64_t rbp = static_cast<std::uint64_t>(regs.rbp);
    constexpr int MAX_DEPTH = 64;
    constexpr std::uint64_t MAX_FRAME_DISTANCE = 1024 * 1024;

    for (int depth = 0; depth < MAX_DEPTH && rbp != 0; ++depth) {
        Frame frame{};
        if (!readProcessMemory(pid, rbp, &frame, sizeof(frame))) break;
        if (frame.returnAddress == 0) break;

        addresses.push_back(frame.returnAddress);

        if (frame.previousRbp <= rbp) break;
        if (frame.previousRbp - rbp > MAX_FRAME_DISTANCE) break;

        rbp = frame.previousRbp;
    }

    return addresses;
}

static std::vector<std::string> sampleStack(pid_t pid) {
    TraceAttach attach(pid);
    Symbolizer symbolizer(pid);

    std::vector<std::uint64_t> addresses = captureAddresses(pid);
    std::vector<std::string> names;
    names.reserve(addresses.size());

    for (std::uint64_t address : addresses) {
        names.push_back(symbolizer.name(address));
    }

    // We collected leaf -> root. DOT graph should be root -> leaf.
    std::reverse(names.begin(), names.end());
    return names;
}

static int toInt(const char* s, int fallback) {
    if (!s) return fallback;
    char* end = nullptr;
    long v = std::strtol(s, &end, 10);
    if (end == s || *end != '\0' || v <= 0) return fallback;
    return static_cast<int>(v);
}

static void printUsage(const char* argv0) {
    std::cerr
        << "Usage:\n"
        << "  " << argv0 << " <pid> [samples=30] [delay_ms=100] [out=callgraph.dot]\n\n"
        << "Example:\n"
        << "  sudo " << argv0 << " 12345 40 100 callgraph.dot\n"
        << "  dot -Tpng callgraph.dot -o callgraph.png\n\n"
        << "Target program should be compiled with frame pointers, for example:\n"
        << "  g++ -O0 -g -fno-omit-frame-pointer -rdynamic target.cpp -o target\n";
}

int main(int argc, char** argv) {
    if (argc < 2 || std::string(argv[1]) == "-h" || std::string(argv[1]) == "--help") {
        printUsage(argv[0]);
        return argc < 2 ? 1 : 0;
    }

    pid_t pid = static_cast<pid_t>(toInt(argv[1], -1));
    int samples = argc > 2 ? toInt(argv[2], 30) : 30;
    int delayMs = argc > 3 ? toInt(argv[3], 100) : 100;
    std::string outFile = argc > 4 ? argv[4] : "callgraph.dot";

    if (pid <= 0) {
        std::cerr << "Invalid PID\n";
        return 1;
    }

    CallGraph graph;

    try {
        for (int i = 0; i < samples; ++i) {
            std::vector<std::string> stack = sampleStack(pid);
            graph.addStack(stack);

            std::cerr << "sample " << (i + 1) << "/" << samples
                      << ": depth=" << stack.size() << "\n";

            if (i + 1 < samples) usleep(static_cast<useconds_t>(delayMs) * 1000);
        }

        if (graph.empty()) {
            std::cerr << "No stack frames collected. Check permissions and frame pointers.\n";
            return 1;
        }

        graph.writeDot(outFile);
        std::cout << "Call graph saved to " << outFile << "\n";
        std::cout << "Render example: dot -Tpng " << outFile << " -o callgraph.png\n";
        return 0;
    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << "\n";
        std::cerr << "Try running with sudo, and make sure the target is built with -fno-omit-frame-pointer.\n";
        return 1;
    }
}