use std::time::Duration;

#[cfg(target_os = "linux")]
use anyhow::Context;
use anyhow::Result;
use lazy_static::lazy_static;
use log::warn;
use prometheus::{IntGaugeVec, register_int_gauge_vec};
use tokio::time::MissedTickBehavior;

const PROCESS_MEMORY_COLLECT_INTERVAL: Duration = Duration::from_secs(10);
#[cfg(target_os = "linux")]
const PROCESS_MEMORY_STATUS_PATH: &str = "/proc/self/status";
#[cfg(target_os = "linux")]
const PROCESS_MEMORY_STATUS_FIELDS: &[&str] = &["VmRSS", "RssAnon", "RssFile", "RssShmem", "VmHWM"];

lazy_static! {
    pub static ref PROCESS_MEMORY_STATUS_GAUGE: IntGaugeVec = register_int_gauge_vec!(
        "tici_process_memory_status_bytes",
        "Process memory status fields from /proc/self/status in bytes",
        &["type"],
    )
    .unwrap_or_else(|err| panic!("failed to register tici_process_memory_status_bytes: {err}"));
}

pub fn monitor_process_memory_status() {
    std::mem::drop(tokio::spawn(async move {
        let mut interval = tokio::time::interval(PROCESS_MEMORY_COLLECT_INTERVAL);
        interval.set_missed_tick_behavior(MissedTickBehavior::Skip);

        loop {
            interval.tick().await;
            if let Err(err) = collect_process_memory_status() {
                warn!("collect process memory status failed, err={:?}", err);
            }
        }
    }));
}

#[cfg(target_os = "linux")]
fn collect_process_memory_status() -> Result<()> {
    let status = std::fs::read_to_string(PROCESS_MEMORY_STATUS_PATH)
        .with_context(|| format!("Failed to read {}", PROCESS_MEMORY_STATUS_PATH))?;
    for line in status.lines() {
        if let Some((memory_type, bytes)) = parse_process_memory_status_line(line) {
            set_process_memory_status_bytes(memory_type, bytes);
        }
    }
    Ok(())
}

#[cfg(not(target_os = "linux"))]
fn collect_process_memory_status() -> Result<()> {
    Ok(())
}

#[cfg(target_os = "linux")]
fn parse_process_memory_status_line(line: &str) -> Option<(&'static str, u64)> {
    for field in PROCESS_MEMORY_STATUS_FIELDS {
        let rest = line
            .strip_prefix(*field)
            .and_then(|rest| rest.strip_prefix(':'));
        if let Some(rest) = rest {
            let mut parts = rest.split_whitespace();
            let value = parts.next()?.parse::<u64>().ok()?;
            if parts.next()? != "kB" {
                return None;
            }
            return Some((*field, value.saturating_mul(1024)));
        }
    }
    None
}

#[cfg(target_os = "linux")]
fn set_process_memory_status_bytes(memory_type: &'static str, bytes: u64) {
    let value = bytes.min(i64::MAX as u64) as i64;
    PROCESS_MEMORY_STATUS_GAUGE
        .with_label_values(&[memory_type])
        .set(value);
}