display-config 0.1.1

use std::{
    collections::HashMap,
    ffi::c_void,
    sync::{Arc, Mutex},
};

use dpi::{LogicalPosition, LogicalSize};
use objc2_core_graphics::{
    CGDirectDisplayID, CGDisplayBounds, CGDisplayChangeSummaryFlags, CGDisplayCopyDisplayMode,
    CGDisplayIsMain, CGDisplayMirrorsDisplay, CGDisplayMode,
    CGDisplayRegisterReconfigurationCallback, CGDisplayRemoveReconfigurationCallback, CGError,
    CGGetActiveDisplayList, kCGNullDirectDisplay,
};
use smallvec::SmallVec;

use crate::{Display, DisplayEventCallback, Event};

/// The type alias for macOS display ID, which is [`CGDirectDisplayID`][CGDirectDisplayID].
///
/// [CGDirectDisplayID]: https://developer.apple.com/documentation/coregraphics/cgdirectdisplayid?language=objc
pub type MacOSDisplayId = CGDirectDisplayID;

/// The error type for macOS-specific operations, which is [`CGError`][CGError].
///
/// [CGError]: https://developer.apple.com/documentation/coregraphics/cgerror?language=objc
pub type MacOSError = CGError;

trait CGErrorToResult {
    fn into_result<T>(self, value: T) -> Result<T, MacOSError>;
}

impl CGErrorToResult for CGError {
    fn into_result<T>(self, value: T) -> Result<T, MacOSError> {
        if self == CGError::Success {
            Ok(value)
        } else {
            Err(self)
        }
    }
}

fn get_scale_factor(id: CGDirectDisplayID) -> f64 {
    let mode = CGDisplayCopyDisplayMode(id);
    let pixel_width = CGDisplayMode::pixel_width(mode.as_deref());
    let point_width = CGDisplayMode::width(mode.as_deref());

    pixel_width as f64 / point_width as f64
}

fn get_macos_display(id: MacOSDisplayId) -> Display {
    let bounds = CGDisplayBounds(id);
    let origin = LogicalPosition::new(bounds.origin.x as i32, bounds.origin.y as i32);
    let size = LogicalSize::new(bounds.size.width as u32, bounds.size.height as u32);
    let is_primary = CGDisplayIsMain(id);
    let is_mirrored = CGDisplayMirrorsDisplay(id) != kCGNullDirectDisplay;
    let scale_factor = get_scale_factor(id);

    Display {
        id: id.into(),
        origin,
        size,
        scale_factor,
        is_primary,
        is_mirrored,
    }
}

pub(crate) fn get_macos_displays() -> Result<Vec<Display>, MacOSError> {
    const MAX_DISPLAYS: u32 = 20;
    let mut active_displays = [0; MAX_DISPLAYS as _];
    let mut display_count = 0;

    unsafe {
        CGGetActiveDisplayList(
            MAX_DISPLAYS,
            &raw mut active_displays as *mut _,
            &mut display_count,
        )
        .into_result(())?;
    }

    let mut displays = Vec::new();
    for &display_id in active_displays.iter().take(display_count as usize) {
        displays.push(get_macos_display(display_id));
    }

    Ok(displays)
}

#[derive(Default)]
struct EventTracker {
    cached_displays: HashMap<MacOSDisplayId, Display>,
}

impl EventTracker {
    fn new() -> Result<Self, MacOSError> {
        Ok(Self {
            cached_displays: Self::collect_new_cached_state()?,
        })
    }

    fn collect_new_cached_state() -> Result<HashMap<MacOSDisplayId, Display>, MacOSError> {
        let displays = get_macos_displays()?;
        let mut cached_state = HashMap::new();

        for display in displays {
            let macos_id = display.id.macos_id();
            cached_state.insert(*macos_id, display);
        }

        Ok(cached_state)
    }

    fn add(&mut self, display: Display) {
        let id = *display.id.macos_id();
        self.cached_displays.insert(id, display);
    }

    fn remove(&mut self, id: MacOSDisplayId) {
        self.cached_displays.remove(&id);
    }

    fn track_changes(&mut self) -> Result<SmallVec<[Event; 4]>, MacOSError> {
        let before =
            std::mem::replace(&mut self.cached_displays, Self::collect_new_cached_state()?);
        let mut events = SmallVec::new();

        for (id, before_display) in before.iter() {
            if let Some(after_display) = self.cached_displays.get(id) {
                if before_display.size != after_display.size {
                    events.push(Event::SizeChanged {
                        display: (*after_display).clone(),
                        before: before_display.size,
                        after: after_display.size,
                    });
                }

                if before_display.origin != after_display.origin {
                    events.push(Event::OriginChanged {
                        display: (*after_display).clone(),
                        before: before_display.origin,
                        after: after_display.origin,
                    });
                }
            }
        }

        Ok(events)
    }
}

struct UserInfo {
    callback: Option<DisplayEventCallback>,
    tracker: EventTracker,
    callback_revision: u64,
}

/// A macOS-specific display observer that monitors changes to the display configuration.
///
/// This observer uses `CGDisplayRegisterReconfigurationCallback` to receive notifications
/// about display changes. It also caches display information to track changes
/// like resolution and origin, which are not directly provided by the callback.
pub(crate) struct MacOSDisplayObserver {
    user_info: Arc<Mutex<UserInfo>>,
}

impl MacOSDisplayObserver {
    /// Creates a new `MacOSDisplayObserver`.
    ///
    /// This function sets up the necessary Core Graphics callbacks to begin observing
    /// display configuration changes.
    pub fn new() -> Result<Self, MacOSError> {
        let user_info = Arc::new(Mutex::new(UserInfo {
            callback: None,
            tracker: EventTracker::new()?,
            callback_revision: 0,
        }));

        unsafe {
            let user_info = Arc::as_ptr(&user_info) as *mut c_void;
            CGDisplayRegisterReconfigurationCallback(Some(display_callback), user_info)
                .into_result(())?;
        }

        Ok(Self { user_info })
    }

    /// Sets the callback function to be invoked when a display event occurs.
    pub fn set_callback(&self, callback: DisplayEventCallback) {
        let mut user_info = self.user_info.lock().unwrap();
        user_info.callback_revision = user_info.callback_revision.wrapping_add(1);
        user_info.callback = Some(callback);
    }

    /// Removes the currently set callback function.
    /// After calling this, no display events will be dispatched.
    pub fn remove_callback(&self) {
        let mut user_info = self.user_info.lock().unwrap();
        user_info.callback_revision = user_info.callback_revision.wrapping_add(1);
        user_info.callback = None;
    }

    /// Runs the [`NSApplication`][NSApplication] event loop to start handling display events.
    ///
    /// This function will block the current thread and dispatch events.
    ///
    /// # Panics
    /// This function must be called on the main thread, otherwise it will panic.
    ///
    /// [NSApplication]: https://developer.apple.com/documentation/appkit/nsapplication
    pub fn run(&self) {
        let mtm =
            objc2::MainThreadMarker::new().expect("This function must be called on main thread");
        objc2_app_kit::NSApplication::sharedApplication(mtm).run();
    }
}

impl Drop for MacOSDisplayObserver {
    fn drop(&mut self) {
        unsafe {
            let user_info = Arc::as_ptr(&self.user_info) as *mut c_void;
            if let Err(e) =
                CGDisplayRemoveReconfigurationCallback(Some(display_callback), user_info)
                    .into_result(())
            {
                panic!(
                    "failed to remove CGDisplay reconfiguration callback in drop: {:?}",
                    e
                );
            }
        }
    }
}

unsafe extern "C-unwind" fn display_callback(
    id: CGDirectDisplayID,
    flags: CGDisplayChangeSummaryFlags,
    user_info: *mut c_void,
) {
    if user_info.is_null() {
        return;
    }

    // We only care about the "after" events, so ignore BeginConfiguration.
    if flags.contains(CGDisplayChangeSummaryFlags::BeginConfigurationFlag) {
        return;
    }

    // We don't own the Arc here, just borrowing the pointer.
    // The `MacOSDisplayObserver` keeps the Arc alive.
    // SAFETY: `user_info` is the pointer to the `Arc<Mutex<CallbackState>>` created in new().
    let user_info = unsafe { &*(user_info as *const Mutex<UserInfo>) };
    let (events, callback_revision, mut callback) = {
        let Ok(mut user_info) = user_info.lock() else {
            return;
        };
        let Some(callback) = user_info.callback.take() else {
            return;
        };

        let mut events: SmallVec<[Event; 4]> = SmallVec::new();
        // Always get the fresh state of the display when an event happens.
        let display_snapshot = get_macos_display(id);

        if flags.contains(CGDisplayChangeSummaryFlags::AddFlag) {
            user_info.tracker.add(display_snapshot.clone());
            events.push(Event::Added(display_snapshot));
        } else if flags.contains(CGDisplayChangeSummaryFlags::RemoveFlag) {
            user_info.tracker.remove(id);
            events.push(Event::Removed(id.into()));
        } else if flags.contains(CGDisplayChangeSummaryFlags::MirrorFlag) {
            events.push(Event::Mirrored(display_snapshot));
        } else if flags.contains(CGDisplayChangeSummaryFlags::UnMirrorFlag) {
            events.push(Event::UnMirrored(display_snapshot));
        } else if (flags.contains(CGDisplayChangeSummaryFlags::SetModeFlag)
            || flags.contains(CGDisplayChangeSummaryFlags::MovedFlag))
            && let Ok(tracked_events) = user_info.tracker.track_changes()
        {
            for event in tracked_events {
                events.push(event);
            }
        }

        if events.is_empty() {
            user_info.callback = Some(callback);
            return;
        }

        (events, user_info.callback_revision, callback)
    };

    for event in events {
        (callback)(event);
    }

    // Only restore the callback when no other thread updated it while we were
    // executing events without the lock. The revision check prevents replacing
    // a newer callback (or a removal) with this older one.
    if let Ok(mut user_info) = user_info.lock()
        && user_info.callback.is_none()
        && user_info.callback_revision == callback_revision
    {
        user_info.callback = Some(callback);
    }
}