/******************************************************************************
 * Remote Debugging Module - Asyncio Functions
 *
 * This file contains functions for parsing asyncio tasks, coroutines,
 * and awaited_by relationships from remote process memory.
 ******************************************************************************/

#include "_remote_debugging.h"
#include "debug_offsets_validation.h"

/* ============================================================================
 * ASYNCIO DEBUG ADDRESS FUNCTIONS
 * ============================================================================ */

uintptr_t
_Py_RemoteDebug_GetAsyncioDebugAddress(proc_handle_t* handle)
{
    uintptr_t address;

#ifdef MS_WINDOWS
    // On Windows, search for asyncio debug in executable or DLL
    address = search_windows_map_for_section(handle, "AsyncioD", L"_asyncio",
                                             NULL);
    if (address == 0) {
        if (!_Py_RemoteDebug_HasPermissionError()) {
            PyObject *exc = PyErr_GetRaisedException();
            PyErr_SetString(PyExc_RuntimeError, "Failed to find the AsyncioDebug section in the process.");
            _PyErr_ChainExceptions1(exc);
        }
    }
#elif defined(__linux__) && HAVE_PROCESS_VM_READV
    // On Linux, search for asyncio debug in executable or DLL
    address = search_linux_map_for_section(handle, "AsyncioDebug", "python",
                                           NULL);
    if (address == 0) {
        if (!_Py_RemoteDebug_HasPermissionError()) {
            PyObject *exc = PyErr_GetRaisedException();
            PyErr_SetString(PyExc_RuntimeError, "Failed to find the AsyncioDebug section in the process.");
            _PyErr_ChainExceptions1(exc);
        }
    }
#elif defined(__APPLE__) && TARGET_OS_OSX
    // On macOS, try libpython first, then fall back to python
    address = search_map_for_section(handle, "AsyncioDebug", "libpython",
                                     NULL);
    if (address == 0 && !_Py_RemoteDebug_HasPermissionError()) {
        PyErr_Clear();
        address = search_map_for_section(handle, "AsyncioDebug", "python",
                                         NULL);
    }
    if (address == 0) {
        if (!_Py_RemoteDebug_HasPermissionError()) {
            PyObject *exc = PyErr_GetRaisedException();
            PyErr_SetString(PyExc_RuntimeError, "Failed to find the AsyncioDebug section in the process.");
            _PyErr_ChainExceptions1(exc);
        }
    }
#else
    Py_UNREACHABLE();
#endif

    return address;
}

int
read_async_debug(RemoteUnwinderObject *unwinder)
{
    uintptr_t async_debug_addr = _Py_RemoteDebug_GetAsyncioDebugAddress(&unwinder->handle);
    if (!async_debug_addr) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to get AsyncioDebug address");
        return -1;
    }

    size_t size = sizeof(struct _Py_AsyncioModuleDebugOffsets);
    int result = _Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle, async_debug_addr, size, &unwinder->async_debug_offsets);
    if (result < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read AsyncioDebug offsets");
        return result;
    }
    if (_PyRemoteDebug_ValidateAsyncDebugOffsets(&unwinder->async_debug_offsets) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Invalid AsyncioDebug offsets");
        return PY_REMOTE_DEBUG_INVALID_ASYNC_DEBUG_OFFSETS;
    }
    return 0;
}

int
ensure_async_debug_offsets(RemoteUnwinderObject *unwinder)
{
    // If already available, nothing to do
    if (unwinder->async_debug_offsets_available) {
        return 0;
    }

    // Try to load async debug offsets (the target process may have
    // loaded asyncio since we last checked)
    int result = read_async_debug(unwinder);
    if (result == PY_REMOTE_DEBUG_INVALID_ASYNC_DEBUG_OFFSETS) {
        return -1;
    }
    if (result < 0) {
        if (!_Py_RemoteDebug_HasPermissionError()) {
            PyErr_Clear();
            PyErr_SetString(PyExc_RuntimeError, "AsyncioDebug section not available");
            set_exception_cause(unwinder, PyExc_RuntimeError,
                "AsyncioDebug section unavailable - asyncio module may not be loaded in target process");
        }
        return -1;
    }

    unwinder->async_debug_offsets_available = 1;
    return 0;
}

/* ============================================================================
 * SET ITERATION FUNCTIONS
 * ============================================================================ */

static int
iterate_set_entries(
    RemoteUnwinderObject *unwinder,
    uintptr_t set_addr,
    PyObject *awaited_by
) {
    char set_object[SIZEOF_SET_OBJ];
    if (_Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle, set_addr,
                                              SIZEOF_SET_OBJ, set_object) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read set object");
        return -1;
    }

    Py_ssize_t num_els = GET_MEMBER(Py_ssize_t, set_object, unwinder->debug_offsets.set_object.used);
    Py_ssize_t mask = GET_MEMBER(Py_ssize_t, set_object, unwinder->debug_offsets.set_object.mask);
    uintptr_t table_ptr = GET_MEMBER(uintptr_t, set_object, unwinder->debug_offsets.set_object.table);

    // Validate mask and num_els to prevent huge loop iterations from garbage data
    if (mask < 0 || mask >= MAX_SET_TABLE_SIZE || num_els < 0 || num_els > mask + 1) {
        PyErr_SetString(PyExc_RuntimeError, "Invalid set object (corrupted remote memory)");
        set_exception_cause(unwinder, PyExc_RuntimeError,
            "Invalid set object (corrupted remote memory)");
        return -1;
    }
    Py_ssize_t set_len = mask + 1;

    Py_ssize_t i = 0;
    Py_ssize_t els = 0;
    while (i < set_len && els < num_els) {
        setentry entry;
        if (_Py_RemoteDebug_PagedReadRemoteMemory(
                &unwinder->handle, table_ptr, sizeof(entry), &entry) < 0)
        {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read set entry");
            return -1;
        }

        uintptr_t key_addr = (uintptr_t)entry.key;
        if (key_addr != 0 && entry.hash != -1) {
            if (parse_task(unwinder, key_addr, awaited_by) < 0) {
                return -1;
            }
            els++;
        }
        table_ptr += sizeof(entry);
        i++;
    }

    return 0;
}

/* ============================================================================
 * TASK NAME PARSING
 * ============================================================================ */

PyObject *
parse_task_name(
    RemoteUnwinderObject *unwinder,
    uintptr_t task_address
) {
    // Read the entire TaskObj at once
    char task_obj[SIZEOF_TASK_OBJ];
    int err = _Py_RemoteDebug_PagedReadRemoteMemory(
        &unwinder->handle,
        task_address,
        (size_t)unwinder->async_debug_offsets.asyncio_task_object.size,
        task_obj);
    if (err < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read task object");
        return NULL;
    }

    uintptr_t task_name_addr = GET_MEMBER_NO_TAG(uintptr_t, task_obj, unwinder->async_debug_offsets.asyncio_task_object.task_name);

    // The task name can be a long or a string so we need to check the type
    char task_name_obj[SIZEOF_PYOBJECT];
    err = _Py_RemoteDebug_PagedReadRemoteMemory(
        &unwinder->handle,
        task_name_addr,
        SIZEOF_PYOBJECT,
        task_name_obj);
    if (err < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read task name object");
        return NULL;
    }

    // Now read the type object to get the flags
    char type_obj[SIZEOF_TYPE_OBJ];
    err = _Py_RemoteDebug_PagedReadRemoteMemory(
        &unwinder->handle,
        GET_MEMBER(uintptr_t, task_name_obj, unwinder->debug_offsets.pyobject.ob_type),
        SIZEOF_TYPE_OBJ,
        type_obj);
    if (err < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read task name type object");
        return NULL;
    }

    if ((GET_MEMBER(unsigned long, type_obj, unwinder->debug_offsets.type_object.tp_flags) & Py_TPFLAGS_LONG_SUBCLASS)) {
        long res = read_py_long(unwinder, task_name_addr);
        if (res == -1 && PyErr_Occurred()) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Task name PyLong parsing failed");
            return NULL;
        }
        return PyUnicode_FromFormat("Task-%ld", res);
    }

    if(!(GET_MEMBER(unsigned long, type_obj, unwinder->debug_offsets.type_object.tp_flags) & Py_TPFLAGS_UNICODE_SUBCLASS)) {
        PyErr_SetString(PyExc_RuntimeError, "Invalid task name object");
        set_exception_cause(unwinder, PyExc_RuntimeError, "Task name object is neither long nor unicode");
        return NULL;
    }

    return read_py_str(
        unwinder,
        task_name_addr,
        255
    );
}

/* ============================================================================
 * COROUTINE CHAIN PARSING
 * ============================================================================ */

// gh-159620
/* Return in *coro_addr the coroutine or async generator to walk for the
   awaitable at addr, or 0 if it is neither (for example, a future). */
static int
get_coro_like_address(
    RemoteUnwinderObject *unwinder,
    uintptr_t addr,
    uintptr_t *coro_addr
) {
    *coro_addr = 0;

    // Skip user classes that share a name with the generator types below
    uintptr_t type_addr;
    unsigned long tp_flags;
    if (read_ptr(unwinder, addr + (uintptr_t)unwinder->debug_offsets.pyobject.ob_type,
                 &type_addr) < 0
        || _Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle,
                type_addr + (uintptr_t)unwinder->debug_offsets.type_object.tp_flags,
                sizeof(tp_flags), &tp_flags) < 0)
    {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read awaitable type");
        return -1;
    }

    if (tp_flags & Py_TPFLAGS_HEAPTYPE) {
        return 0;
    }

    uintptr_t name_addr;
    char name[32] = {0};
    if (read_ptr(unwinder, type_addr + (uintptr_t)unwinder->debug_offsets.type_object.tp_name,
                 &name_addr) < 0
        || _Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle, name_addr,
                                                 sizeof(name) - 1, name) < 0)
    {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read awaitable type name");
        return -1;
    }

    if (strcmp(name, "coroutine") == 0) {
        *coro_addr = addr;
    }
    else if (strcmp(name, "async_generator_asend") == 0
             || strcmp(name, "async_generator_athrow") == 0)
    {
        // asend and athrow are separate wrapper objects: read the pointer
        // to their parent async generator, stored right after PyObject_HEAD
        return read_ptr(unwinder, addr + (uintptr_t)unwinder->debug_offsets.pyobject.size,
                        coro_addr);
    }
    return 0;
}

static int
get_awaited_coro_address(
    RemoteUnwinderObject *unwinder,
    uintptr_t gi_iframe_addr,
    uintptr_t *next_coro
) {
    *next_coro = 0;

    // Read the entire interpreter frame at once
    char iframe[SIZEOF_INTERP_FRAME];
    int err = _Py_RemoteDebug_PagedReadRemoteMemory(
        &unwinder->handle,
        gi_iframe_addr,
        SIZEOF_INTERP_FRAME,
        iframe);
    if (err < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read interpreter frame in yield_from handler");
        return -1;
    }

    if (GET_MEMBER(char, iframe, unwinder->debug_offsets.interpreter_frame.owner) != FRAME_OWNED_BY_GENERATOR) {
        PyErr_SetString(
            PyExc_RuntimeError,
            "generator doesn't own its frame \\_o_/");
        set_exception_cause(unwinder, PyExc_RuntimeError, "Frame ownership mismatch in yield_from");
        return -1;
    }

    uintptr_t stackpointer_addr = GET_MEMBER_NO_TAG(uintptr_t, iframe, unwinder->debug_offsets.interpreter_frame.stackpointer);

    if ((void*)stackpointer_addr != NULL) {
        err = read_py_ptr(
            unwinder,
            stackpointer_addr - sizeof(void*) * 2,
            next_coro);
        if (err) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read gi_await address");
            return -1;
        }
    }

    return 0;
}

static int
parse_coro_chain(
    RemoteUnwinderObject *unwinder,
    uintptr_t coro_address,
    PyObject *render_to
) {
    assert((void*)coro_address != NULL);

    for (size_t depth = 0; (void*)coro_address != NULL; depth++) {
        if (depth >= MAX_FRAME_CHAIN_DEPTH) {
            PyErr_SetString(PyExc_RuntimeError,
                "Too many coroutine frames (possible infinite loop)");
            set_exception_cause(unwinder, PyExc_RuntimeError,
                "Coroutine chain depth limit exceeded");
            return -1;
        }

        if (get_coro_like_address(unwinder, coro_address, &coro_address) < 0) {
            return -1;
        }
        if ((void*)coro_address == NULL) {
            return 0;
        }

        // Read the entire generator object at once
        char gen_object[SIZEOF_GEN_OBJ];
        int err = _Py_RemoteDebug_PagedReadRemoteMemory(
            &unwinder->handle,
            coro_address,
            SIZEOF_GEN_OBJ,
            gen_object);
        if (err < 0) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read generator object in coro chain");
            return -1;
        }

        int8_t frame_state = GET_MEMBER(int8_t, gen_object, unwinder->debug_offsets.gen_object.gi_frame_state);
        if (frame_state == FRAME_CLEARED) {
            return 0;
        }

        PyObject* name = NULL;

        // Parse the previous frame using the gi_iframe from local copy
        uintptr_t prev_frame;
        uintptr_t gi_iframe_addr = coro_address + (uintptr_t)unwinder->debug_offsets.gen_object.gi_iframe;
        uintptr_t address_of_code_object = 0;
        if (parse_frame_object(unwinder, &name, gi_iframe_addr, &address_of_code_object, &prev_frame) < 0) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse frame object in coro chain");
            return -1;
        }

        if (!name) {
            return 0;
        }

        if (PyList_Append(render_to, name)) {
            Py_DECREF(name);
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to append frame to coro chain");
            return -1;
        }
        Py_DECREF(name);

        if (frame_state != FRAME_SUSPENDED_YIELD_FROM) {
            return 0;
        }

        if (get_awaited_coro_address(unwinder, gi_iframe_addr, &coro_address) < 0) {
            return -1;
        }
    }

    return 0;
}

/* ============================================================================
 * TASK PARSING FUNCTIONS
 * ============================================================================ */

static PyObject*
create_task_result(
    RemoteUnwinderObject *unwinder,
    uintptr_t task_address
) {
    PyObject* result = NULL;
    PyObject *call_stack = NULL;
    PyObject *tn = NULL;
    char task_obj[SIZEOF_TASK_OBJ];
    uintptr_t coro_addr;

    // Create call_stack first since it's the first tuple element
    call_stack = PyList_New(0);
    if (call_stack == NULL) {
        set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create call stack list");
        goto error;
    }

    // Create task name/address for second tuple element
    tn = PyLong_FromUnsignedLongLong(task_address);
    if (tn == NULL) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to create task name/address");
        goto error;
    }

    // Parse coroutine chain
    if (_Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle, task_address,
                                              (size_t)unwinder->async_debug_offsets.asyncio_task_object.size,
                                              task_obj) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read task object for coro chain");
        goto error;
    }

    coro_addr = GET_MEMBER_NO_TAG(uintptr_t, task_obj, unwinder->async_debug_offsets.asyncio_task_object.task_coro);

    if ((void*)coro_addr != NULL) {
        if (parse_coro_chain(unwinder, coro_addr, call_stack) < 0) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse coroutine chain");
            goto error;
        }

        if (PyList_Reverse(call_stack)) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to reverse call stack");
            goto error;
        }
    }

    // Create final CoroInfo result
    RemoteDebuggingState *state = RemoteDebugging_GetStateFromObject((PyObject*)unwinder);
    result = PyStructSequence_New(state->CoroInfo_Type);
    if (result == NULL) {
        set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create CoroInfo");
        goto error;
    }

    // PyStructSequence_SetItem steals references, so we don't need to DECREF on success
    PyStructSequence_SetItem(result, 0, call_stack);  // This steals the reference
    PyStructSequence_SetItem(result, 1, tn);  // This steals the reference

    return result;

error:
    Py_XDECREF(result);
    Py_XDECREF(call_stack);
    Py_XDECREF(tn);
    return NULL;
}

int
parse_task(
    RemoteUnwinderObject *unwinder,
    uintptr_t task_address,
    PyObject *render_to
) {
    char is_task;
    PyObject* result = NULL;
    int err;

    err = read_char(
        unwinder,
        task_address + (uintptr_t)unwinder->async_debug_offsets.asyncio_task_object.task_is_task,
        &is_task);
    if (err) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read is_task flag");
        goto error;
    }

    if (is_task) {
        result = create_task_result(unwinder, task_address);
        if (!result) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to create task result");
            goto error;
        }
    } else {
        // Create an empty CoroInfo for non-task objects
        RemoteDebuggingState *state = RemoteDebugging_GetStateFromObject((PyObject*)unwinder);
        result = PyStructSequence_New(state->CoroInfo_Type);
        if (result == NULL) {
            set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create empty CoroInfo");
            goto error;
        }
        PyObject *empty_list = PyList_New(0);
        if (empty_list == NULL) {
            set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create empty list");
            goto error;
        }
        PyObject *task_name = PyLong_FromUnsignedLongLong(task_address);
        if (task_name == NULL) {
            Py_DECREF(empty_list);
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to create task name");
            goto error;
        }
        PyStructSequence_SetItem(result, 0, empty_list);  // This steals the reference
        PyStructSequence_SetItem(result, 1, task_name);  // This steals the reference
    }
    if (PyList_Append(render_to, result)) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to append task result to render list");
        goto error;
    }

    Py_DECREF(result);
    return 0;

error:
    Py_XDECREF(result);
    return -1;
}

/* ============================================================================
 * TASK AWAITED_BY PROCESSING
 * ============================================================================ */

static int
parse_task_awaited_by(
    RemoteUnwinderObject *unwinder,
    uintptr_t task_address,
    PyObject *awaited_by
) {
    // Read the entire TaskObj at once
    char task_obj[SIZEOF_TASK_OBJ];
    if (_Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle, task_address,
                                              (size_t)unwinder->async_debug_offsets.asyncio_task_object.size,
                                              task_obj) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read task object");
        return -1;
    }

    uintptr_t task_ab_addr = GET_MEMBER_NO_TAG(uintptr_t, task_obj, unwinder->async_debug_offsets.asyncio_task_object.task_awaited_by);
    if ((void*)task_ab_addr == NULL) {
        return 0;  // No tasks waiting for this one
    }

    char awaited_by_is_a_set = GET_MEMBER(char, task_obj, unwinder->async_debug_offsets.asyncio_task_object.task_awaited_by_is_set);

    if (awaited_by_is_a_set) {
        return iterate_set_entries(unwinder, task_ab_addr, awaited_by);
    } else {
        // Single task waiting
        return parse_task(unwinder, task_ab_addr, awaited_by);
    }
}

int
process_single_task_node(
    RemoteUnwinderObject *unwinder,
    uintptr_t task_addr,
    PyObject **task_info,
    PyObject *result
) {
    PyObject *tn = NULL;
    PyObject *current_awaited_by = NULL;
    PyObject *task_id = NULL;
    PyObject *result_item = NULL;
    PyObject *coroutine_stack = NULL;

    tn = parse_task_name(unwinder, task_addr);
    if (tn == NULL) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse task name in single task node");
        goto error;
    }

    current_awaited_by = PyList_New(0);
    if (current_awaited_by == NULL) {
        set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create awaited_by list in single task node");
        goto error;
    }

    // Extract the coroutine stack for this task
    coroutine_stack = PyList_New(0);
    if (coroutine_stack == NULL) {
        set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create coroutine stack list in single task node");
        goto error;
    }

    if (parse_task(unwinder, task_addr, coroutine_stack) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse task coroutine stack in single task node");
        goto error;
    }

    task_id = PyLong_FromUnsignedLongLong(task_addr);
    if (task_id == NULL) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to create task ID in single task node");
        goto error;
    }

    RemoteDebuggingState *state = RemoteDebugging_GetStateFromObject((PyObject*)unwinder);
    result_item = PyStructSequence_New(state->TaskInfo_Type);
    if (result_item == NULL) {
        set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create TaskInfo in single task node");
        goto error;
    }

    PyStructSequence_SetItem(result_item, 0, task_id);  // steals ref
    PyStructSequence_SetItem(result_item, 1, tn);  // steals ref
    PyStructSequence_SetItem(result_item, 2, coroutine_stack);  // steals ref
    PyStructSequence_SetItem(result_item, 3, current_awaited_by);  // steals ref

    // References transferred to tuple
    task_id = NULL;
    tn = NULL;
    coroutine_stack = NULL;
    current_awaited_by = NULL;

    if (PyList_Append(result, result_item)) {
        Py_DECREF(result_item);
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to append result item in single task node");
        return -1;
    }
    if (task_info != NULL) {
        *task_info = result_item;
    }
    Py_DECREF(result_item);

    // Get back current_awaited_by reference for parse_task_awaited_by
    current_awaited_by = PyStructSequence_GetItem(result_item, 3);
    if (parse_task_awaited_by(unwinder, task_addr, current_awaited_by) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to parse awaited_by in single task node");
        // No cleanup needed here since all references were transferred to result_item
        // and result_item was already added to result list and decreffed
        return -1;
    }

    return 0;

error:
    Py_XDECREF(tn);
    Py_XDECREF(current_awaited_by);
    Py_XDECREF(task_id);
    Py_XDECREF(result_item);
    Py_XDECREF(coroutine_stack);
    return -1;
}

static int
process_task_waiters(
    RemoteUnwinderObject *unwinder,
    PyObject *result
) {
    for (Py_ssize_t i = 0; i < PyList_GET_SIZE(result); i++) {
        PyObject *task_info = PyList_GET_ITEM(result, i);
        PyObject *waiters = PyStructSequence_GET_ITEM(task_info, 3);
        for (Py_ssize_t j = 0; j < PyList_GET_SIZE(waiters); j++) {
            if (PyList_GET_SIZE(result) >= MAX_TASK_WAITER_WALK_TASKS) {
                PyErr_SetString(PyExc_RuntimeError,
                    "Too many task waiters (possible infinite loop)");
                set_exception_cause(unwinder, PyExc_RuntimeError,
                    "Task waiter walk size limit exceeded");
                return -1;
            }
            PyObject *waiter = PyList_GET_ITEM(waiters, j);
            // CoroInfo item 1 holds the waiter task address stored by parse_task().
            PyObject *task_id = PyStructSequence_GET_ITEM(waiter, 1);
            void *task_ptr = PyLong_AsVoidPtr(task_id);
            if (task_ptr == NULL && PyErr_Occurred()) {
                set_exception_cause(unwinder, PyExc_RuntimeError,
                                    "Failed to parse waiter task ID");
                return -1;
            }
            if (process_single_task_node(
                    unwinder, (uintptr_t)task_ptr, NULL, result) < 0)
            {
                return -1;
            }
        }
    }

    return 0;
}

/* ============================================================================
 * RUNNING TASK FUNCTIONS
 * ============================================================================ */

int
find_running_task_in_thread(
    RemoteUnwinderObject *unwinder,
    uintptr_t thread_state_addr,
    uintptr_t *running_task_addr
) {
    *running_task_addr = (uintptr_t)NULL;

    uintptr_t address_of_running_loop;
    int bytes_read = read_py_ptr(
        unwinder,
        thread_state_addr + (uintptr_t)unwinder->async_debug_offsets.asyncio_thread_state.asyncio_running_loop,
        &address_of_running_loop);
    if (bytes_read == -1) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read running loop address");
        return -1;
    }

    // no asyncio loop is now running
    if ((void*)address_of_running_loop == NULL) {
        return 0;
    }

    int err = read_ptr(
        unwinder,
        thread_state_addr + (uintptr_t)unwinder->async_debug_offsets.asyncio_thread_state.asyncio_running_task,
        running_task_addr);
    if (err) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read running task address");
        return -1;
    }

    return 0;
}

int
get_task_frame(RemoteUnwinderObject *unwinder, uintptr_t task_addr, uintptr_t *frame_addr) {
    uintptr_t running_coro_addr = 0;

    if(read_py_ptr(
        unwinder,
        task_addr + (uintptr_t)unwinder->async_debug_offsets.asyncio_task_object.task_coro,
        &running_coro_addr) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Running task coro read failed");
        return -1;
    }

    if (running_coro_addr == 0) {
        PyErr_SetString(PyExc_RuntimeError, "Running task coro is NULL");
        set_exception_cause(unwinder, PyExc_RuntimeError, "Running task coro address is NULL");
        return -1;
    }

    if (get_coro_like_address(unwinder, running_coro_addr, &running_coro_addr) < 0) {
        return -1;
    }

    *frame_addr = running_coro_addr + (uintptr_t)unwinder->debug_offsets.gen_object.gi_iframe;
    return 0;
}

/* ============================================================================
 * ASYNC FRAME CHAIN PARSING
 * ============================================================================ */

int
parse_async_frame_chain(
    RemoteUnwinderObject *unwinder,
    PyObject *calls,
    uintptr_t address_of_thread,
    uintptr_t running_task_frame
) {
    uintptr_t address_of_current_frame;
    if (find_running_frame(unwinder, address_of_thread, &address_of_current_frame) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Running frame search failed in async chain");
        return -1;
    }

    size_t frame_count = 0;
    while ((void*)address_of_current_frame != NULL) {
        if (++frame_count > MAX_FRAME_CHAIN_DEPTH) {
            PyErr_SetString(PyExc_RuntimeError, "Too many async stack frames (possible infinite loop)");
            set_exception_cause(unwinder, PyExc_RuntimeError, "Async frame chain iteration limit exceeded");
            return -1;
        }
        PyObject* frame_info = NULL;
        uintptr_t address_of_code_object;

        uintptr_t this_frame = address_of_current_frame;
        int res = parse_frame_object(
            unwinder,
            &frame_info,
            address_of_current_frame,
            &address_of_code_object,
            &address_of_current_frame
        );

        if (res < 0) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Async frame object parsing failed in chain");
            return -1;
        }

        if (!frame_info) {
            continue;
        }

        if (PyList_Append(calls, frame_info) == -1) {
            Py_DECREF(frame_info);
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to append frame info to async chain");
            return -1;
        }

        Py_DECREF(frame_info);

        // Stop at the task's own frame. Code objects are shared by
        // recursive calls, so they cannot identify it.
        if (this_frame == running_task_frame) {
            break;
        }
    }

    return 0;
}

/* ============================================================================
 * AWAITED BY PARSING FUNCTIONS
 * ============================================================================ */

static int
append_awaited_by_for_thread(
    RemoteUnwinderObject *unwinder,
    uintptr_t head_addr,
    PyObject *result
) {
    char task_node[SIZEOF_LLIST_NODE];

    if (_Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle, head_addr,
                                              sizeof(task_node), task_node) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read task node head");
        return -1;
    }

    size_t iteration_count = 0;
    const size_t MAX_ITERATIONS = 2 << 15;  // A reasonable upper bound

    while (GET_MEMBER(uintptr_t, task_node, unwinder->debug_offsets.llist_node.next) != head_addr) {
        if (++iteration_count > MAX_ITERATIONS) {
            PyErr_SetString(PyExc_RuntimeError, "Task list appears corrupted");
            set_exception_cause(unwinder, PyExc_RuntimeError, "Task list iteration limit exceeded");
            return -1;
        }

        uintptr_t next_node = GET_MEMBER(uintptr_t, task_node, unwinder->debug_offsets.llist_node.next);
        if (next_node == 0) {
            PyErr_SetString(PyExc_RuntimeError,
                           "Invalid linked list structure reading remote memory");
            set_exception_cause(unwinder, PyExc_RuntimeError, "NULL pointer in task linked list");
            return -1;
        }

        uintptr_t task_addr = next_node
            - (uintptr_t)unwinder->async_debug_offsets.asyncio_task_object.task_node;

        if (process_single_task_node(unwinder, task_addr, NULL, result) < 0) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to process task node in awaited_by");
            return -1;
        }

        // Read next node
        if (_Py_RemoteDebug_PagedReadRemoteMemory(
                &unwinder->handle,
                next_node,
                sizeof(task_node),
                task_node) < 0) {
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read next task node in awaited_by");
            return -1;
        }
    }

    return 0;
}

int
append_awaited_by(
    RemoteUnwinderObject *unwinder,
    unsigned long tid,
    uintptr_t head_addr,
    PyObject *result)
{
    PyObject *tid_py = PyLong_FromUnsignedLong(tid);
    if (tid_py == NULL) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to create thread ID object");
        return -1;
    }

    PyObject* awaited_by_for_thread = PyList_New(0);
    if (awaited_by_for_thread == NULL) {
        Py_DECREF(tid_py);
        set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create awaited_by thread list");
        return -1;
    }

    RemoteDebuggingState *state = RemoteDebugging_GetStateFromObject((PyObject*)unwinder);
    PyObject *result_item = PyStructSequence_New(state->AwaitedInfo_Type);
    if (result_item == NULL) {
        Py_DECREF(tid_py);
        Py_DECREF(awaited_by_for_thread);
        set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create AwaitedInfo");
        return -1;
    }

    PyStructSequence_SetItem(result_item, 0, tid_py);  // steals ref
    PyStructSequence_SetItem(result_item, 1, awaited_by_for_thread);  // steals ref
    if (PyList_Append(result, result_item)) {
        Py_DECREF(result_item);
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to append awaited_by result item");
        return -1;
    }
    Py_DECREF(result_item);

    if (append_awaited_by_for_thread(unwinder, head_addr, awaited_by_for_thread))
    {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to append awaited_by for thread");
        return -1;
    }

    return 0;
}

/* ============================================================================
 * THREAD PROCESSOR FUNCTIONS FOR ASYNCIO
 * ============================================================================ */

int
process_thread_for_awaited_by(
    RemoteUnwinderObject *unwinder,
    uintptr_t thread_state_addr,
    unsigned long tid,
    void *context
) {
    PyObject *result = (PyObject *)context;
    uintptr_t head_addr = thread_state_addr + (uintptr_t)unwinder->async_debug_offsets.asyncio_thread_state.asyncio_tasks_head;
    return append_awaited_by(unwinder, tid, head_addr, result);
}

static int
process_running_task_chain(
    RemoteUnwinderObject *unwinder,
    uintptr_t running_task_addr,
    uintptr_t thread_state_addr,
    PyObject *result
) {
    uintptr_t running_task_frame = 0;
    if(get_task_frame(unwinder, running_task_addr, &running_task_frame) < 0) {
        return -1;
    }

    // First, add this task to the result
    PyObject *task_info = NULL;
    if (process_single_task_node(unwinder, running_task_addr, &task_info, result) < 0) {
        return -1;
    }

    // Get the chain from the current frame to this task
    PyObject *coro_chain = PyStructSequence_GET_ITEM(task_info, 2);
    assert(coro_chain != NULL);
    if (PyList_GET_SIZE(coro_chain) != 1) {
        PyErr_Format(PyExc_RuntimeError,
            "Expected single-item coro chain, got %zd items",
            PyList_GET_SIZE(coro_chain));
        set_exception_cause(unwinder, PyExc_RuntimeError, "Coro chain is not a single item");
        return -1;
    }
    PyObject *coro_info = PyList_GET_ITEM(coro_chain, 0);
    assert(coro_info != NULL);
    PyObject *frame_chain = PyStructSequence_GET_ITEM(coro_info, 0);
    assert(frame_chain != NULL);

    // Clear the coro_chain
    if (PyList_Clear(frame_chain) < 0) {
        set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to clear coroutine chain");
        return -1;
    }

    // Add the chain from the current frame to this task
    if (parse_async_frame_chain(unwinder, frame_chain, thread_state_addr, running_task_frame) < 0) {
        return -1;
    }

    // Now find all tasks that are waiting for this task and process them
    if (process_task_waiters(unwinder, result) < 0) {
        return -1;
    }

    return 0;
}

int
process_thread_for_async_stack_trace(
    RemoteUnwinderObject *unwinder,
    uintptr_t thread_state_addr,
    unsigned long tid,
    void *context
) {
    PyObject *result = (PyObject *)context;

    // Find running task in this thread
    uintptr_t running_task_addr;
    if (find_running_task_in_thread(unwinder, thread_state_addr, &running_task_addr) < 0) {
        return 0;
    }

    // If we found a running task, process it and its waiters
    if ((void*)running_task_addr != NULL) {
        PyObject *task_list = PyList_New(0);
        if (task_list == NULL) {
            set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create task list for thread");
            return -1;
        }

        if (process_running_task_chain(unwinder, running_task_addr, thread_state_addr, task_list) < 0) {
            Py_DECREF(task_list);
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to process running task chain");
            return -1;
        }

        // Create AwaitedInfo structure for this thread
        PyObject *tid_py = PyLong_FromUnsignedLong(tid);
        if (tid_py == NULL) {
            Py_DECREF(task_list);
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to create thread ID");
            return -1;
        }

        RemoteDebuggingState *state = RemoteDebugging_GetStateFromObject((PyObject*)unwinder);
        PyObject *awaited_info = PyStructSequence_New(state->AwaitedInfo_Type);
        if (awaited_info == NULL) {
            Py_DECREF(tid_py);
            Py_DECREF(task_list);
            set_exception_cause(unwinder, PyExc_MemoryError, "Failed to create AwaitedInfo");
            return -1;
        }

        PyStructSequence_SetItem(awaited_info, 0, tid_py);  // steals ref
        PyStructSequence_SetItem(awaited_info, 1, task_list);  // steals ref

        if (PyList_Append(result, awaited_info)) {
            Py_DECREF(awaited_info);
            set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to append AwaitedInfo to result");
            return -1;
        }
        Py_DECREF(awaited_info);
    }

    return 0;
}