#pragma once

#include <iostream>
#include <string_view>
#include <source_location>
#include <cxxabi.h>
// execinfo.h not available in Alpine Linux - using alternative approach
#include <dlfcn.h>
#include <cstring>
#include <memory>
#include <sstream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <unistd.h>
#include <vector>

// ANSI color codes
namespace colors {
    constexpr const char* reset = "\033[0m";
    constexpr const char* red = "\033[31m";
    constexpr const char* green = "\033[32m";
    constexpr const char* yellow = "\033[33m";
    constexpr const char* blue = "\033[34m";
    constexpr const char* magenta = "\033[35m";
    constexpr const char* cyan = "\033[36m";
}

// Simple RAII wrapper for file descriptor
class FileDescriptor {
    int fd_ = -1;
public:
    explicit FileDescriptor(int fd) : fd_(fd) {}
    ~FileDescriptor() { if (fd_ != -1) close(fd_); }
    operator int() const { return fd_; }
    FileDescriptor(const FileDescriptor&) = delete;
    FileDescriptor& operator=(const FileDescriptor&) = delete;
};

struct SourceInfo {
    std::string function;
    std::string file;
    int line = 0;
};

std::vector<SourceInfo> get_source_info_batch(const char* executable, const std::vector<void*>& addresses) {
    if (addresses.empty()) return {};

    // Build command with all addresses
    std::string cmd = std::string("addr2line -f -e ") + executable;
    for (void* addr : addresses) {
        char addr_buf[32];
        snprintf(addr_buf, sizeof(addr_buf), " %p", addr);
        cmd += addr_buf;
    }
    cmd += " 2>/dev/null";

    FILE* pipe = popen(cmd.c_str(), "r");
    if (!pipe) return std::vector<SourceInfo>(addresses.size());

    std::vector<SourceInfo> results(addresses.size());
    char buffer[1024] = {0};
    size_t current_frame = 0;

    // Read function names and file/line info for each address
    while (current_frame < addresses.size() && fgets(buffer, sizeof(buffer), pipe)) {
        // Remove trailing newline
        size_t len = strlen(buffer);
        if (len > 0 && buffer[len-1] == '\n') {
            buffer[len-1] = '\0';
        }
        
        // Demangle function name
        int status = 0;
        char* demangled = abi::__cxa_demangle(buffer, nullptr, nullptr, &status);
        std::string func_name;
        if (demangled) {
            func_name = demangled;
            free(demangled);
        } else {
            func_name = buffer;
        }
        
        // Simplify function signature: replace argument list with (...)
        size_t paren_pos = func_name.find('(');
        if (paren_pos != std::string::npos) {
            func_name = func_name.substr(0, paren_pos) + "(...)";
        }
        
        results[current_frame].function = func_name;

        // Read file and line number
        if (fgets(buffer, sizeof(buffer), pipe)) {
            char* colon = strchr(buffer, ':');
            if (colon) {
                *colon = '\0';
                results[current_frame].line = atoi(colon + 1);
                
                // Extract just the filename
                const char* slash = strrchr(buffer, '/');
                results[current_frame].file = (slash ? slash + 1 : buffer);
            }
        }
        
        current_frame++;
    }

    pclose(pipe);
    return results;
}

// Helper function to get the program counter (return address) in a portable way
inline void* get_pc() {
    void* pc;
    #if defined(__aarch64__) || defined(__arm__)
        // For ARM/AArch64, we can use the link register (x30 on AArch64, r14 on ARM)
        #if defined(__aarch64__)
            asm volatile ("mov %0, x30" : "=r" (pc));
        #else
            asm volatile ("mov %0, lr" : "=r" (pc));
        #endif
    #else
        // For x86/x86_64, use __builtin_return_address
        pc = __builtin_return_address(0);
    #endif
    return pc;
}

void print_stacktrace() {
    std::cerr << colors::yellow << "Stack trace:" << colors::reset << "\n";
    
    char exe_path[1024] = {0};
    ssize_t count = readlink("/proc/self/exe", exe_path, sizeof(exe_path) - 1);
    if (count == -1) {
        strcpy(exe_path, "[unknown]");
    } else {
        exe_path[count] = '\0';
    }

    std::vector<void*> addresses;
    
    // Get current frame pointer and return address
    void* frame_ptr = __builtin_frame_address(0);
    void* return_addr = get_pc();
    
    // Add current return address
    if (return_addr) {
        addresses.push_back(return_addr);
    }
    
    // Try to walk up the stack if we have a valid frame pointer
    if (frame_ptr) {
        void** frame = static_cast<void**>(frame_ptr);
        // Limit the number of frames to prevent potential issues
        for (int i = 0; i < 5; ++i) {
            // The return address is typically at frame[1]
            if (frame[1] == nullptr) break;
            addresses.push_back(frame[1]);
            // Move to the next frame
            if (frame[0] == nullptr) break;
            frame = static_cast<void**>(frame[0]);
        }
    }
    
    if (addresses.empty()) {
        std::cerr << "  " << colors::red << "[no frames available]" << colors::reset << "\n";
        return;
    }

    // Get source info
    std::vector<SourceInfo> results = get_source_info_batch(exe_path, addresses);
    
    // Filter out frames from assert.hpp
    results.erase(std::remove_if(results.begin(), results.end(),
        [](const SourceInfo& frame) {
            if (frame.file.empty()) return false;
            std::string_view file(frame.file);
            return file.find("assert.hpp") != std::string_view::npos;
        }),
        results.end()
    );
    
    if (results.empty()) {
        std::cerr << "  " << colors::red << "[no user frames available]" << colors::reset << "\n";
        return;
    }
    
    // First pass: find the maximum function name length for alignment
    size_t max_func_length = 0;
    for (const auto& frame : results) {
        size_t length = frame.function.empty() 
            ? strlen("[unknown function]") 
            : frame.function.length();
        if (length > max_func_length) {
            max_func_length = length;
        }
    }
    
    // Second pass: print with aligned output
    for (size_t i = 0; i < results.size(); ++i) {
        const auto& frame = results[i];
        std::cerr << "  ";
        
        // Print function name with alignment
        size_t current_length;
        if (!frame.function.empty()) {
            std::cerr << colors::green << frame.function << colors::reset;
            current_length = frame.function.length();
        } else {
            std::cerr << colors::red << "[unknown function]" << colors::reset;
            current_length = strlen("[unknown function]");
        }
        
        // Add padding to align the file/line text
        if (!frame.file.empty() && frame.line > 0) {
            // Calculate padding (minimum 1 space)
            size_t padding = (max_func_length > current_length) 
                ? (max_func_length - current_length + 1) 
                : 1;
            std::cerr << std::string(padding, ' ')
                     << " " << colors::blue << frame.file << colors::reset 
                     << ":" << colors::magenta << frame.line << colors::reset;
        }
        
        std::cerr << "\n";
    }
}

void assert_failed(
    bool condition,
    std::string_view message,
    std::source_location location
) {
    if (!condition) {
        std::cout << std::flush;
        std::cerr << std::flush;

        std::cerr << colors::red << "Assertion failed at " << location.file_name() << ":" << location.line() << ": " 
                  << location.function_name() << ": " << message << colors::reset << "\n";
        print_stacktrace();
        std::abort();
    }
}

#define ASSERT(condition, message) assert_failed(condition, message, std::source_location::current())
