use std::collections::{HashMap, HashSet};

use drm::control::{Device as ControlDevice, connector};

/// Responsible for tracking connected/disconnected events.
///
/// ### Example
/// ```no_run
/// # mod helpers { include!("../docs/doctest_helpers.rs"); };
/// # let drm_device: helpers::FakeDevice = todo!();
/// use smithay_drm_extras::drm_scanner::{ConnectorScanner, ConnectorScanEvent};
///
/// let mut scanner = ConnectorScanner::new();
///
/// for event in scanner.scan(&drm_device).expect("failed to scan connectors") {
///     match event {
///         ConnectorScanEvent::Connected(conn) => {},
///         ConnectorScanEvent::Disconnected(conn) => {},
///         ConnectorScanEvent::Changed(conn) => {},
///     }
/// }
#[derive(Debug, Default)]
pub struct ConnectorScanner {
    connectors: HashMap<connector::Handle, connector::Info>,
}

impl ConnectorScanner {
    /// Create new [`ConnectorScanner`]
    pub fn new() -> Self {
        Default::default()
    }

    /// Should be called on every device changed event
    pub fn scan(&mut self, drm: &impl ControlDevice) -> std::io::Result<ConnectorScanResult> {
        let res_handles = drm.resource_handles()?;
        let connector_handles = res_handles.connectors();

        let mut added = Vec::new();
        let mut removed = Vec::new();
        let mut changed = Vec::new();

        for conn in connector_handles
            .iter()
            .filter_map(|conn| drm.get_connector(*conn, true).ok())
        {
            let curr_state = conn.state();

            use connector::State;
            if let Some(old) = self.connectors.insert(conn.handle(), conn.clone()) {
                match (old.state(), curr_state) {
                    (State::Connected, State::Disconnected) => removed.push(conn),
                    (State::Disconnected | State::Unknown, State::Connected) => added.push(conn),
                    //
                    // Emit Changed when the mode list changes while staying
                    // connected. This covers the EDID race where the kernel
                    // initially reports a connector as Connected with an
                    // empty or fallback mode list (EDID not yet read) and a
                    // later rescan finds the real modes populated.
                    (State::Connected, State::Connected) => {
                        if old.modes() != conn.modes() {
                            changed.push(conn);
                        }
                    }
                    (State::Disconnected, State::Disconnected) => {}
                    //
                    (State::Unknown, _) => {}
                    (_, State::Unknown) => {}
                }
            } else if curr_state == State::Connected {
                added.push(conn)
            }
        }

        // Check for connectors that disappeared entirely (e.g., USB-C dock disconnected while
        // suspended). These won't be in resource_handles() anymore, so handle them specially.
        let connector_handles_set: HashSet<_> = connector_handles.iter().copied().collect();
        let to_remove: Vec<_> = self
            .connectors
            .keys()
            .filter(|handle| !connector_handles_set.contains(handle))
            .copied()
            .collect();

        for handle in to_remove {
            if let Some(info) = self.connectors.remove(&handle) {
                if info.state() == connector::State::Connected {
                    removed.push(info);
                }
            }
        }

        Ok(ConnectorScanResult {
            connected: added,
            disconnected: removed,
            changed,
        })
    }

    /// Get map of all connectors, connected and disconnected ones.
    pub fn connectors(&self) -> &HashMap<connector::Handle, connector::Info> {
        &self.connectors
    }
}

/// Result of [`ConnectorScanner::scan`]
///
/// You can use `added` and `removed` fields of this result manually,
/// or you can just iterate (using [`IntoIterator`] or [`ConnectorScanResult::iter`])
/// over this result to get [`ConnectorScanEvent`].
#[derive(Debug, Default, Clone)]
pub struct ConnectorScanResult {
    /// Connectors that got plugged in since last scan
    pub connected: Vec<connector::Info>,
    /// Connectors that got unplugged since last scan
    pub disconnected: Vec<connector::Info>,
    /// Connectors whose mode list changed while staying connected
    pub changed: Vec<connector::Info>,
}

/// Created from [`ConnectorScanResult`], informs about connector events.
#[derive(Debug, Clone)]
pub enum ConnectorScanEvent {
    /// A new connector got plugged in since last scan
    Connected(connector::Info),
    /// A connector got unplugged in since last scan
    Disconnected(connector::Info),
    /// The connector's mode list changed while staying connected
    Changed(connector::Info),
}

impl ConnectorScanResult {
    /// Creates event iterator for this result
    ///
    /// Internally this clones the data so it is equivalent to [`IntoIterator`]
    pub fn iter(&self) -> impl Iterator<Item = ConnectorScanEvent> {
        self.clone().into_iter()
    }
}

impl IntoIterator for ConnectorScanResult {
    type Item = ConnectorScanEvent;
    type IntoIter = std::vec::IntoIter<ConnectorScanEvent>;

    fn into_iter(self) -> Self::IntoIter {
        self.disconnected
            .into_iter()
            .map(ConnectorScanEvent::Disconnected)
            .chain(self.connected.into_iter().map(ConnectorScanEvent::Connected))
            .chain(self.changed.into_iter().map(ConnectorScanEvent::Changed))
            .collect::<Vec<_>>()
            .into_iter()
    }
}
