#![cfg_attr(target_os = "yggdrasil", feature(yggdrasil_os, rustc_private))]
#![allow(clippy::single_match, clippy::collapsible_else_if)]
#![feature(if_let_guard, iter_chain, map_many_mut, let_chains)]

use std::{
    collections::{BTreeMap, HashMap, HashSet},
    marker::PhantomData,
    os::fd::{AsRawFd, RawFd},
    process::{Command, ExitCode},
};

use cross::mem::SharedMemory;
use input::InputState;
use libcolors::{
    event::{
        EventData, KeyEvent, KeyModifiers, KeyboardKeyEvent, WindowEvent, WindowInfo,
        WindowManagementEvent,
    },
    input::Key,
    message::{ClientMessage, CreateWindowInfo},
};
use sys::{Backend, DisplaySurface, FromClient, Point, ServerSender, WindowServer};
use uipc::PeerAddress;
use window::Window;
use wm::{display::Display, layout::Direction};

pub mod input;
pub mod sys;
pub mod window;
pub mod wm;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("{0}")]
    Backend(#[from] sys::Error),
}

pub struct Server<'a> {
    input_state: InputState,
    last_client_id: u32,
    client_map: HashMap<u32, PeerAddress>,
    wm_clients: HashSet<PeerAddress>,

    display: Display<u32>,
    windows: BTreeMap<u32, Window<'a>>,
    last_window_id: u32,

    background: u32,
    _pd: PhantomData<&'a ()>,
}

impl Server<'_> {
    fn new() -> Result<Self, Error> {
        Ok(Self {
            input_state: InputState::default(),

            background: 0xCCCCCC,

            last_client_id: 0,
            client_map: HashMap::new(),
            wm_clients: HashSet::new(),

            windows: BTreeMap::new(),
            display: Display::new(800, 600),
            last_window_id: 1,

            _pd: PhantomData,
        })
    }

    fn create_window(
        &mut self,
        surface: &mut DisplaySurface,
        tx: &mut ServerSender,
        peer: &PeerAddress,
        info: CreateWindowInfo,
    ) -> Result<(WindowInfo, RawFd, Option<usize>), Error> {
        let wid = self.last_window_id;
        self.last_window_id += 1;
        let (need_focus, workspace) = self
            .display
            .create_window(wid, &info.ty)
            .expect("TODO: handle create window failed");
        let frame = self.display.window_frame(wid).unwrap();
        let mapping_size = surface.width() * surface.height() * 4;
        // let mapping_size = self.display.width() * self.display.height() * 4;
        let surface_shm = SharedMemory::new(mapping_size).unwrap();
        let fd = surface_shm.as_raw_fd();
        let mut surface_mapping = surface_shm.map().unwrap();

        let surface_data = unsafe {
            std::slice::from_raw_parts_mut(
                surface_mapping.as_mut_ptr() as *mut u32,
                (frame.w * frame.h) as usize,
            )
        };

        let window = Window {
            wid,
            peer: peer.clone(),
            surface_mapping,
            surface_data,
        };

        self.windows.insert(wid, window);

        let info = WindowInfo {
            window_id: wid,
            surface_stride: surface.width() * 4, // self.display.width() * 4,
            surface_mapping_size: mapping_size,
            width: frame.w,
            height: frame.h,
        };

        surface.fill(self.background);
        if need_focus {
            self.focus_window(tx, wid);
        }
        self.redraw_all(tx);

        Ok((info, fd, workspace))
    }

    fn focus_window(&mut self, tx: &mut ServerSender, wid: u32) {
        let (old_wid, new) = self.display.current_workspace_mut().focus_window(wid);
        let old = old_wid.and_then(|wid| {
            let window = self.windows.get(&wid)?;
            Some((wid, window))
        });
        let new = new.and_then(|_| self.windows.get(&wid));

        if let Some((wid, window)) = old {
            log::info!("wid #{wid} focus=false");
            tx.send_event(
                EventData::WindowEvent(wid, WindowEvent::FocusChanged(false)),
                &window.peer,
            );
        }
        if let Some(window) = new {
            log::info!("wid #{wid} focus=true");
            tx.send_event(
                EventData::WindowEvent(wid, WindowEvent::FocusChanged(true)),
                &window.peer,
            );
        }
    }

    fn move_focus(&mut self, tx: &mut ServerSender, direction: Direction) {
        if let Some(wid) = self
            .display
            .current_workspace_mut()
            .window_towards_wrap(direction)
        {
            self.focus_window(tx, wid);
        }
    }

    fn redraw_all(&mut self, tx: &mut ServerSender) {
        self.display.all_windows().for_each(|(wid, _)| {
            if let Some(window) = self.windows.get(&wid) {
                tx.send_event(
                    EventData::WindowEvent(wid, WindowEvent::RedrawRequested),
                    &window.peer,
                );
            }
        });
    }

    fn move_window(
        &mut self,
        surface: &mut DisplaySurface,
        tx: &mut ServerSender,
        direction: Direction,
    ) {
        if self.display.move_window(direction) {
            surface.fill(self.background);
            self.flush_dirty_frames(tx);
            self.redraw_all(tx);
        }
    }

    fn flush_dirty_frames(&mut self, tx: &mut ServerSender) {
        for (wid, rect) in self.display.dirty_windows() {
            let Some(window) = self.windows.get_mut(&wid) else {
                continue;
            };
            log::info!("Resize #{}: {}x{}", wid, rect.w, rect.h);

            window.resize(rect.w, rect.h);

            tx.send_event(
                EventData::WindowEvent(
                    wid,
                    WindowEvent::Resized {
                        width: rect.w,
                        height: rect.h,
                    },
                ),
                &window.peer,
            );
        }
    }

    fn switch_to_workspace(&mut self, tx: &mut ServerSender, next: usize) -> bool {
        if next == self.display.current_workspace || next >= Display::<u32>::MAX {
            return false;
        }

        log::info!(
            "Switch workspace {} -> {}",
            self.display.current_workspace,
            next
        );

        let old_workspace = self.display.current_workspace;
        let old_focus_wid = self.display.current_workspace().focused_window();
        let old_focus = old_focus_wid.and_then(|wid| {
            let window = self.windows.get(&wid)?;
            Some((wid, window))
        });

        if let Some((old_wid, old_window)) = old_focus {
            log::info!("Unfocus #{old_wid}");
            tx.send_event(
                EventData::WindowEvent(old_wid, WindowEvent::FocusChanged(false)),
                &old_window.peer,
            );
        }

        self.display.current_workspace = next;

        let new_focus_wid = self.display.current_workspace().focused_window();
        let new_focus = new_focus_wid.and_then(|wid| {
            let window = self.windows.get(&wid)?;
            Some((wid, window))
        });

        if let Some((new_wid, new_window)) = new_focus {
            log::info!("Focus #{new_wid}");
            tx.send_event(
                EventData::WindowEvent(new_wid, WindowEvent::FocusChanged(true)),
                &new_window.peer,
            );
        }

        // Send redraw to all frames
        self.flush_dirty_frames(tx);
        self.redraw_all(tx);

        self.send_window_management_event(
            tx,
            WindowManagementEvent::WorkspaceChanged {
                old: old_workspace,
                new: next,
            },
        );

        true
    }

    fn send_focus_to_workspace(&mut self, tx: &mut ServerSender, target: usize) -> bool {
        if target == self.display.current_workspace && target >= Display::<u32>::MAX {
            return false;
        }

        if !self.display.send_focused_window(target) {
            return false;
        }

        self.flush_dirty_frames(tx);
        self.redraw_all(tx);
        self.send_workspace_activity_event(tx);
        true
    }

    fn send_window_management_event(&self, tx: &mut ServerSender, event: WindowManagementEvent) {
        for client in self.wm_clients.iter() {
            tx.send_event(EventData::WindowManagementEvent(event.clone()), client);
        }
    }

    fn send_workspace_activity_event(&self, tx: &mut ServerSender) {
        let mask = self.display.activity_mask();
        self.send_window_management_event(
            tx,
            WindowManagementEvent::WorkspaceActivity {
                mask,
                max: Display::<u32>::MAX,
            },
        );
    }
}

impl WindowServer for Server<'_> {
    fn handle_initial(&mut self, mut surface: DisplaySurface) {
        self.display
            .resize(surface.width() as u32, surface.height() as u32);
        surface.fill(self.background);
        surface.present();

        Command::new("/bin/colors-bar").spawn().ok();
    }

    fn handle_display_resize(&mut self, tx: &mut ServerSender, mut surface: DisplaySurface) {
        log::info!("Display resize: {}x{}", surface.width(), surface.height());
        self.display
            .resize(surface.width() as u32, surface.height() as u32);
        surface.fill(self.background);
        self.redraw_all(tx);
        surface.present();
    }

    fn handle_client_message(
        &mut self,
        mut surface: DisplaySurface,
        tx: &mut ServerSender,
        message: FromClient,
    ) {
        let peer = message.peer;
        match message.message {
            ClientMessage::ClientHello => {
                log::info!("{:?}: ClientHello", peer);
                // Echo the ID back
                self.last_client_id += 1;
                let id = self.last_client_id;
                self.client_map.insert(id, peer.clone());
                tx.send_event(EventData::ServerHello(id), &peer);
            }
            ClientMessage::SubscribeToWindowManagement => {
                self.wm_clients.insert(peer);
            }
            ClientMessage::CreateWindow(info) => {
                log::info!("{:?}: CreateWindow", peer);
                let (info, shm_fd, workspace) =
                    self.create_window(&mut surface, tx, &peer, info).unwrap();
                let window_id = info.window_id;
                tx.send_event_with_file(EventData::NewWindowInfo(info), &shm_fd, &peer);
                self.flush_dirty_frames(tx);
                tx.send_event(
                    EventData::WindowEvent(window_id, WindowEvent::RedrawRequested),
                    &peer,
                );
                surface.present();
                self.send_window_management_event(
                    tx,
                    WindowManagementEvent::WindowCreated {
                        workspace,
                        wid: window_id,
                    },
                );
                self.send_workspace_activity_event(tx);
            }
            ClientMessage::BlitWindow {
                window_id,
                x,
                y,
                w,
                h,
            } => {
                log::trace!("{:?}: BlitWindow", peer);
                if let Some(window) = self.windows.get(&window_id) {
                    let Some(frame) = self.display.window_frame(window_id) else {
                        return;
                    };

                    let x = x.min(frame.w);
                    let y = y.min(frame.h);
                    let w = w.min(frame.w.saturating_sub(x));
                    let h = h.min(frame.h.saturating_sub(y));

                    if w == 0 || h == 0 {
                        // Invalid rectangle, skip it
                        return;
                    }

                    let dst = Point(frame.x as _, frame.y as _);
                    let src = Point(x as _, y as _);

                    log::trace!("Blit {src:?} {w}x{h} -> {dst:?}");

                    surface.blit_buffer(
                        window.surface_data,
                        dst,
                        src,
                        w as _,
                        h as _,
                        frame.w as usize,
                    );
                }
            }
            ClientMessage::DestroyWindow(wid) => {
                log::info!("{:?}: DestroyWindow", peer);
                let window = self.windows.remove(&wid);
                let old_focus = self.display.current_workspace().focused_window();
                if window.is_some() {
                    let (removed, workspace) = self.display.remove_window(wid);

                    if removed {
                        surface.fill(self.background);
                        self.flush_dirty_frames(tx);
                        self.redraw_all(tx);

                        let focus = self.display.current_workspace().focused_window();

                        if focus != old_focus {
                            let focus = focus.and_then(|wid| {
                                let window = self.windows.get(&wid)?;
                                Some((wid, window))
                            });

                            if let Some((wid, window)) = focus {
                                tx.send_event(
                                    EventData::WindowEvent(wid, WindowEvent::FocusChanged(true)),
                                    &window.peer,
                                );
                            }
                        }

                        surface.present();

                        self.send_window_management_event(
                            tx,
                            WindowManagementEvent::WindowRemoved { workspace, wid },
                        );
                        self.send_workspace_activity_event(tx);
                    }
                }
            }
        }
    }

    fn handle_keyboard_event(
        &mut self,
        mut surface: DisplaySurface,
        tx: &mut ServerSender,
        event: KeyboardKeyEvent,
    ) {
        self.input_state.update(event.key, event.state);

        let input = self.input_state.make_input(event.key);

        if event.state {
            // Non-window keys
            #[allow(clippy::single_match)]
            match (input.modifiers, input.key) {
                (KeyModifiers::ALT, Key::Enter) => {
                    // TODO do something with spawned child
                    Command::new("/bin/term").spawn().ok();
                    return;
                }
                _ => (),
            }

            match (input.modifiers, input.key) {
                (KeyModifiers::ALT, Key::Char(c)) if c.is_ascii_digit() => {
                    let number = ((c + 9) - b'0') % 10;
                    if self.switch_to_workspace(tx, number as usize) {
                        surface.fill(self.background);
                        surface.present();
                    }
                    return;
                }
                (KeyModifiers::ALT_SHIFT, Key::Char(c)) if c.is_ascii_digit() => {
                    let number = ((c + 9) - b'0') % 10;
                    if self.send_focus_to_workspace(tx, number as usize) {
                        surface.fill(self.background);
                        surface.present();
                    }
                    return;
                }
                (KeyModifiers::ALT, Key::Char(b'l')) => {
                    self.move_focus(tx, Direction::Right);
                    return;
                }
                (KeyModifiers::ALT, Key::Char(b'h')) => {
                    self.move_focus(tx, Direction::Left);
                    return;
                }
                (KeyModifiers::ALT, Key::Char(b'j')) => {
                    self.move_focus(tx, Direction::Down);
                    return;
                }
                (KeyModifiers::ALT, Key::Char(b'k')) => {
                    self.move_focus(tx, Direction::Up);
                    return;
                }
                (KeyModifiers::ALT_SHIFT, Key::Char(b'l')) => {
                    self.move_window(&mut surface, tx, Direction::Right);
                    surface.present();
                    return;
                }
                (KeyModifiers::ALT_SHIFT, Key::Char(b'h')) => {
                    self.move_window(&mut surface, tx, Direction::Left);
                    surface.present();
                    return;
                }
                (KeyModifiers::ALT_SHIFT, Key::Char(b'j')) => {
                    self.move_window(&mut surface, tx, Direction::Down);
                    surface.present();
                    return;
                }
                (KeyModifiers::ALT_SHIFT, Key::Char(b'k')) => {
                    self.move_window(&mut surface, tx, Direction::Up);
                    surface.present();
                    return;
                }
                _ => (),
            }
        }

        let focus = self
            .display
            .current_workspace()
            .focused_window()
            .and_then(|wid| {
                let window = self.windows.get(&wid)?;
                Some((wid, window))
            });

        if let Some((wid, window)) = focus {
            if event.state {
                match (input.modifiers, input.key) {
                    (KeyModifiers::ALT_SHIFT, Key::Char(b'q')) => {
                        log::info!("Send close request -> #{wid}");
                        tx.send_event(
                            EventData::WindowEvent(wid, WindowEvent::CloseRequested),
                            &window.peer,
                        );
                        return;
                    }
                    _ => (),
                }
            }

            if event.state {
                tx.send_event(
                    EventData::WindowEvent(
                        wid,
                        WindowEvent::KeyPressed(KeyEvent {
                            key: event.key,
                            modifiers: input.modifiers,
                        }),
                    ),
                    &window.peer,
                );

                // Deliver event to the window
                tx.send_event(
                    EventData::WindowEvent(wid, WindowEvent::KeyInput(input)),
                    &window.peer,
                );
            } else {
                tx.send_event(
                    EventData::WindowEvent(
                        wid,
                        WindowEvent::KeyReleased(KeyEvent {
                            key: event.key,
                            modifiers: input.modifiers,
                        }),
                    ),
                    &window.peer,
                );
            }
        }
    }
}

fn main() -> ExitCode {
    logsink::setup_logging(false);
    log::info!("Colors starting");

    let server = Server::new().unwrap();
    let backend = Backend::new(server).unwrap();
    log::info!("Run backend");
    match backend.run() {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("{error}");
            ExitCode::FAILURE
        }
    }
}
