mechutil 0.8.11

use std::collections::HashMap;
use std::sync::Arc;
use shared_memory::Shmem;
use serde::{Deserialize, Serialize};

/// Layout configuration for a single variable (received from Server)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShmVariableConfig {
    pub name: String,      // Local name (e.g., "holding_0")
    pub gm_name: String,   // Global Memory name
    pub offset: usize,
    pub size: usize,
    #[serde(rename = "type")] // Handle "type" keyword in JSON
    pub data_type: String,
    pub direction: String,
}

/// Manages the shared memory connection and pointers
pub struct ShmContext {
    #[allow(dead_code)] // Keep shmem alive
    shmem: Shmem,
    pointers: HashMap<String, usize>, // variable_name -> offset
    /// Declared layout per variable, kept so a module can check that the type it
    /// is about to read/write through a raw pointer matches what the server
    /// actually allocated. Writing the wrong width corrupts neighbouring
    /// variables silently, so the declaration has to survive resolution.
    layout: HashMap<String, ShmVariableConfig>,
}

// SAFETY: `Shmem`'s mapping is process-wide (created with mmap), so its
// raw pointer is valid from any thread in the process. The `pointers` map
// stores offsets, not raw pointers. We expose the resolved pointers through
// `unsafe` accessors that already require the caller to manage thread safety.
unsafe impl Send for ShmContext {}
unsafe impl Sync for ShmContext {}

impl std::fmt::Debug for ShmContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ShmContext")
            .field("shm_id", &self.shmem.get_os_id())
            .field("pointer_count", &self.pointers.len())
            .finish()
    }
}

impl ShmContext {
    pub fn new(os_id: &str, configs: Vec<ShmVariableConfig>) -> Result<Self, Box<dyn std::error::Error>> {
        // Open existing SHM (created by Server)
        let shmem = shared_memory::ShmemConf::new().os_id(os_id).open()?;
        let mut pointers = HashMap::new();
        let mut layout = HashMap::new();

        for config in configs {
            // Validate bounds (optional but recommended)
            if config.offset + config.size > shmem.len() {
                // Log warning or error
                continue;
            }
            pointers.insert(config.name.clone(), config.offset);
            layout.insert(config.name.clone(), config);
        }

        Ok(Self { shmem, pointers, layout })
    }

    /// Get a raw pointer to a variable.
    /// UNSAFE: Caller must ensure thread safety and type correctness.
    pub unsafe fn get_pointer(&self, name: &str) -> Option<*mut u8> {
        self.pointers.get(name).map(|&offset| unsafe { self.shmem.as_ptr().add(offset) })
    }

    /// The layout the server declared for `name` (offset/size/type/direction).
    pub fn var(&self, name: &str) -> Option<&ShmVariableConfig> {
        self.layout.get(name)
    }

    /// The declared data type of `name` (the project's variable type, e.g.
    /// `"f64"`/`"bool"`). Check this before reading or writing through
    /// [`get_pointer`](Self::get_pointer): the pointer is untyped, so a
    /// mismatched write silently corrupts adjacent variables.
    pub fn data_type(&self, name: &str) -> Option<&str> {
        self.layout.get(name).map(|c| c.data_type.as_str())
    }

    /// The number of bytes the server allocated for `name`.
    pub fn size_of(&self, name: &str) -> Option<usize> {
        self.layout.get(name).map(|c| c.size)
    }

    /// Every variable name the server mapped for this module.
    ///
    /// The GM only distributes a variable to the domain that owns its `link`
    /// prefix, so this is exactly this module's slice — never the whole project.
    pub fn names(&self) -> Vec<String> {
        self.pointers.keys().cloned().collect()
    }

    /// Resolve a list of variable names to an `ShmMap` of pointers.
    ///
    /// Names not found in the layout are silently skipped.
    ///
    /// The returned `ShmMap` keeps an `Arc<ShmContext>` so the underlying
    /// `Shmem` mapping stays alive for the lifetime of the map. Without this,
    /// dropping or replacing the original `Arc<ShmContext>` (e.g. when a new
    /// `configure_shm` arrives) would unmap the segment and leave the map's
    /// raw pointers dangling — and the next write through one of them would
    /// segfault.
    pub fn resolve(self: &Arc<Self>, names: &[String]) -> ShmMap {
        let mut map = HashMap::new();
        for name in names {
            if let Some(ptr) = unsafe { self.get_pointer(name) } {
                map.insert(name.clone(), ShmPtr::new(ptr));
            }
        }
        ShmMap::with_context(map, Arc::clone(self))
    }
}

/// Wrapper for a shared-memory pointer, stored as `usize` for Send+Sync safety.
#[derive(Debug, Clone, Copy)]
pub struct ShmPtr(pub usize);

impl ShmPtr {
    pub fn new<T>(ptr: *mut T) -> Self {
        Self(ptr as usize)
    }

    pub fn as_ptr<T>(&self) -> *mut T {
        self.0 as *mut T
    }
}

// SAFETY: ShmPtr stores an address as usize (no raw pointer) — Send+Sync is safe.
unsafe impl Send for ShmPtr {}
unsafe impl Sync for ShmPtr {}

/// A typed map of shared-memory variable names to their resolved pointers.
///
/// `_context` keeps the underlying `Shmem` mapping alive for as long as any
/// `ShmMap` clone exists. The raw pointers in `pointers` are only valid while
/// that mapping is mapped into this process's address space, so the map
/// itself must own a reference to it.
#[derive(Debug, Clone)]
pub struct ShmMap {
    pointers: Arc<HashMap<String, ShmPtr>>,
    context: Option<Arc<ShmContext>>,
}

impl ShmMap {
    /// Build an `ShmMap` with no anchored context. Only safe when the caller
    /// guarantees that the pointers' backing memory outlives every `ShmMap`
    /// clone (e.g. pointers into a static buffer, or in unit tests).
    /// Prefer `with_context` for anything sourced from `ShmContext::resolve`.
    pub fn new(pointers: HashMap<String, ShmPtr>) -> Self {
        Self {
            pointers: Arc::new(pointers),
            context: None,
        }
    }

    /// Build an `ShmMap` that anchors its lifetime to `context`. The `Arc`
    /// keeps the `Shmem` mapping alive even if the original owner (e.g.
    /// `IpcClient::shm_context`) replaces or drops its reference.
    pub fn with_context(pointers: HashMap<String, ShmPtr>, context: Arc<ShmContext>) -> Self {
        Self {
            pointers: Arc::new(pointers),
            context: Some(context),
        }
    }

    /// Get a typed pointer to a variable.
    pub fn get<T>(&self, name: &str) -> Option<*mut T> {
        self.pointers.get(name).map(|p| p.as_ptr())
    }

    /// Get the raw address of a variable.
    pub fn get_raw(&self, name: &str) -> Option<usize> {
        self.pointers.get(name).map(|p| p.0)
    }

    /// The data type the server declared for `name`, when this map was built
    /// from an [`ShmContext`]. `None` for a context-less map, or an unknown name.
    ///
    /// [`get`](Self::get)/[`read`](Self::read)/[`write`](Self::write) are
    /// untyped: nothing checks that `T` matches the variable, and a too-wide
    /// write silently corrupts whatever the server allocated next. Check this
    /// first when the type comes from config rather than a compile-time constant.
    pub fn data_type(&self, name: &str) -> Option<&str> {
        self.context.as_ref().and_then(|c| c.data_type(name))
    }

    /// Read a value from shared memory.
    ///
    /// # Safety
    /// Caller must ensure `T` matches the type of data at the pointer.
    pub unsafe fn read<T: Copy>(&self, name: &str) -> Option<T> {
        self.get::<T>(name).map(|ptr| unsafe { *ptr })
    }

    /// Write a value to shared memory.
    ///
    /// # Safety
    /// Caller must ensure `T` matches the type of data at the pointer.
    pub unsafe fn write<T: Copy>(&self, name: &str, value: T) -> bool {
        if let Some(ptr) = self.get::<T>(name) {
            unsafe { *ptr = value };
            true
        } else {
            false
        }
    }

    /// Check if a variable name exists in the map.
    pub fn contains(&self, name: &str) -> bool {
        self.pointers.contains_key(name)
    }

    /// Return the number of resolved pointers.
    pub fn len(&self) -> usize {
        self.pointers.len()
    }

    /// Return true if no pointers were resolved.
    pub fn is_empty(&self) -> bool {
        self.pointers.is_empty()
    }

    /// Iterate over all (name, pointer) entries.
    pub fn iter(&self) -> impl Iterator<Item = (&String, &ShmPtr)> {
        self.pointers.iter()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use shared_memory::ShmemConf;

    /// Regression test for the cold-boot segfault in autocore-ethercat.
    ///
    /// Reproduces the scenario where an `IpcClient::shm_context` slot is
    /// replaced (a second `configure_shm` arrives, dropping the original
    /// `Arc<ShmContext>`) while an `ShmMap` resolved from the original
    /// context is still in use elsewhere. Before the fix, the underlying
    /// `Shmem` was `munmap`'d when the original Arc was dropped, and writes
    /// through the map's raw pointers segfaulted. With the fix, the map
    /// holds its own `Arc<ShmContext>` and the mapping stays valid.
    #[test]
    fn shmmap_survives_original_context_drop() {
        // Create a real SHM segment owned by this test.
        let owner = ShmemConf::new()
            .size(64)
            .create()
            .expect("create SHM segment");
        let os_id = owner.get_os_id().to_string();

        let configs = vec![ShmVariableConfig {
            name: "flag".to_string(),
            gm_name: "gm.flag".to_string(),
            offset: 0,
            size: 1,
            data_type: "bool".to_string(),
            direction: "read".to_string(),
        }];

        // Build a context that opens the same segment (mirrors the
        // module-side mapping).
        let ctx = Arc::new(ShmContext::new(&os_id, configs).expect("open SHM"));
        let map = ctx.resolve(&["flag".to_string()]);

        // Drop the caller's reference — this is what the IPC client does
        // when it replaces `shm_context` on a second `configure_shm`.
        drop(ctx);

        // The map's anchor must keep the mapping alive. Writing through
        // it must not segfault. Read it back via the owner's pointer to
        // confirm the byte actually landed in the segment.
        unsafe { assert!(map.write::<u8>("flag", 0xA5)); }
        let observed = unsafe { *owner.as_ptr() };
        assert_eq!(observed, 0xA5);
    }
}