#include <stdio.h>    // For printf, perror
#include <stdlib.h>   // For exit
#include <unistd.h>   // For sysconf
#include <sys/mman.h> // For mmap, munmap, PROT_*, MAP_*
#include <string.h>   // For strcpy, memset
#include <sanitizer/asan_interface.h>

// Function to print a separator for clarity in output
void print_separator(const char* title) {
    printf("\n--- %s ---\n", title);
}

int main() {
    long page_size = sysconf(_SC_PAGESIZE);
    if (page_size == -1) {
        perror("sysconf failed to get page size");
        exit(EXIT_FAILURE);
    }
    printf("System page size: %ld bytes\n", page_size);

    void *mem_ptr = NULL;
    char *char_ptr = NULL; // Use char* for byte-level access

    print_separator("1. Initial mmap allocation");
    // Allocate one page of anonymous, private, read/write memory
    mem_ptr = mmap(NULL, page_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);

    if (mem_ptr == MAP_FAILED) {
        perror("First mmap failed");
        exit(EXIT_FAILURE);
    }
    char_ptr = (char *)mem_ptr;
    printf("First mapped address: %p\n", mem_ptr);

    // Write some recognizable data to the first mapped region
    strcpy(char_ptr, "Hello from first map!");
    printf("Content after first map: '%s'\n", char_ptr);

    printf("Explictly poison the memory region\n");
    __asan_poison_memory_region(mem_ptr, page_size);

    print_separator("2. munmap() - Releasing memory and poisoning by ASan");
    // Unmap the memory region
    if (munmap(mem_ptr, page_size) == -1) {
        perror("munmap failed");
        exit(EXIT_FAILURE);
    }
    printf("Memory at %p unmapped.\n", mem_ptr);
    printf("ASan has now poisoned this memory region.\n");

    print_separator("3. Second mmap() - Remapping the same virtual address (MAP_FIXED)");
    // Try to remap the *exact same* virtual address.
    // MAP_FIXED is crucial for this, but it can fail if the address is not truly available.
    void *re_mapped_ptr = mmap(mem_ptr, page_size, PROT_READ | PROT_WRITE,
                               MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);

    if (re_mapped_ptr == MAP_FAILED) {
        perror("Second mmap (MAP_FIXED) failed. This might happen if the OS couldn't reuse the address.");
        // We'll try a regular mmap if MAP_FIXED fails, though it won't be the same address.
        printf("Trying regular mmap for demonstration purposes...\n");
        re_mapped_ptr = mmap(NULL, page_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (re_mapped_ptr == MAP_FAILED) {
            perror("Fallback mmap failed");
            exit(EXIT_FAILURE);
        }
    }

    char_ptr = (char *)re_mapped_ptr; // Update char_ptr to the new mapping
    printf("Second mapped address: %p (should be same as first if MAP_FIXED worked)\n", re_mapped_ptr);
    printf("ASan has now unpoisoned this memory region as it's a new valid mapping.\n");

    // Write new data to the remapped region
    memset(char_ptr, 0, page_size); // Clear previous content (if any)
    strcpy(char_ptr, "Hello from second map!");
    printf("Content after second map: '%s'\n", char_ptr);

    print_separator("4. Final Cleanup");
    if (munmap(re_mapped_ptr, page_size) == -1) {
        perror("Final munmap failed");
        exit(EXIT_FAILURE);
    }
    printf("Memory at %p unmapped again. Program exiting gracefully.\n", re_mapped_ptr);

    return 0;
}