// SPDX-License-Identifier: BSD-3-Clause
// SPDX-FileCopyrightText: 2025 Raspberry Pi Ltd

use std::env;
use std::error::Error;
use std::fmt;
use std::os::unix::net::UnixStream;
use std::path::PathBuf;

use nix::sys::epoll::{Epoll, EpollEvent, EpollFlags, EpollTimeout};

// Epoll event token for display server socket
const DISPLAY_SOCKET_TOKEN: u64 = 1;

#[derive(Debug)]
enum DisplayError {
    NoDisplay,
    ConnectionFailed(String),
}

impl fmt::Display for DisplayError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DisplayError::NoDisplay => write!(
                f,
                "No display server found (neither WAYLAND_DISPLAY nor DISPLAY set)"
            ),
            DisplayError::ConnectionFailed(msg) => {
                write!(f, "Failed to connect to display server socket: {}", msg)
            }
        }
    }
}

impl Error for DisplayError {}

#[derive(Debug)]
enum DisplayServer {
    Wayland(PathBuf),
    X11(PathBuf),
}

fn get_wayland_socket_path() -> Option<PathBuf> {
    let display = env::var("WAYLAND_DISPLAY").ok()?;

    // Try XDG_RUNTIME_DIR first (standard location)
    if let Ok(runtime_dir) = env::var("XDG_RUNTIME_DIR") {
        let socket_path = PathBuf::from(runtime_dir).join(&display);
        if socket_path.exists() {
            return Some(socket_path);
        }
    }

    // Fallback to /tmp
    let socket_path = PathBuf::from("/tmp").join(&display);
    if socket_path.exists() {
        return Some(socket_path);
    }

    None
}

fn get_x11_socket_path() -> Option<PathBuf> {
    let display = env::var("DISPLAY").ok()?;

    // Extract display number from DISPLAY (e.g., ":0" -> "0", ":1.0" -> "1")
    let display_number = display
        .trim_start_matches(':')
        .split('.')
        .next()
        .unwrap_or("0");

    let socket_path = PathBuf::from(format!("/tmp/.X11-unix/X{}", display_number));
    if socket_path.exists() {
        return Some(socket_path);
    }

    None
}

fn detect_display_server() -> Result<DisplayServer, DisplayError> {
    // Try Wayland first
    if let Some(socket_path) = get_wayland_socket_path() {
        return Ok(DisplayServer::Wayland(socket_path));
    }

    // Try X11 next
    if let Some(socket_path) = get_x11_socket_path() {
        return Ok(DisplayServer::X11(socket_path));
    }

    Err(DisplayError::NoDisplay)
}

fn monitor_socket(socket_path: &PathBuf, server_type: &str) -> Result<(), Box<dyn Error>> {
    println!(
        "Monitoring {} socket: {}",
        server_type,
        socket_path.display()
    );

    // Connect to the display server socket
    let stream = UnixStream::connect(socket_path)
        .map_err(|e| DisplayError::ConnectionFailed(e.to_string()))?;

    println!("Connected to {}", server_type);

    // Create epoll instance
    let epoll = Epoll::new(nix::sys::epoll::EpollCreateFlags::EPOLL_CLOEXEC)?;

    // Add the socket to epoll for monitoring
    let event = EpollEvent::new(
        EpollFlags::EPOLLHUP | EpollFlags::EPOLLERR,
        DISPLAY_SOCKET_TOKEN,
    );
    epoll.add(&stream, event)?;

    println!("Waiting for {} to close...", server_type);

    // Simple event loop
    let mut events = [EpollEvent::empty(); 1];
    loop {
        match epoll.wait(&mut events, EpollTimeout::NONE) {
            Ok(n) if n > 0 => {
                // Socket closed or error occurred
                println!("{} socket closed", server_type);
                break;
            }
            Ok(_) => {
                // Timeout - continue monitoring
                continue;
            }
            Err(e) => {
                eprintln!("Error monitoring socket: {}", e);
                break;
            }
        }
    }
    Ok(())
}

fn monitor_display_server() -> Result<(), Box<dyn Error>> {
    let display_server = detect_display_server()?;

    match display_server {
        DisplayServer::Wayland(socket_path) => {
            monitor_socket(&socket_path, "Wayland compositor")?;
        }
        DisplayServer::X11(socket_path) => {
            monitor_socket(&socket_path, "X11 server")?;
        }
    }

    println!("Exiting cleanly");
    Ok(())
}

fn main() -> Result<(), Box<dyn Error>> {
    monitor_display_server()
}
