use std::error::Error;

use clap::Parser;
use serde::Serialize;
use xcb::{XidNew, x};

#[derive(Debug, Parser)]
#[command(name = "xdi", about = "Print X11 desktop information as JSON")]
struct Args {
    /// Print a new JSON object whenever the desktop state changes.
    #[arg(short, long)]
    follow: bool,
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
struct State {
    active_window: Option<WindowInfo>,
    workspaces: Vec<Workspace>,
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
struct Workspace {
    id: u32,
    current: bool,
    windows: Vec<WindowInfo>,
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
struct WindowInfo {
    id: u32,
    name: Option<String>,
    class: Option<String>,
}

xcb::atoms_struct! {
    #[derive(Debug, Copy, Clone)]
    struct Atoms {
        net_number_of_desktops => b"_NET_NUMBER_OF_DESKTOPS",
        net_current_desktop => b"_NET_CURRENT_DESKTOP",
        net_client_list => b"_NET_CLIENT_LIST",
        net_wm_desktop => b"_NET_WM_DESKTOP",
        net_active_window => b"_NET_ACTIVE_WINDOW",
        net_wm_name => b"_NET_WM_NAME",
        utf8_string => b"UTF8_STRING",
        wm_class => b"WM_CLASS",
    }
}

struct X11 {
    conn: xcb::Connection,
    root: x::Window,
    atoms: Atoms,
}

impl X11 {
    fn connect() -> Result<Self, Box<dyn Error>> {
        let (conn, screen_num) = xcb::Connection::connect(None)?;

        #[allow(clippy::cast_sign_loss)]
        let screen = conn
            .get_setup()
            .roots()
            .nth(screen_num as usize)
            .ok_or("X11 screen not found")?;

        let atoms = Atoms::intern_all(&conn)?;
        let root = screen.root();

        Ok(Self { conn, root, atoms })
    }

    fn state(&self) -> Result<State, Box<dyn Error>> {
        let workspace_count = self
            .property_u32(self.root, self.atoms.net_number_of_desktops)?
            .first()
            .copied()
            .unwrap_or(0);

        let current = self
            .property_u32(self.root, self.atoms.net_current_desktop)?
            .first()
            .copied()
            .unwrap_or(0);

        let active_id = self
            .property_u32(self.root, self.atoms.net_active_window)?
            .first()
            .copied();

        let client_ids = self.property_u32(self.root, self.atoms.net_client_list)?;

        let mut workspaces: Vec<Workspace> = (0..workspace_count)
            .map(|id| Workspace {
                id,
                current: id == current,
                windows: Vec::new(),
            })
            .collect();

        let mut active_window = None;

        for id in client_ids {
            let window = self.window_info(id)?;

            let desktop = self
                .property_u32(x::Window::new(id), self.atoms.net_wm_desktop)?
                .first()
                .copied();

            let Some(desktop) = desktop else {
                continue;
            };

            // EWMH uses 0xFFFFFFFF for windows that appear on
            // every desktop. Do not assign them to a specific workspace.
            if desktop == u32::MAX {
                continue;
            }

            let Some(workspace) = workspaces.get_mut(desktop as usize) else {
                continue;
            };

            if active_id == Some(id) {
                active_window = Some(window.clone());
            }

            workspace.windows.push(window);
        }

        // Keep the active window available even if its desktop could
        // not be resolved.
        if active_window.is_none()
            && let Some(id) = active_id
        {
            active_window = Some(self.window_info(id)?);
        }

        Ok(State {
            active_window,
            workspaces,
        })
    }

    fn window_info(&self, id: u32) -> Result<WindowInfo, Box<dyn Error>> {
        let window = x::Window::new(id);

        let name = self
            .property_bytes(window, self.atoms.net_wm_name, self.atoms.utf8_string)?
            .and_then(|bytes| String::from_utf8(bytes).ok());

        let class = self
            .property_bytes(window, self.atoms.wm_class, x::ATOM_STRING)?
            .map(|bytes| parse_wm_class(&bytes));

        Ok(WindowInfo { id, name, class })
    }

    fn property_u32(
        &self,
        window: x::Window,
        property: x::Atom,
    ) -> Result<Vec<u32>, Box<dyn Error>> {
        let cookie = self.conn.send_request(&x::GetProperty {
            delete: false,
            window,
            property,
            r#type: x::ATOM_ANY,
            long_offset: 0,
            long_length: u32::MAX,
        });

        let reply = self.conn.wait_for_reply(cookie)?;

        Ok(reply.value::<u32>().to_vec())
    }

    fn property_bytes(
        &self,
        window: x::Window,
        property: x::Atom,
        property_type: x::Atom,
    ) -> Result<Option<Vec<u8>>, Box<dyn Error>> {
        let cookie = self.conn.send_request(&x::GetProperty {
            delete: false,
            window,
            property,
            r#type: property_type,
            long_offset: 0,
            long_length: u32::MAX,
        });

        let reply = self.conn.wait_for_reply(cookie)?;

        let value = reply.value::<u8>();

        if value.is_empty() {
            return Ok(None);
        }

        Ok(Some(value.to_vec()))
    }

    fn watch(&self) -> Result<(), Box<dyn Error>> {
        self.select_events(
            self.root,
            x::EventMask::PROPERTY_CHANGE | x::EventMask::SUBSTRUCTURE_NOTIFY,
        )?;

        let client_ids = self.property_u32(self.root, self.atoms.net_client_list)?;

        for id in client_ids {
            self.select_events(x::Window::new(id), x::EventMask::PROPERTY_CHANGE)?;
        }

        self.conn.flush()?;

        Ok(())
    }

    fn select_events(
        &self,
        window: x::Window,
        event_mask: x::EventMask,
    ) -> Result<(), Box<dyn Error>> {
        let cookie = self.conn.send_request_checked(&x::ChangeWindowAttributes {
            window,
            value_list: &[x::Cw::EventMask(event_mask)],
        });

        self.conn.check_request(cookie)?;

        Ok(())
    }

    fn wait_for_relevant_event(&self) -> Result<(), Box<dyn Error>> {
        loop {
            let event = self.conn.wait_for_event()?;

            match event {
                xcb::Event::X(x::Event::PropertyNotify(event)) => {
                    if event.window() == self.root || self.is_relevant_property(event.atom()) {
                        return Ok(());
                    }
                }

                xcb::Event::X(x::Event::CreateNotify(event)) => {
                    let _ = self.select_events(event.window(), x::EventMask::PROPERTY_CHANGE);

                    self.conn.flush()?;

                    return Ok(());
                }

                xcb::Event::X(
                    x::Event::DestroyNotify(_) | x::Event::MapNotify(_) | x::Event::UnmapNotify(_),
                ) => {
                    return Ok(());
                }

                _ => {}
            }
        }
    }

    fn is_relevant_property(&self, atom: x::Atom) -> bool {
        atom == self.atoms.net_number_of_desktops
            || atom == self.atoms.net_current_desktop
            || atom == self.atoms.net_client_list
            || atom == self.atoms.net_wm_desktop
            || atom == self.atoms.net_active_window
            || atom == self.atoms.net_wm_name
            || atom == self.atoms.wm_class
    }
}

fn parse_wm_class(bytes: &[u8]) -> String {
    let mut parts = bytes.split(|byte| *byte == 0);

    let _instance = parts.next();
    let class = parts.next().unwrap_or_default();

    String::from_utf8_lossy(class).into_owned()
}

fn print_state(state: &State) -> Result<(), Box<dyn Error>> {
    println!("{}", serde_json::to_string(state)?);
    Ok(())
}

fn main() -> Result<(), Box<dyn Error>> {
    let args = Args::parse();
    let x11 = X11::connect()?;

    if !args.follow {
        print_state(&x11.state()?)?;
        return Ok(());
    }

    x11.watch()?;

    let mut previous = None;

    loop {
        let state = x11.state()?;

        if previous.as_ref() != Some(&state) {
            print_state(&state)?;
            previous = Some(state);
        }

        x11.wait_for_relevant_event()?;
    }
}