use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::OnceLock;
use std::time::Duration;

#[derive(Clone, Copy)]
pub enum Stage {
    VcpuRun,
    DataAbort,
    Hvc,
    Svc,
    Vtimer,
    MmioSerialRead,
    MmioSerialWrite,
    MmioVsockRead,
    MmioVsockWrite,
    MmioVirtioRead,
    MmioVirtioWrite,
    MmioOtherRead,
    MmioOtherWrite,
}

struct Counter {
    calls: AtomicU64,
    us: AtomicU64,
}

impl Counter {
    const fn new() -> Self {
        Self {
            calls: AtomicU64::new(0),
            us: AtomicU64::new(0),
        }
    }
}

static VCPU_RUN: Counter = Counter::new();
static DATA_ABORT: Counter = Counter::new();
static HVC: Counter = Counter::new();
static SVC: Counter = Counter::new();
static VTIMER: Counter = Counter::new();
static MMIO_SERIAL_READ: Counter = Counter::new();
static MMIO_SERIAL_WRITE: Counter = Counter::new();
static MMIO_VSOCK_READ: Counter = Counter::new();
static MMIO_VSOCK_WRITE: Counter = Counter::new();
static MMIO_VIRTIO_READ: Counter = Counter::new();
static MMIO_VIRTIO_WRITE: Counter = Counter::new();
static MMIO_OTHER_READ: Counter = Counter::new();
static MMIO_OTHER_WRITE: Counter = Counter::new();

pub fn enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    *ENABLED.get_or_init(|| {
        let on = matches!(
            std::env::var("SUPERMACHINE_EXIT_PROFILE").as_deref(),
            Ok("1") | Ok("true") | Ok("yes") | Ok("on")
        );
        if on {
            start_reporter();
        }
        on
    })
}

pub fn record(stage: Stage, us: u64) {
    if !enabled() {
        return;
    }
    let c = counter(stage);
    c.calls.fetch_add(1, Ordering::Relaxed);
    c.us.fetch_add(us, Ordering::Relaxed);
}

pub fn mmio_stage(gpa: u64, write: bool) -> Stage {
    use crate::arch::aarch64::layout;
    if gpa >= layout::SERIAL_MMIO_BASE && gpa < layout::SERIAL_MMIO_BASE + layout::SERIAL_MMIO_SIZE
    {
        if write {
            Stage::MmioSerialWrite
        } else {
            Stage::MmioSerialRead
        }
    } else if gpa >= layout::VIRTIO_MMIO_BASE
        && gpa < layout::VIRTIO_MMIO_BASE + layout::VIRTIO_MMIO_STRIDE
    {
        if write {
            Stage::MmioVsockWrite
        } else {
            Stage::MmioVsockRead
        }
    } else if gpa >= layout::VIRTIO_MMIO_BASE {
        if write {
            Stage::MmioVirtioWrite
        } else {
            Stage::MmioVirtioRead
        }
    } else if write {
        Stage::MmioOtherWrite
    } else {
        Stage::MmioOtherRead
    }
}

fn counter(stage: Stage) -> &'static Counter {
    match stage {
        Stage::VcpuRun => &VCPU_RUN,
        Stage::DataAbort => &DATA_ABORT,
        Stage::Hvc => &HVC,
        Stage::Svc => &SVC,
        Stage::Vtimer => &VTIMER,
        Stage::MmioSerialRead => &MMIO_SERIAL_READ,
        Stage::MmioSerialWrite => &MMIO_SERIAL_WRITE,
        Stage::MmioVsockRead => &MMIO_VSOCK_READ,
        Stage::MmioVsockWrite => &MMIO_VSOCK_WRITE,
        Stage::MmioVirtioRead => &MMIO_VIRTIO_READ,
        Stage::MmioVirtioWrite => &MMIO_VIRTIO_WRITE,
        Stage::MmioOtherRead => &MMIO_OTHER_READ,
        Stage::MmioOtherWrite => &MMIO_OTHER_WRITE,
    }
}

fn start_reporter() {
    static STARTED: OnceLock<()> = OnceLock::new();
    STARTED.get_or_init(|| {
        std::thread::Builder::new()
            .name("exit-profile".into())
            .spawn(|| {
                let mut prev = snapshot();
                loop {
                    std::thread::sleep(Duration::from_secs(1));
                    let now = snapshot();
                    eprintln!("[exit-profile] {}", format_delta(&prev, &now));
                    prev = now;
                }
            })
            .ok();
    });
}

#[derive(Clone, Copy)]
struct Snap {
    calls: [u64; 13],
    us: [u64; 13],
}

fn snapshot() -> Snap {
    let counters = [
        &VCPU_RUN,
        &DATA_ABORT,
        &HVC,
        &SVC,
        &VTIMER,
        &MMIO_SERIAL_READ,
        &MMIO_SERIAL_WRITE,
        &MMIO_VSOCK_READ,
        &MMIO_VSOCK_WRITE,
        &MMIO_VIRTIO_READ,
        &MMIO_VIRTIO_WRITE,
        &MMIO_OTHER_READ,
        &MMIO_OTHER_WRITE,
    ];
    let mut calls = [0; 13];
    let mut us = [0; 13];
    for (i, c) in counters.iter().enumerate() {
        calls[i] = c.calls.load(Ordering::Relaxed);
        us[i] = c.us.load(Ordering::Relaxed);
    }
    Snap { calls, us }
}

fn format_delta(prev: &Snap, now: &Snap) -> String {
    const NAMES: [&str; 13] = [
        "vcpu_run",
        "data_abort",
        "hvc",
        "svc",
        "vtimer",
        "mmio_serial_r",
        "mmio_serial_w",
        "mmio_vsock_r",
        "mmio_vsock_w",
        "mmio_virtio_r",
        "mmio_virtio_w",
        "mmio_other_r",
        "mmio_other_w",
    ];
    let mut out = String::new();
    for i in 0..NAMES.len() {
        let calls = now.calls[i].saturating_sub(prev.calls[i]);
        if calls == 0 {
            continue;
        }
        let us = now.us[i].saturating_sub(prev.us[i]);
        if !out.is_empty() {
            out.push_str(" | ");
        }
        out.push_str(&format!(
            "{} calls={} avg_us={:.2}",
            NAMES[i],
            calls,
            us as f64 / calls as f64
        ));
    }
    if out.is_empty() {
        "idle".to_owned()
    } else {
        out
    }
}
