oxwm

use crate::Config;

use crate::bar::Bar;

use crate::errors::WmError;

use crate::keyboard::{self, Arg, KeyAction, handlers};

use crate::layout::GapConfig;

use crate::layout::Layout;

use crate::layout::tiling::TilingLayout;

use std::collections::HashSet;

use x11rb::cursor::Handle as CursorHandle;

use x11rb::connection::Connection;

use x11rb::protocol::Event;

use x11rb::protocol::xproto::*;

use x11rb::rust_connection::RustConnection;

pub type TagMask = u32;

pub fn tag_mask(tag: usize) -> TagMask {

    1 << tag

}

pub struct WindowManager {

    config: Config,

    connection: RustConnection,

    screen_number: usize,

    root: Window,

    screen: Screen,

    windows: Vec<Window>,

    focused_window: Option<Window>,

    layout: Box<dyn Layout>,

    window_tags: std::collections::HashMap<Window, TagMask>,

    selected_tags: TagMask,

    gaps_enabled: bool,

    fullscreen_window: Option<Window>,

    floating_windows: HashSet<Window>,

    bar: Bar,

}

type WmResult<T> = Result<T, WmError>;

impl WindowManager {

    pub fn new(config: Config) -> WmResult<Self> {

        let (connection, screen_number) = x11rb::connect(None)?;

        let root = connection.setup().roots[screen_number].root;

        let screen = connection.setup().roots[screen_number].clone();

        let normal_cursor = CursorHandle::new(

            &connection,

            screen_number,

            &x11rb::resource_manager::new_from_default(&connection)?,

        )?

        .reply()?

        .load_cursor(&connection, "left_ptr")?;

        connection

            .change_window_attributes(

                root,

                &ChangeWindowAttributesAux::new()

                    .cursor(normal_cursor)

                    .event_mask(

                        EventMask::SUBSTRUCTURE_REDIRECT

                            | EventMask::SUBSTRUCTURE_NOTIFY

                            | EventMask::PROPERTY_CHANGE

                            | EventMask::KEY_PRESS

                            | EventMask::BUTTON_PRESS,

                    ),

            )?

            .check()?;

        connection.grab_button(

            false,

            root,

            EventMask::BUTTON_PRESS | EventMask::BUTTON_RELEASE,

            GrabMode::SYNC,

            GrabMode::ASYNC,

            x11rb::NONE,

            x11rb::NONE,

            ButtonIndex::M1,

            u16::from(config.modkey).into(),

        )?;

        connection.grab_button(

            false,

            root,

            EventMask::BUTTON_PRESS | EventMask::BUTTON_RELEASE,

            GrabMode::SYNC,

            GrabMode::ASYNC,

            x11rb::NONE,

            x11rb::NONE,

            ButtonIndex::M3,

            u16::from(config.modkey).into(),

        )?;

        let bar = Bar::new(&connection, &screen, screen_number, &config)?;

        let selected_tags = Self::get_saved_selected_tags(&connection, root, config.tags.len())?;

        let gaps_enabled = config.gaps_enabled;

        let mut window_manager = Self {

            config,

            connection,

            screen_number,

            root,

            screen,

            windows: Vec::new(),

            focused_window: None,

            layout: Box::new(TilingLayout),

            window_tags: std::collections::HashMap::new(),

            selected_tags,

            gaps_enabled,

            fullscreen_window: None,

            floating_windows: HashSet::new(),

            bar,

        };

        window_manager.scan_existing_windows()?;

        window_manager.update_bar()?;

        Ok(window_manager)

    }

    fn get_saved_selected_tags(

        connection: &RustConnection,

        root: Window,

        tag_count: usize,

    ) -> WmResult<TagMask> {

        let net_current_desktop = connection

            .intern_atom(false, b"_NET_CURRENT_DESKTOP")?

            .reply()?

            .atom;

        match connection

            .get_property(false, root, net_current_desktop, AtomEnum::CARDINAL, 0, 1)?

            .reply()

        {

            Ok(prop) if prop.value.len() >= 4 => {

                let desktop = u32::from_ne_bytes([

                    prop.value[0],

                    prop.value[1],

                    prop.value[2],

                    prop.value[3],

                ]);

                if desktop < tag_count as u32 {

                    let mask = tag_mask(desktop as usize);

                    return Ok(mask);

                }

            }

            Ok(_) => {}

            Err(e) => {

                eprintln!("No _NET_CURRENT_DESKTOP property ({})", e);

            }

        }

        Ok(tag_mask(0))

    }

    fn scan_existing_windows(&mut self) -> WmResult<()> {

        let tree = self.connection.query_tree(self.root)?.reply()?;

        let net_client_info = self

            .connection

            .intern_atom(false, b"_NET_CLIENT_INFO")?

            .reply()?

            .atom;

        let wm_state_atom = self

            .connection

            .intern_atom(false, b"WM_STATE")?

            .reply()?

            .atom;

        for &window in &tree.children {

            if window == self.bar.window() {

                continue;

            }

            let Ok(attrs) = self.connection.get_window_attributes(window)?.reply() else {

                continue;

            };

            if attrs.override_redirect {

                continue;

            }

            if attrs.map_state == MapState::VIEWABLE {

                let tag = self.get_saved_tag(window, net_client_info)?;

                self.windows.push(window);

                self.window_tags.insert(window, tag);

                continue;

            }

            if attrs.map_state == MapState::UNMAPPED {

                let has_wm_state = self

                    .connection

                    .get_property(false, window, wm_state_atom, AtomEnum::ANY, 0, 2)?

                    .reply()

                    .is_ok_and(|prop| !prop.value.is_empty());

                if !has_wm_state {

                    continue;

                }

                let has_wm_class = self

                    .connection

                    .get_property(false, window, AtomEnum::WM_CLASS, AtomEnum::STRING, 0, 1024)?

                    .reply()

                    .is_ok_and(|prop| !prop.value.is_empty());

                if has_wm_class {

                    let tag = self.get_saved_tag(window, net_client_info)?;

                    self.windows.push(window);

                    self.window_tags.insert(window, tag);

                }

            }

        }

        if let Some(&first) = self.windows.first() {

            self.set_focus(Some(first))?;

        }

        self.apply_layout()?;

        Ok(())

    }

    fn get_saved_tag(&self, window: Window, net_client_info: Atom) -> WmResult<TagMask> {

        match self

            .connection

            .get_property(false, window, net_client_info, AtomEnum::CARDINAL, 0, 2)?

            .reply()

        {

            Ok(prop) if prop.value.len() >= 4 => {

                let tags = u32::from_ne_bytes([

                    prop.value[0],

                    prop.value[1],

                    prop.value[2],

                    prop.value[3],

                ]);

                if tags != 0 && tags < (1 << self.config.tags.len()) {

                    return Ok(tags);

                }

            }

            Ok(_) => {}

            Err(e) => {

                eprintln!("No _NET_CLIENT_INFO property ({})", e);

            }

        }

        Ok(self.selected_tags)

    }

    fn save_client_tag(&self, window: Window, tag: TagMask) -> WmResult<()> {

        let net_client_info = self

            .connection

            .intern_atom(false, b"_NET_CLIENT_INFO")?

            .reply()?

            .atom;

        let bytes = tag.to_ne_bytes().to_vec();

        self.connection.change_property(

            PropMode::REPLACE,

            window,

            net_client_info,

            AtomEnum::CARDINAL,

            32,

            1,

            &bytes,

        )?;

        self.connection.flush()?;

        Ok(())

    }

    fn set_wm_state(&self, window: Window, state: u32) -> WmResult<()> {

        let wm_state_atom = self

            .connection

            .intern_atom(false, b"WM_STATE")?

            .reply()?

            .atom;

        let data = [state, 0u32];

        let bytes: Vec<u8> = data.iter().flat_map(|&v| v.to_ne_bytes()).collect();

        self.connection.change_property(

            PropMode::REPLACE,

            window,

            wm_state_atom,

            wm_state_atom,

            32,

            2,

            &bytes,

        )?;

        self.connection.flush()?;

        Ok(())

    }

    pub fn run(&mut self) -> WmResult<bool> {

        println!("oxwm started on display {}", self.screen_number);

        keyboard::setup_keybinds(&self.connection, self.root, &self.config.keybindings)?;

        self.update_bar()?;

        loop {

            self.bar.update_blocks();

            if let Ok(Some(event)) = self.connection.poll_for_event() {

                if let Some(should_restart) = self.handle_event(event)? {

                    return Ok(should_restart);

                }

            }

            if self.bar.needs_redraw() {

                self.update_bar()?;

            }

            std::thread::sleep(std::time::Duration::from_millis(10));

        }

    }

    fn toggle_floating(&mut self) -> WmResult<()> {

        if let Some(focused) = self.focused_window {

            if self.floating_windows.contains(&focused) {

                self.floating_windows.remove(&focused);

                self.apply_layout()?;

            } else {

                self.floating_windows.insert(focused);

                self.connection.configure_window(

                    focused,

                    &ConfigureWindowAux::new().stack_mode(StackMode::ABOVE),

                )?;

                self.apply_layout()?;

                self.connection.flush()?;

            }

        }

        Ok(())

    }

    fn toggle_fullscreen(&mut self) -> WmResult<()> {

        if let Some(focused) = self.focused_window {

            if self.fullscreen_window == Some(focused) {

                self.fullscreen_window = None;

                self.connection.map_window(self.bar.window())?;

                self.apply_layout()?;

                self.update_focus_visuals()?;

            } else {

                self.fullscreen_window = Some(focused);

                self.connection.unmap_window(self.bar.window())?;

                let screen_width = self.screen.width_in_pixels as u32;

                let screen_height = self.screen.height_in_pixels as u32;

                self.connection.configure_window(

                    focused,

                    &ConfigureWindowAux::new()

                        .x(0)

                        .y(0)

                        .width(screen_width)

                        .height(screen_height)

                        .border_width(0)

                        .stack_mode(StackMode::ABOVE),

                )?;

                self.connection.flush()?;

            }

        }

        Ok(())

    }

    fn update_bar(&mut self) -> WmResult<()> {

        let mut occupied_tags: TagMask = 0;

        for &tags in self.window_tags.values() {

            occupied_tags |= tags;

        }

        self.bar.invalidate();

        self.bar

            .draw(&self.connection, self.selected_tags, occupied_tags)?;

        Ok(())

    }

    fn handle_key_action(&mut self, action: KeyAction, arg: &Arg) -> WmResult<()> {

        match action {

            KeyAction::Spawn => handlers::handle_spawn_action(action, arg)?,

            KeyAction::KillClient => {

                if let Some(focused) = self.focused_window {

                    match self.connection.kill_client(focused) {

                        Ok(_) => {

                            self.connection.flush()?;

                        }

                        Err(e) => {

                            eprintln!("Failed to kill window {}: {}", focused, e);

                        }

                    }

                }

            }

            KeyAction::ToggleFullScreen => {

                self.toggle_fullscreen()?;

            }

            KeyAction::ToggleFloating => {

                self.toggle_floating()?;

            }

            KeyAction::FocusStack => {

                if let Arg::Int(direction) = arg {

                    self.cycle_focus(*direction)?;

                }

            }

            KeyAction::Quit | KeyAction::Restart => {

                // Handled in handle_event

            }

            KeyAction::Recompile => {

                match std::process::Command::new("oxwm")

                    .arg("--recompile")

                    .spawn()

                {

                    Ok(_) => eprintln!("Recompiling in background"),

                    Err(e) => eprintln!("Failed to spawn recompile: {}", e),

                }

            }

            KeyAction::ViewTag => {

                if let Arg::Int(tag_index) = arg {

                    self.view_tag(*tag_index as usize)?;

                }

            }

            KeyAction::MoveToTag => {

                if let Arg::Int(tag_index) = arg {

                    self.move_to_tag(*tag_index as usize)?;

                }

            }

            KeyAction::ToggleGaps => {

                self.gaps_enabled = !self.gaps_enabled;

                self.apply_layout()?;

            }

            KeyAction::None => {}

        }

        Ok(())

    }

    fn is_window_visible(&self, window: Window) -> bool {

        if let Some(&tags) = self.window_tags.get(&window) {

            (tags & self.selected_tags) != 0

        } else {

            false

        }

    }

    fn visible_windows(&self) -> Vec<Window> {

        self.windows

            .iter()

            .filter(|&&w| self.is_window_visible(w))

            .copied()

            .collect()

    }

    fn update_window_visibility(&self) -> WmResult<()> {

        for &window in &self.windows {

            if self.is_window_visible(window) {

                self.connection.map_window(window)?;

            } else {

                self.connection.unmap_window(window)?;

            }

        }

        self.connection.flush()?;

        Ok(())

    }

    pub fn view_tag(&mut self, tag_index: usize) -> WmResult<()> {

        if tag_index >= self.config.tags.len() {

            return Ok(());

        }

        if self.fullscreen_window.is_some() {

            self.fullscreen_window = None;

            self.connection.map_window(self.bar.window())?;

        }

        self.selected_tags = tag_mask(tag_index);

        self.save_selected_tags()?;

        self.update_window_visibility()?;

        self.apply_layout()?;

        self.update_bar()?;

        let visible = self.visible_windows();

        self.set_focus(visible.first().copied())?;

        Ok(())

    }

    fn save_selected_tags(&self) -> WmResult<()> {

        let net_current_desktop = self

            .connection

            .intern_atom(false, b"_NET_CURRENT_DESKTOP")?

            .reply()?

            .atom;

        let desktop = self.selected_tags.trailing_zeros();

        let bytes = (desktop as u32).to_ne_bytes();

        self.connection.change_property(

            PropMode::REPLACE,

            self.root,

            net_current_desktop,

            AtomEnum::CARDINAL,

            32,

            1,

            &bytes,

        )?;

        self.connection.flush()?;

        Ok(())

    }

    pub fn move_to_tag(&mut self, tag_index: usize) -> WmResult<()> {

        if tag_index >= self.config.tags.len() {

            return Ok(());

        }

        if let Some(focused) = self.focused_window {

            let mask = tag_mask(tag_index);

            self.window_tags.insert(focused, mask);

            let _ = self.save_client_tag(focused, mask);

            self.update_window_visibility()?;

            self.apply_layout()?;

            self.update_bar()?;

        }

        Ok(())

    }

    pub fn cycle_focus(&mut self, direction: i32) -> WmResult<()> {

        let visible = self.visible_windows();

        if visible.is_empty() {

            return Ok(());

        }

        let next_window = if let Some(current) = self.focused_window {

            if let Some(current_index) = visible.iter().position(|&w| w == current) {

                let next_index = if direction > 0 {

                    (current_index + 1) % visible.len()

                } else {

                    (current_index + visible.len() - 1) % visible.len()

                };

                visible[next_index]

            } else {

                visible[0]

            }

        } else {

            visible[0]

        };

        self.set_focus(Some(next_window))?;

        Ok(())

    }

    pub fn set_focus(&mut self, window: Option<Window>) -> WmResult<()> {

        self.focused_window = window;

        if let Some(win) = window {

            self.connection

                .set_input_focus(InputFocus::POINTER_ROOT, win, x11rb::CURRENT_TIME)?;

            self.connection.flush()?;

        }

        self.update_focus_visuals()?;

        Ok(())

    }

    fn update_focus_visuals(&self) -> WmResult<()> {

        for &window in &self.windows {

            let (border_color, border_width) = if self.focused_window == Some(window) {

                (self.config.border_focused, self.config.border_width)

            } else {

                (self.config.border_unfocused, self.config.border_width)

            };

            self.connection.configure_window(

                window,

                &ConfigureWindowAux::new().border_width(border_width),

            )?;

            self.connection.change_window_attributes(

                window,

                &ChangeWindowAttributesAux::new().border_pixel(border_color),

            )?;

        }

        self.connection.flush()?;

        Ok(())

    }

    fn move_mouse(&mut self, window: Window) -> WmResult<()> {

        self.floating_windows.insert(window);

        let geometry = self.connection.get_geometry(window)?.reply()?;

        self.connection

            .grab_pointer(

                false,

                self.root,

                (EventMask::POINTER_MOTION | EventMask::BUTTON_RELEASE).into(),

                GrabMode::ASYNC,

                GrabMode::ASYNC,

                x11rb::NONE,

                x11rb::NONE,

                x11rb::CURRENT_TIME,

            )?

            .reply()?;

        let pointer = self.connection.query_pointer(self.root)?.reply()?;

        let (start_x, start_y) = (pointer.root_x, pointer.root_y);

        self.connection.configure_window(

            window,

            &ConfigureWindowAux::new().stack_mode(StackMode::ABOVE),

        )?;

        loop {

            let event = self.connection.wait_for_event()?;

            match event {

                Event::MotionNotify(e) => {

                    let new_x = geometry.x + (e.root_x - start_x);

                    let new_y = geometry.y + (e.root_y - start_y);

                    self.connection.configure_window(

                        window,

                        &ConfigureWindowAux::new().x(new_x as i32).y(new_y as i32),

                    )?;

                    self.connection.flush()?;

                }

                Event::ButtonRelease(_) => break,

                _ => {}

            }

        }

        self.connection

            .ungrab_pointer(x11rb::CURRENT_TIME)?

            .check()?;

        self.connection

            .allow_events(Allow::REPLAY_POINTER, x11rb::CURRENT_TIME)?

            .check()?;

        Ok(())

    }

    fn resize_mouse(&mut self, window: Window) -> WmResult<()> {

        self.floating_windows.insert(window);

        let geometry = self.connection.get_geometry(window)?.reply()?;

        self.connection.warp_pointer(

            x11rb::NONE,

            window,

            0,

            0,

            0,

            0,

            geometry.width as i16,

            geometry.height as i16,

        )?;

        self.connection

            .grab_pointer(

                false,

                self.root,

                (EventMask::POINTER_MOTION | EventMask::BUTTON_RELEASE).into(),

                GrabMode::ASYNC,

                GrabMode::ASYNC,

                x11rb::NONE,

                x11rb::NONE,

                x11rb::CURRENT_TIME,

            )?

            .reply()?;

        self.connection.configure_window(

            window,

            &ConfigureWindowAux::new().stack_mode(StackMode::ABOVE),

        )?;

        loop {

            let event = self.connection.wait_for_event()?;

            match event {

                Event::MotionNotify(e) => {

                    let new_width = (e.root_x - geometry.x).max(1) as u32;

                    let new_height = (e.root_y - geometry.y).max(1) as u32;

                    self.connection.configure_window(

                        window,

                        &ConfigureWindowAux::new()

                            .width(new_width)

                            .height(new_height),

                    )?;

                    self.connection.flush()?;

                }

                Event::ButtonRelease(_) => break,

                _ => {}

            }

        }

        self.connection

            .ungrab_pointer(x11rb::CURRENT_TIME)?

            .check()?;

        self.connection

            .allow_events(Allow::REPLAY_POINTER, x11rb::CURRENT_TIME)?

            .check()?;

        Ok(())

    }

    fn handle_event(&mut self, event: Event) -> WmResult<Option<bool>> {

        match event {

            Event::MapRequest(event) => {

                let attrs = match self.connection.get_window_attributes(event.window)?.reply() {

                    Ok(attrs) => attrs,

                    Err(_) => return Ok(None),

                };

                if attrs.override_redirect {

                    return Ok(None);

                }

                if self.windows.contains(&event.window) {

                    return Ok(None);

                }

                self.connection.map_window(event.window)?;

                self.connection.change_window_attributes(

                    event.window,

                    &ChangeWindowAttributesAux::new().event_mask(EventMask::ENTER_WINDOW),

                )?;

                self.windows.push(event.window);

                self.window_tags.insert(event.window, self.selected_tags);

                self.set_wm_state(event.window, 1)?;

                let _ = self.save_client_tag(event.window, self.selected_tags);

                self.apply_layout()?;

                self.update_bar()?;

                self.set_focus(Some(event.window))?;

            }

            Event::UnmapNotify(event) => {

                if self.windows.contains(&event.window) && self.is_window_visible(event.window) {

                    self.remove_window(event.window)?;

                }

            }

            Event::DestroyNotify(event) => {

                if self.windows.contains(&event.window) {

                    self.remove_window(event.window)?;

                }

            }

            Event::EnterNotify(event) => {

                if event.mode != x11rb::protocol::xproto::NotifyMode::NORMAL {

                    return Ok(None);

                }

                if self.windows.contains(&event.event) {

                    self.set_focus(Some(event.event))?;

                }

            }

            Event::KeyPress(event) => {

                let (action, arg) = keyboard::handle_key_press(event, &self.config.keybindings);

                match action {

                    KeyAction::Quit => return Ok(Some(false)),

                    KeyAction::Restart => return Ok(Some(true)),

                    _ => self.handle_key_action(action, &arg)?,

                }

            }

            Event::ButtonPress(event) => {

                if event.event == self.bar.window() {

                    if let Some(tag_index) = self.bar.handle_click(event.event_x) {

                        self.view_tag(tag_index)?;

                    }

                } else if event.child != x11rb::NONE {

                    self.set_focus(Some(event.child))?;

                    if event.detail == ButtonIndex::M1.into() {

                        self.move_mouse(event.child)?;

                    } else if event.detail == ButtonIndex::M3.into() {

                        self.resize_mouse(event.child)?;

                    }

                }

            }

            Event::Expose(event) => {

                if event.window == self.bar.window() {

                    self.bar.invalidate();

                    self.update_bar()?;

                }

            }

            _ => {}

        }

        Ok(None)

    }

    fn apply_layout(&self) -> WmResult<()> {

        if self.fullscreen_window.is_some() {

            return Ok(());

        }

        let screen_width = self.screen.width_in_pixels as u32;

        let screen_height = self.screen.height_in_pixels as u32;

        let border_width = self.config.border_width;

        let bar_height = self.bar.height() as u32;

        let usable_height = screen_height.saturating_sub(bar_height);

        let gaps = if self.gaps_enabled {

            GapConfig {

                inner_horizontal: self.config.gap_inner_horizontal,

                inner_vertical: self.config.gap_inner_vertical,

                outer_horizontal: self.config.gap_outer_horizontal,

                outer_vertical: self.config.gap_outer_vertical,

            }

        } else {

            GapConfig {

                inner_horizontal: 0,

                inner_vertical: 0,

                outer_horizontal: 0,

                outer_vertical: 0,

            }

        };

        let visible: Vec<Window> = self

            .visible_windows()

            .into_iter()

            .filter(|w| !self.floating_windows.contains(w))

            .collect();

        let geometries = self

            .layout

            .arrange(&visible, screen_width, usable_height, &gaps);

        for (window, geometry) in visible.iter().zip(geometries.iter()) {

            let adjusted_width = geometry.width.saturating_sub(2 * border_width);

            let adjusted_height = geometry.height.saturating_sub(2 * border_width);

            let adjusted_y = geometry.y_coordinate + bar_height as i32;

            self.connection.configure_window(

                *window,

                &ConfigureWindowAux::new()

                    .x(geometry.x_coordinate)

                    .y(adjusted_y)

                    .width(adjusted_width)

                    .height(adjusted_height),

            )?;

        }

        self.connection.flush()?;

        Ok(())

    }

    fn remove_window(&mut self, window: Window) -> WmResult<()> {

        let initial_count = self.windows.len();

        self.windows.retain(|&w| w != window);

        self.window_tags.remove(&window);

        self.floating_windows.remove(&window);

        if self.fullscreen_window == Some(window) {

            self.fullscreen_window = None;

            self.connection.map_window(self.bar.window())?;

        }

        if self.windows.len() < initial_count {

            if self.focused_window == Some(window) {

                let new_focus = if self.windows.is_empty() {

                    None

                } else {

                    Some(self.windows[self.windows.len() - 1])

                };

                self.set_focus(new_focus)?;

            }

            self.apply_layout()?;

            self.update_bar()?;

        }

        Ok(())

    }

}