// SPDX-License-Identifier: GPL-2.0
// Copyright (c) 2024 Dennis Kobert <dennis@kobert.dev>

//! Perf Event.
//!
//! C header: [`include/linux/perf_event.h`](srctree/include/linux/perf_event.h)

use crate::alloc::AllocError;
use crate::error::from_err_ptr;
use crate::prelude::*;
use crate::types::ARef;
use crate::{bindings, task::Task, types::Opaque};

pub mod events;
use events::*;

/// Perf Event Attr wrapper
#[repr(transparent)]
pub struct PerfEventAttr {
    inner: Opaque<bindings::perf_event_attr>,
}

impl PerfEventAttr {
    /// Create a new PerfEventAttr from raw attr
    pub fn from_raw(attr: bindings::perf_event_attr) -> Self {
        Self {
            inner: Opaque::new(attr),
        }
    }

    /// Get a mutable pointer to the inner attr
    pub fn as_inner(&self) -> *mut bindings::perf_event_attr {
        self.inner.get()
    }
}

/// Wrapper for sample data
pub struct SampleData {
    inner: Opaque<bindings::perf_sample_data>,
}

impl SampleData {
    /// Returns a reference to the underlying data
    pub fn get(&self) -> Option<&bindings::perf_sample_data> {
        let ptr = self.inner.get();
        if ptr.is_null() {
            return None;
        }
        Some(unsafe { &*ptr })
    }
}

/// Wrapper for the current register values when the overflow handler is called
pub struct Registers {
    inner: Opaque<bindings::pt_regs>,
}

impl Registers {
    /// Returns a reference to the underlying data
    pub fn get(&self) -> Option<&bindings::pt_regs> {
        let ptr = self.inner.get();
        if ptr.is_null() {
            return None;
        }
        Some(unsafe { &*ptr })
    }
}

/// Handler function for overflow events

/// Perf Event wrapper
#[repr(transparent)]
pub struct PerfEvent {
    inner: *mut bindings::perf_event,
}

/// Perf Event wrapper
#[repr(transparent)]
pub struct PerfEventRef {
    inner: *const bindings::perf_event,
}

// SAFETY: perf_event has internal locking for thread safety
unsafe impl Send for PerfEvent {}
unsafe impl Sync for PerfEvent {}

impl PerfEvent {
    /// Returns a raw pointer to the inner C struct.
    #[inline]
    pub fn as_ptr(&self) -> *mut bindings::perf_event {
        self.inner
    }

    /// Enable the event for counting
    pub fn enable(&self) {
        unsafe {
            bindings::perf_event_enable(self.as_ptr());
        }
    }

    /// Disable the event
    pub fn disable(&self) {
        unsafe {
            bindings::perf_event_disable(self.as_ptr());
        }
    }

    /// Read the current value of the event counter
    pub fn read(&self) -> u64 {
        let mut enabled = 0;
        let mut running = 0;

        unsafe { bindings::perf_event_read_value(self.as_ptr(), &mut enabled, &mut running) }
    }
}

impl Drop for PerfEvent {
    fn drop(&mut self) {
        if !self.inner.is_null() {
            let context_ptr =
                unsafe { *self.inner }.overflow_handler_context as *mut OverflowHandler;
            if !context_ptr.is_null() {
                let OverflowHandler { closure, dyn_fn } = unsafe { context_ptr.read() };
                let _ = dyn_fn;
                unsafe { KBox::from_raw(closure) };
            }
        }

        unsafe {
            bindings::perf_event_release_kernel(self.inner);
        }
    }
}

/// Builder for configuring a performance event
pub struct EventBuilder {
    event_type: EventType,
    sample_period: Option<u64>,
    sample_freq: Option<u64>,
    disabled: bool,
    inherit: bool,
    pinned: bool,
    exclusive: bool,
    exclude_user: bool,
    exclude_kernel: bool,
    exclude_hv: bool,
    exclude_idle: bool,
    cpu: Option<i32>,
    task: Option<ARef<Task>>,
    overflow_handler: Option<OverflowHandler>,
}

/// Error type for performance event operations
#[derive(Debug)]
pub enum Error {
    /// Event creation failed
    InvalidConfig,
    /// Invalid CPU specified
    InvalidCpu,
    /// Invalid task specified
    InvalidTask,
    /// MemoryAllocation Error
    Alloc(AllocError),
}

impl From<AllocError> for Error {
    fn from(value: AllocError) -> Self {
        Self::Alloc(value)
    }
}

impl EventBuilder {
    /// Create a new event builder for the given event type
    pub fn new(event_type: impl Into<EventType>) -> Self {
        Self {
            event_type: event_type.into(),
            sample_period: None,
            sample_freq: None,
            disabled: false,
            inherit: false,
            pinned: false,
            exclusive: false,
            exclude_user: false,
            exclude_kernel: false,
            exclude_hv: false,
            exclude_idle: false,
            cpu: None,
            task: None,
            overflow_handler: None,
        }
    }

    /// Set the sampling period (number of events between samples)
    pub fn sample_period(mut self, period: u64) -> Self {
        self.sample_period = Some(period);
        self.sample_freq = None; // Period and frequency are mutually exclusive
        self
    }

    /// Set the sampling frequency (samples per second)
    pub fn sample_freq(mut self, freq: u64) -> Self {
        self.sample_freq = Some(freq);
        self.sample_period = None; // Period and frequency are mutually exclusive
        self
    }

    /// Start the event disabled (must be explicitly enabled)
    pub fn disabled(mut self) -> Self {
        self.disabled = true;
        self
    }

    /// Child tasks inherit this event
    pub fn inherit(mut self) -> Self {
        self.inherit = true;
        self
    }

    /// Event must always be on PMU
    pub fn pinned(mut self) -> Self {
        self.pinned = true;
        self
    }

    /// Only group on PMU
    pub fn exclusive(mut self) -> Self {
        self.exclusive = true;
        self
    }

    /// Don't count user-space events
    pub fn exclude_user(mut self) -> Self {
        self.exclude_user = true;
        self
    }

    /// Don't count kernel events
    pub fn exclude_kernel(mut self) -> Self {
        self.exclude_kernel = true;
        self
    }

    /// Don't count hypervisor events
    pub fn exclude_hv(mut self) -> Self {
        self.exclude_hv = true;
        self
    }

    /// Don't count when CPU is idle
    pub fn exclude_idle(mut self) -> Self {
        self.exclude_idle = true;
        self
    }

    /// Monitor events on a specific CPU (-1 for all CPUs)
    pub fn cpu(mut self, cpu: i32) -> Self {
        self.cpu = Some(cpu);
        self
    }

    /// Monitor events for a specific task (None for per-CPU mode)
    pub fn task(mut self, task: ARef<Task>) -> Self {
        self.task = Some(task);
        self
    }

    /// Set handler for overflow events
    pub fn on_overflow(mut self, handler: OverflowHandler) -> Self {
        self.overflow_handler = Some(handler);
        self
    }

    /// Build the perf event
    pub fn build(self) -> Result<PerfEvent, Error> {
        // Create the perf_event_attr structure
        let mut attr = bindings::perf_event_attr::default();

        // Configure event-specific attributes
        self.event_type.configure_attr(&mut attr);
        attr.size = core::mem::size_of::<bindings::perf_event_attr>() as u32;

        // Set sampling configuration
        if let Some(period) = self.sample_period {
            attr.__bindgen_anon_1.sample_period = period;
        } else if let Some(freq) = self.sample_freq {
            attr.__bindgen_anon_1.sample_freq = freq;
            attr.set_freq(1);
        }

        // Set the configuration bits using the bindgen-generated setters
        attr.set_disabled(self.disabled as u64);
        attr.set_inherit(self.inherit as u64);
        attr.set_pinned(self.pinned as u64);
        attr.set_exclusive(self.exclusive as u64);
        attr.set_exclude_user(self.exclude_user as u64);
        attr.set_exclude_kernel(self.exclude_kernel as u64);
        attr.set_exclude_hv(self.exclude_hv as u64);
        attr.set_exclude_idle(self.exclude_idle as u64);

        let perf_event_attr = PerfEventAttr::from_raw(attr);
        let cpu = self.cpu.unwrap_or(-1);

        // Create the perf event using the existing kernel interface
        let event = perf_event_create_kernel_counter(
            perf_event_attr,
            cpu,
            self.task,
            self.overflow_handler,
        )
        .map_err(|_| Error::InvalidConfig)?;

        Ok(event)
    }
}

/// Registers a new perf event counter
///
/// # Arguments
/// * `attr`: attributes of the counter to create
/// * `cpu`: cpu to which the counter is bound (-1 for all CPUs)
/// * `task`: task to profile (None for per-cpu)
/// * `overflow_handler`: callback to trigger when we hit the event
pub fn perf_event_create_kernel_counter(
    perf_event_attr: PerfEventAttr,
    cpu: i32,
    task: Option<ARef<Task>>,
    overflow: Option<OverflowHandler>,
) -> Result<PerfEvent, Error> {
    // Convert handler to C callback if provided
    // Create the perf event using kernel functions
    let raw_perf_event = unsafe {
        bindings::perf_event_create_kernel_counter(
            perf_event_attr.as_inner(),
            cpu,
            task.as_ref().map_or(core::ptr::null_mut(), |t| t.as_ptr()),
            overflow.is_some().then_some(overflow_trampoline),
            overflow.map_or(core::ptr::null_mut(), |x| {
                KBox::into_raw(KBox::new(x, crate::alloc::flags::GFP_KERNEL).unwrap())
                    as *mut crate::ffi::c_void
            }),
        )
    };

    if raw_perf_event.is_null() {
        pr_err!("event null");
        return Err(Error::InvalidConfig);
    }

    let result = from_err_ptr(raw_perf_event);
    match result {
        Err(e) => {
            pr_err!("Encountered error during creation of perf event");
            pr_err!("Error: {e:?}");
            Err(Error::InvalidConfig)
        }
        Ok(raw_event) => Ok(PerfEvent { inner: raw_event }),
    }
}

unsafe extern "C" fn overflow_trampoline(
    perf_event: *mut bindings::perf_event,
    sample_data: *mut bindings::perf_sample_data,
    registers: *mut bindings::pt_regs,
) {
    if perf_event.is_null() {
        return;
    }
    let context_ptr = unsafe { *perf_event }.overflow_handler_context as *mut OverflowHandler;
    if context_ptr.is_null() {
        return;
    }
    let context = &mut unsafe { context_ptr.read() };
    overflow_wrapper(perf_event, sample_data, registers, &mut *context.dyn_fn);
}
fn overflow_wrapper(
    perf_event: *mut bindings::perf_event,
    sample_data: *mut bindings::perf_sample_data,
    registers: *mut bindings::pt_regs,
    mut handler: impl FnMut(&PerfEventRef, &mut SampleData, &mut Registers),
) {
    if perf_event.is_null() || sample_data.is_null() || registers.is_null() {
        return;
    }
    handler(
        &PerfEventRef { inner: perf_event },
        &mut SampleData {
            inner: Opaque::new(unsafe { *sample_data }),
        },
        &mut Registers {
            inner: Opaque::new(unsafe { *registers }),
        },
    )
}
fn into_dyn(
    handler: impl Fn(&PerfEventRef, &mut SampleData, &mut Registers) + Send + Sync + 'static,
) -> Result<OverflowHandler, Error> {
    let b = KBox::new(handler, crate::alloc::flags::GFP_KERNEL)?;
    let b = Box::leak(b);
    let b_ptr = (b as *mut _) as *mut crate::ffi::c_void;
    let c = b as &'static mut (dyn FnMut(&PerfEventRef, &mut SampleData, &mut Registers)
                      + Send
                      + Sync
                      + 'static);

    Ok(OverflowHandler {
        closure: b_ptr,
        dyn_fn: c,
    })
}
/// Workaround for the missing support of using KBox as a fat pointer
pub struct OverflowHandler {
    closure: *mut crate::ffi::c_void,
    dyn_fn: &'static mut (dyn FnMut(&PerfEventRef, &mut SampleData, &mut Registers)
                      + Send
                      + Sync
                      + 'static),
}

impl OverflowHandler {
    /// Constructs a new overflow handler callback which is run when a performance counter overflows.
    ///
    /// # Safety
    /// The callback function is run in an NMI context:
    /// - Handler must be interrupt-safe
    /// - Handler must not block
    /// - Handler must not alloc
    /// - Handler must not panic
    pub unsafe fn new(
        handler: impl Fn(&PerfEventRef, &mut SampleData, &mut Registers) + Send + Sync + 'static,
    ) -> Result<Self, Error> {
        into_dyn(handler)
    }
}