use std::sync::{Mutex, OnceLock};

#[derive(Debug, Default, Clone, Copy)]
pub struct KernelProfileSnapshot {
    pub calls: u64,
    pub launches: u64,
    pub total_ns: u128,
    pub setup_ns: u128,
    pub copy_ns: u128,
    pub dispatch_ns: u128,
    pub transient_allocations: u64,
    pub metadata_upload_bytes: u64,
    pub metadata_reuse_hits: u64,
    pub metadata_reuse_bytes: u64,
    pub resident_rollout_steps: u64,
}

#[derive(Debug, Default)]
pub(crate) struct KernelProfileState {
    pub(crate) calls: u64,
    pub(crate) launches: u64,
    pub(crate) total_ns: u128,
    pub(crate) setup_ns: u128,
    pub(crate) copy_ns: u128,
    pub(crate) dispatch_ns: u128,
    pub(crate) transient_allocations: u64,
    pub(crate) metadata_upload_bytes: u64,
    pub(crate) metadata_reuse_hits: u64,
    pub(crate) metadata_reuse_bytes: u64,
    pub(crate) resident_rollout_steps: u64,
}

pub(crate) struct KernelProfileSite {
    inner: OnceLock<Mutex<KernelProfileState>>,
}

impl KernelProfileSite {
    pub const fn new() -> Self {
        Self {
            inner: OnceLock::new(),
        }
    }

    fn state(&self) -> &Mutex<KernelProfileState> {
        self.inner
            .get_or_init(|| Mutex::new(KernelProfileState::default()))
    }
}

pub(crate) fn profile_enabled() -> bool {
    std::env::var_os("DragonModel_STAGE_PROFILE").is_some()
}

pub(crate) fn profile_reset(site: &'static KernelProfileSite) {
    if let Ok(mut state) = site.state().lock() {
        *state = KernelProfileState::default();
    }
}

pub(crate) fn profile_snapshot(site: &'static KernelProfileSite) -> KernelProfileSnapshot {
    if let Ok(state) = site.state().lock() {
        return KernelProfileSnapshot {
            calls: state.calls,
            launches: state.launches,
            total_ns: state.total_ns,
            setup_ns: state.setup_ns,
            copy_ns: state.copy_ns,
            dispatch_ns: state.dispatch_ns,
            transient_allocations: state.transient_allocations,
            metadata_upload_bytes: state.metadata_upload_bytes,
            metadata_reuse_hits: state.metadata_reuse_hits,
            metadata_reuse_bytes: state.metadata_reuse_bytes,
            resident_rollout_steps: state.resident_rollout_steps,
        };
    }
    KernelProfileSnapshot::default()
}

pub(crate) fn profile_record(
    site: &'static KernelProfileSite,
    mutator: impl FnOnce(&mut KernelProfileState),
) {
    if let Ok(mut state) = site.state().lock() {
        mutator(&mut state);
    }
}
