// Copyright (c) 2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module profiler

import time

// profiler_alloc is the allocation function that records allocations
fn profiler_alloc(size int, ctx voidptr) voidptr {
    ptr := unsafe { C.malloc(size) }
    if ptr == unsafe { nil } {
        return ptr
    }

    if profiler_state.mu == unsafe { nil } || !profiler_state.enabled {
        return ptr
    }

    profiler_state.mu.lock()
    defer {
        profiler_state.mu.unlock()
    }

    idx := profiler_state.allocs.len
    profiler_state.allocs << AllocRecord{
        ptr:       ptr
        size:      size
        frame:     profiler_state.current_frame
        file:      '' // Will be set by caller if needed
        line:      0
        timestamp: time.sys_mono_now()
        freed:     false
    }
    profiler_state.alloc_map[ptr] = idx
    profiler_state.live_bytes += u64(size)
    profiler_state.total_allocs++

    if profiler_state.live_bytes > profiler_state.peak_bytes {
        profiler_state.peak_bytes = profiler_state.live_bytes
    }

    // Add to current frame's new allocs
    if profiler_state.frames.len > 0 {
        frame_idx := profiler_state.frames.len - 1
        profiler_state.frames[frame_idx].new_allocs << idx
        profiler_state.frames[frame_idx].new_bytes += u64(size)
    }

    return ptr
}

// profiler_alloc_with_location records allocation with source location
pub fn profiler_alloc_with_location(size int, file string, line int) voidptr {
    ptr := unsafe { C.malloc(size) }
    if ptr == unsafe { nil } {
        return ptr
    }

    if profiler_state.mu == unsafe { nil } || !profiler_state.enabled {
        return ptr
    }

    profiler_state.mu.lock()
    defer {
        profiler_state.mu.unlock()
    }

    idx := profiler_state.allocs.len
    profiler_state.allocs << AllocRecord{
        ptr:       ptr
        size:      size
        frame:     profiler_state.current_frame
        file:      file
        line:      line
        timestamp: time.sys_mono_now()
        freed:     false
    }
    profiler_state.alloc_map[ptr] = idx
    profiler_state.live_bytes += u64(size)
    profiler_state.total_allocs++

    if profiler_state.live_bytes > profiler_state.peak_bytes {
        profiler_state.peak_bytes = profiler_state.live_bytes
    }

    // Add to current frame's new allocs
    if profiler_state.frames.len > 0 {
        frame_idx := profiler_state.frames.len - 1
        profiler_state.frames[frame_idx].new_allocs << idx
        profiler_state.frames[frame_idx].new_bytes += u64(size)
    }

    return ptr
}

// profiler_free records a free operation
fn profiler_free(ptr voidptr, ctx voidptr) {
    if ptr == unsafe { nil } {
        return
    }

    if profiler_state.mu != unsafe { nil } && profiler_state.enabled {
        profiler_state.mu.lock()

        if idx := profiler_state.alloc_map[ptr] {
            profiler_state.allocs[idx].freed = true
            profiler_state.allocs[idx].free_frame = profiler_state.current_frame
            alloc_size := u64(profiler_state.allocs[idx].size)
            if profiler_state.live_bytes >= alloc_size {
                profiler_state.live_bytes -= alloc_size
            }
            profiler_state.total_frees++

            // Add to current frame's freed list
            if profiler_state.frames.len > 0 {
                frame_idx := profiler_state.frames.len - 1
                profiler_state.frames[frame_idx].freed_idxs << idx
                profiler_state.frames[frame_idx].freed_bytes += alloc_size
            }

            // Remove from map
            profiler_state.alloc_map.delete(ptr)
        }

        profiler_state.mu.unlock()
    }

    unsafe { C.free(ptr) }
}

// profiler_realloc records a realloc operation
fn profiler_realloc(ptr voidptr, new_size int, ctx voidptr) voidptr {
    if ptr == unsafe { nil } {
        return profiler_alloc(new_size, ctx)
    }
    if new_size == 0 {
        profiler_free(ptr, ctx)
        return unsafe { nil }
    }

    // Record the free of old allocation
    old_size := 0
    if profiler_state.mu != unsafe { nil } && profiler_state.enabled {
        profiler_state.mu.lock()
        if idx := profiler_state.alloc_map[ptr] {
            unsafe {
                *(&old_size) = profiler_state.allocs[idx].size
            }
        }
        profiler_state.mu.unlock()
    }

    new_ptr := unsafe { C.realloc(ptr, new_size) }
    if new_ptr == unsafe { nil } {
        return new_ptr
    }

    if profiler_state.mu != unsafe { nil } && profiler_state.enabled {
        profiler_state.mu.lock()
        defer {
            profiler_state.mu.unlock()
        }

        // Remove old entry
        if idx := profiler_state.alloc_map[ptr] {
            profiler_state.allocs[idx].freed = true
            profiler_state.allocs[idx].free_frame = profiler_state.current_frame
            if profiler_state.live_bytes >= u64(old_size) {
                profiler_state.live_bytes -= u64(old_size)
            }
            profiler_state.alloc_map.delete(ptr)
        }

        // Add new entry
        new_idx := profiler_state.allocs.len
        profiler_state.allocs << AllocRecord{
            ptr:       new_ptr
            size:      new_size
            frame:     profiler_state.current_frame
            file:      ''
            line:      0
            timestamp: time.sys_mono_now()
            freed:     false
        }
        profiler_state.alloc_map[new_ptr] = new_idx
        profiler_state.live_bytes += u64(new_size)

        if profiler_state.live_bytes > profiler_state.peak_bytes {
            profiler_state.peak_bytes = profiler_state.live_bytes
        }
    }

    return new_ptr
}

// profiler_allocator returns an Allocator that tracks all allocations
pub fn profiler_allocator() Allocator {
    return Allocator{
        alloc_fn:   profiler_alloc
        free_fn:    profiler_free
        realloc_fn: profiler_realloc
        ctx:        unsafe { nil }
    }
}