#ifndef DROPSHELL_ASSERT_HPP
#define DROPSHELL_ASSERT_HPP

#include <iostream>
#include <string_view>
#include <cstdlib> // For std::exit and EXIT_FAILURE
#include <execinfo.h>
#include <cxxabi.h>
#include <unistd.h>
#include <cstdio>
#include <memory>
#include <sstream>
#include <cstdint>
#include <cstring>

namespace ds {

struct SourceLocation {
    const char* file_name;
    int line;
    const char* function_name;
};

// Helper macro to create a SourceLocation with current context
#define DS_CURRENT_LOCATION ds::SourceLocation{__FILE__, __LINE__, __func__}

inline uintptr_t get_base_address(const char* exe_path) {
    // Open /proc/self/maps
    FILE* maps = fopen("/proc/self/maps", "r");
    if (!maps) return 0;
    char line[512];
    uintptr_t base_addr = 0;
    while (fgets(line, sizeof(line), maps)) {
        // Look for a line containing the exe path and 'r-xp'
        if (strstr(line, exe_path) && strstr(line, "r-xp")) {
            // Parse the start address
            char* endptr = nullptr;
            base_addr = strtoull(line, &endptr, 16);
            break;
        }
    }
    fclose(maps);
    return base_addr;
}

inline void print_stacktrace() {
    constexpr int max_frames = 64;
    void* addrlist[max_frames];
    int addrlen = backtrace(addrlist, max_frames);

    if (addrlen == 0) {
        std::cerr << "  <empty, possibly corrupt>\n";
        return;
    }

    // Get the program name
    char exe_path[1024];
    ssize_t len = readlink("/proc/self/exe", exe_path, sizeof(exe_path) - 1);
    if (len == -1) {
        std::cerr << "  <could not read exe path>\n";
        return;
    }
    exe_path[len] = '\0';

    uintptr_t base_addr = get_base_address(exe_path);

    int frame_num = 1;
    for (int i = 0; i < addrlen; i++) {
        uintptr_t addr = (uintptr_t)addrlist[i];
        uintptr_t offset = base_addr ? (addr - base_addr) : addr;
        std::ostringstream cmd;
        cmd << "addr2line -e " << exe_path << " -f -C 0x" << std::hex << offset;
        std::unique_ptr<FILE, int(*)(FILE*)> pipe(popen(cmd.str().c_str(), "r"), pclose);
        if (!pipe) {
            std::cerr << "  <addr2line failed>\n\n";
            frame_num++;
            continue;
        }
        // addr2line -f prints two lines per address: function name, then file:line
        char func[256] = {0};
        char fileline[256] = {0};
        if (fgets(func, sizeof(func), pipe.get()) && fgets(fileline, sizeof(fileline), pipe.get())) {
            // Remove trailing newlines
            size_t len1 = strlen(func);
            if (len1 > 0 && func[len1-1] == '\n') func[len1-1] = 0;
            size_t len2 = strlen(fileline);
            if (len2 > 0 && fileline[len2-1] == '\n') fileline[len2-1] = 0;
            // Remove argument list from function name (truncate at first '(')
            char* paren = strchr(func, '(');
            if (paren) *paren = 0;
            std::cerr << "  [" << frame_num << "] " << func << " at " << fileline << "\n\n";
        }
        frame_num++;
    }
}

[[noreturn]] inline void assert_fail(
    const char* expression,
    const SourceLocation& location,
    const char* message = nullptr) {
    
    std::cerr << "\033[1;31mAssertion failed!\033[0m\n"
              << "Expression: \033[1;33m" << expression << "\033[0m\n"
              << "Location:   \033[1;36m" << location.file_name << ":" 
                                           << location.line << "\033[0m\n"
              << "Function:   \033[1;36m" << location.function_name << "\033[0m\n";
              
    if (message) {
        std::cerr << "Message:    \033[1;37m" << message << "\033[0m\n";
    }
    
    std::cerr << std::endl;
    
    std::cerr << "Stack trace:" << std::endl;
    print_stacktrace();
    std::cerr << "----------------------------------------" << std::endl;
    
    // Exit the program without creating a core dump
    std::exit(EXIT_FAILURE);
}

} // namespace ds

// Standard assertion
#define ASSERT(condition) \
    do { \
        if (!(condition)) { \
            ds::assert_fail(#condition, DS_CURRENT_LOCATION); \
        } \
    } while (false)
    
// Assertion with custom message
#define ASSERT_MSG(condition, message) \
    do { \
        if (!(condition)) { \
            ds::assert_fail(#condition, DS_CURRENT_LOCATION, message); \
        } \
    } while (false)

#endif // DROPSHELL_ASSERT_HPP