// Status: minimal — TX FIFO direct-to stdout-or-sink, RX permanently
// empty, line-aware heartbeat detection.
//
// PL011 register map (subset we implement):
//   0x000  UARTDR     data register; write = TX byte, read = RX byte
//   0x018  UARTFR     flag register; bit 7 TXFE, bit 4 RXFE always
//                     report TXFE=1 (TX always empty/ready), RXFE=1
//                     (RX always empty)
//   0x024  UARTIBRD   integer baud divisor (write-ignored)
//   0x028  UARTFBRD   fractional baud divisor (write-ignored)
//   0x02c  UARTLCR_H  line control (write-ignored)
//   0x030  UARTCR     control (write-ignored)
//   0x038  UARTIMSC   interrupt mask
//   0xfe0  UARTPeriphID0 = 0x11 (PL011 magic)
//   0xfe4  UARTPeriphID1 = 0x10
//   0xfe8  UARTPeriphID2 = 0x14 (rev 4)
//   0xfec  UARTPeriphID3 = 0x00
//   0xff0  UARTPCellID0  = 0x0d
//   0xff4  UARTPCellID1  = 0xf0
//   0xff8  UARTPCellID2  = 0x05
//   0xffc  UARTPCellID3  = 0xb1

use std::io::Write;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Mutex, OnceLock};

use super::mmio_bus::MmioDevice;

/// Process-global heartbeat count. Bumped each time a complete line
/// containing the literal "heartbeat counter=" arrives on PL011.
/// Used by `--snapshot-at N` (capture once HEARTBEAT_COUNT >= N) and
/// `--quiesce-ms` (post-heartbeat wait-for-WFI).
pub static HEARTBEAT_COUNT: AtomicU64 = AtomicU64::new(0);

/// Set when a complete line containing "parking PID 1" arrives on
/// PL011 — i.e. init-oci's workload exited and pid 1 is now idle.
/// The bake worker uses this as an "early fallback" snapshot trigger
/// for non-service images: instead of waiting out the full
/// `--snapshot-after-ms` timeout (default 7 s) when no listener
/// appears, snapshot as soon as init reports parking. Service images
/// (nginx, redis, etc.) never trip this — their workload doesn't exit.
pub static WORKLOAD_PARKED: AtomicBool = AtomicBool::new(false);

/// Set when a complete line containing "init-oci: workload-pre-exec"
/// arrives on PL011 — i.e. init-oci has finished its setup and is
/// about to fork+exec the workload (it then nanosleeps 100 ms to
/// give us a stable WFI window). Used as the bake-ready trigger for
/// the always-pipelined-skip-warm `.build()` path: capture the
/// guest BEFORE the workload runs, so each restore re-executes the
/// workload fresh (which is what agent-runtime users want). Saves
/// 50-150 ms of bake time vs waiting for the workload's listener.
///
/// Gated host-side by the bake driver's `on_pre_exec` flag — only
/// active for skip_warm_snapshot=true bakes. With_warmup + service-
/// image bakes leave `on_pre_exec=false` so this marker is detected
/// but ignored, preserving the listener-ready snapshot semantics.
pub static PRE_EXEC_READY: AtomicBool = AtomicBool::new(false);

/// Optional sink: when set, PL011 bytes go here instead of stdout.
/// Wired by `--log-sink FILE` at startup.
pub static LOG_SINK: Mutex<Option<std::fs::File>> = Mutex::new(None);
static LOG_SINK_SET: AtomicBool = AtomicBool::new(false);
static HEARTBEAT_DETECTION_ENABLED: AtomicBool = AtomicBool::new(false);

/// Sliding buffer of the current PL011 output line, reset on '\n'.
static PL011_LINE: Mutex<Vec<u8>> = Mutex::new(Vec::new());

/// Marker substring we look for to bump HEARTBEAT_COUNT. The match is
/// `windows().any(...)` so it can sit anywhere in the line.
const HEARTBEAT_MARKER: &[u8] = b"heartbeat counter=";

/// Marker for [`WORKLOAD_PARKED`] — init-oci writes "parking PID 1
/// (exit=N)" when its forked workload exits and pid 1 goes idle.
const PARKED_MARKER: &[u8] = b"parking PID 1";

/// Marker for [`PRE_EXEC_READY`] — init-oci writes
/// "init-oci: workload-pre-exec" right before forking its workload
/// child (then nanosleeps 100 ms so the host can capture).
const PRE_EXEC_MARKER: &[u8] = b"init-oci: workload-pre-exec";

fn console_log_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    *ENABLED.get_or_init(|| {
        !matches!(
            std::env::var("SUPERMACHINE_CONSOLE_LOG").as_deref(),
            Ok("0") | Ok("false") | Ok("no") | Ok("off")
        )
    })
}

pub fn set_log_sink(path: &str) -> std::io::Result<()> {
    let f = std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(path)?;
    *LOG_SINK.lock().unwrap() = Some(f);
    LOG_SINK_SET.store(true, Ordering::Release);
    Ok(())
}

pub fn set_heartbeat_detection(enabled: bool) {
    HEARTBEAT_DETECTION_ENABLED.store(enabled, Ordering::Release);
}

pub struct SerialPl011;

impl SerialPl011 {
    pub fn new() -> Self {
        Self
    }
}

impl MmioDevice for SerialPl011 {
    fn read(&self, offset: u64, _size: u8) -> u64 {
        match offset {
            0x000 => 0,                   // UARTDR — empty RX
            0x018 => (1 << 7) | (1 << 4), // UARTFR — TXFE=1, RXFE=1
            0xfe0 => 0x11,
            0xfe4 => 0x10,
            0xfe8 => 0x14,
            0xfec => 0x00,
            0xff0 => 0x0d,
            0xff4 => 0xf0,
            0xff8 => 0x05,
            0xffc => 0xb1,
            _ => 0,
        }
    }

    fn write(&self, offset: u64, value: u64, _size: u8) {
        if offset != 0x000 {
            return;
        }
        let b = (value & 0xff) as u8;

        if !console_log_enabled()
            && !LOG_SINK_SET.load(Ordering::Acquire)
            && !HEARTBEAT_DETECTION_ENABLED.load(Ordering::Acquire)
        {
            return;
        }

        // Route output: log sink wins if set, else stdout. Locking
        // the sink mutex per byte is cheap (single-vCPU contention is
        // none; multi-vCPU contention is rare since the kernel
        // serializes PL011 writes via a port spinlock anyway).
        {
            let mut sink = LOG_SINK.lock().unwrap();
            if let Some(f) = sink.as_mut() {
                let _ = f.write_all(&[b]);
            } else if console_log_enabled() {
                drop(sink);
                let mut out = std::io::stdout().lock();
                let _ = out.write_all(&[b]);
                let _ = out.flush();
            }
        }

        if !HEARTBEAT_DETECTION_ENABLED.load(Ordering::Acquire) {
            return;
        }

        // Line-aware marker detection (heartbeat + workload-parked).
        let mut line = PL011_LINE.lock().unwrap();
        if b == b'\n' {
            if line
                .windows(HEARTBEAT_MARKER.len())
                .any(|w| w == HEARTBEAT_MARKER)
            {
                HEARTBEAT_COUNT.fetch_add(1, Ordering::SeqCst);
            }
            if line
                .windows(PARKED_MARKER.len())
                .any(|w| w == PARKED_MARKER)
            {
                WORKLOAD_PARKED.store(true, Ordering::SeqCst);
            }
            if line
                .windows(PRE_EXEC_MARKER.len())
                .any(|w| w == PRE_EXEC_MARKER)
            {
                PRE_EXEC_READY.store(true, Ordering::SeqCst);
            }
            line.clear();
        } else if line.len() < 4096 {
            line.push(b);
        }
    }

    fn len(&self) -> u64 {
        super::super::arch::aarch64::layout::SERIAL_MMIO_SIZE
    }
}
