Overflow_Page :: struct {
    next: *Overflow_Page;
    allocator: Allocator;
    size: s64;
}

Temporary_Storage_Paged_Alloc :: struct {
    #as using base: Temporary_Storage;

    #overlay (implementation_data) using derived: struct {
        overflow_allocator : Allocator;

        overflow_pages: *Overflow_Page;
        original_data: *u8;  // Data to restore after clearing overflow pages. @Simplify: Maybe could be an Overflow_Page, but maybe we want to be able to assert on overflow_pages==null to ensure performance.
        original_size: s64;

        current_page_bytes_occupied: s64;
    }
}

Temporary_Storage_State :: struct {
    top_overflow_page: *Overflow_Page;
    current_page_bytes_occupied: s64;
    total_bytes_occupied: s64;
}

temporary_allocator_proc :: (mode: Allocator_Mode, requested_size: s64, old_size: s64, old_memory: *void, allocator_data: *void) -> *void {
    if #complete mode == {
      case .FREE;
        return null;
      case .RESIZE;
        real_old_size := (old_size + 7) & ~7; // We know old_size was padded
        if requested_size <= real_old_size return old_memory;
        #through;
      case .ALLOCATE;
        // Pad to 8-byte alignment.
        size := (requested_size + 7) & ~7;

        ts := cast(*Temporary_Storage_Paged_Alloc) context.temporary_storage;

        // If someone gave us a context with no temporary_storage, fall back to context.allocator.
        if !ts {
            a := context.allocator;
            if a.proc == temporary_allocator_proc {
                // context.allocator is temporary_allocator, so, fall back to default_allocator.
                a = Context_Base.default_allocator;
            }

            return a.proc(mode, requested_size, old_size, old_memory, a.data);
        }

        if size > ts.size - ts.current_page_bytes_occupied {
            success := add_new_page(ts, size);
            if !success  return null;
        }

        result := ts.data + ts.current_page_bytes_occupied;
        if result && (mode == .RESIZE) {
            memcpy(result, old_memory, min(old_size, requested_size));
        }

        ts.current_page_bytes_occupied += size;
        ts.total_bytes_occupied += size;
        ts.high_water_mark = max(ts.high_water_mark, ts.total_bytes_occupied);

        return result;

      case .STARTUP;      #through;
      case .SHUTDOWN;
        return null;

      case .THREAD_START; #through;
      case .THREAD_STOP;
        assert(false, "Multithreaded access is not supported by temporary_allocator_proc.\n");
        return null;

      case .CREATE_HEAP; #through;
      case .DESTROY_HEAP;
        assert(false, "Create/Destroy heap are not supported by temporary_allocator_proc.\n");
        return null;

      case .IS_THIS_YOURS;
        ts := cast(*Temporary_Storage_Paged_Alloc) context.temporary_storage;
        if !ts return null;

        if (old_memory >= ts.data) && (old_memory < ts.data + ts.size) return cast(*void) true;
        if (old_memory >= ts.original_data) && (old_memory < ts.original_data + ts.original_size) return cast(*void) true;

        page := ts.overflow_pages;
        while page {
            if (old_memory >= page) && (old_memory < cast(*void)page + page.size + ALIGNED_OVERFLOW_PAGE_ALLOCATION) return cast(*void) true;
            page = page.next;
        }

        // Not found!
        return null;

      case .CAPS;
//        if old_memory { (.*) cast(*string)old_memory = "modules/Basic temporary_allocator_proc"; }
        if old_memory { (cast(*string)old_memory).* = "modules/Basic temporary_allocator_proc"; }
        return cast(*void)(Allocator_Caps.HINT_I_AM_PER_FRAME_TEMPORARY_STORAGE);
    }
}

auto_release_temp :: (loc := #caller_location) #expand {
    __mark := get_temporary_storage_mark();
    `defer set_temporary_storage_mark(__mark, loc);
}

get_temporary_storage_mark :: inline () -> Temporary_Storage_State {
    ts := cast(*Temporary_Storage_Paged_Alloc) context.temporary_storage;
    if !ts return .{};  // Just to have something to return!

    state: Temporary_Storage_State = ---;
    state.current_page_bytes_occupied = ts.current_page_bytes_occupied;
    state.total_bytes_occupied        = ts.total_bytes_occupied;
    state.top_overflow_page           = ts.overflow_pages;

    return state;
}

// Typically used as follows:
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// mark := get_temporary_storage_mark();
// defer set_temporary_storage_mark(mark);
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// You can also use the `auto_release_temp` macro.
set_temporary_storage_mark :: inline (state: Temporary_Storage_State, loc := #caller_location) {
    ts := cast(*Temporary_Storage_Paged_Alloc) context.temporary_storage;
    if !ts return;

    ts.last_set_mark_location = loc;

    if state.top_overflow_page {
        free_pages_down_to(ts, state.top_overflow_page);

        page := state.top_overflow_page;
        ts.data = (cast(*u8) page) + ALIGNED_OVERFLOW_PAGE_ALLOCATION;
        ts.size = page.size;
        assert(page.size >= 0);
    } else {
        if ts.overflow_pages {
            free_pages_down_to(ts, null);
            ts.data = ts.original_data;
            ts.size = ts.original_size;
        }
        assert(ts.overflow_pages == null);
    }

    ts.current_page_bytes_occupied = state.current_page_bytes_occupied;
    ts.total_bytes_occupied        = state.total_bytes_occupied;

    #if TEMP_ALLOCATOR_OVERWRITE_MEMORY {
        memset(cast(*u8) ts.data + ts.current_page_bytes_occupied, SET_STORAGE_MARK_STAMP, ts.size - ts.current_page_bytes_occupied);
    }
}

reset_temporary_storage :: (_ts: *Temporary_Storage) {
    ts := cast(*Temporary_Storage_Paged_Alloc)_ts;

    if ts.overflow_pages {
        ts.data = ts.original_data;
        ts.size = ts.original_size;
    }
    ts.current_page_bytes_occupied = 0;
    ts.total_bytes_occupied = 0;
    ts.high_water_mark = 0;

    // @Speed: Maybe we should keep the allocated pages on the side,
    // rather than freeing them every frame. It probably would not
    // add much complexity here.
    free_pages_down_to(ts, null);

    #if TEMP_ALLOCATOR_OVERWRITE_MEMORY {
        memset(ts.data, RESET_STORAGE_STAMP, ts.size);
    }
}

free_pages_down_to :: (ts: *Temporary_Storage_Paged_Alloc, page: *Overflow_Page) {
    cursor := ts.overflow_pages;
    while cursor != page {
        if !cursor {
            log_error("Attempt to free down to page %, but that page was not found! This is probably an error in using set_temporary_storage_mark, or a general memory corruption.\n", page);
            ts.overflow_pages = null;
            return;
        }

        next := cursor.next;
        allocator := cursor.allocator;

        #if TEMP_ALLOCATOR_OVERWRITE_MEMORY {
            size_including_header := get_size_including_overflow_page_header(cursor.size);
            memset(cursor, FREE_PAGE_STAMP, size_including_header);
        }

        free(cursor,, allocator=allocator);
        cursor = next;
    }

    ts.overflow_pages = page;
}

ALIGNED_OVERFLOW_PAGE_ALLOCATION :: 32;

get_size_including_overflow_page_header :: (data_size: s64) -> s64 {
    size := data_size + ALIGNED_OVERFLOW_PAGE_ALLOCATION;
    // Pad to 8-byte alignment.
    size = (size + 7) & ~7;

    return size;
}

add_new_page :: (ts: *Temporary_Storage_Paged_Alloc, minimum_size: s64) -> bool {
    #assert,stallable(size_of(Overflow_Page) <= ALIGNED_OVERFLOW_PAGE_ALLOCATION);

    data_size : s64 = ts.original_size - ALIGNED_OVERFLOW_PAGE_ALLOCATION;
    if data_size < minimum_size  data_size = minimum_size;

    size_including_header := get_size_including_overflow_page_header(data_size);

    // It's legal for overflow_allocator to be null, in which case
    // we do whatever alloc does.
    allocator := ts.overflow_allocator;
    if !allocator.proc {
        // We don't know what allocator might be pushed onto the context,
        // so, fall back to default_allocator.
        allocator = Context_Base.default_allocator;
    }

    memory := alloc(size_including_header,, allocator=allocator);
    if !memory {
        log_error("Failed to allocate % bytes for a Temporary_Storage overflow page.\n", size_including_header);
        return false;  // This is going to be bad!
    }

    if ts.overflow_pages == null {
        // Save out the original size and data,
        // so we can restore them later.
        ts.original_data = ts.data;
        ts.original_size = ts.size;
    }

    page := cast(*Overflow_Page) memory;
    page.next           = ts.overflow_pages;
    page.allocator      = allocator;
    page.size           = data_size;

    ts.overflow_pages = page;

    ts.data = memory + ALIGNED_OVERFLOW_PAGE_ALLOCATION;
    ts.size = data_size;
    ts.current_page_bytes_occupied = 0;
    assert(ts.size >= 0);

    return true;
}