mod river_protocols;

use river_protocols::zriver_seat_status_v1::ZriverSeatStatusV1;
use river_protocols::{
    zriver_output_status_v1, zriver_seat_status_v1,
    zriver_status_manager_v1::ZriverStatusManagerV1,
};
use serde::ser::{SerializeSeq, SerializeStruct, Serializer};
use serde::Serialize;
use wayland_client::protocol::wl_registry::WlRegistry;
use std::collections::HashMap;
use wayland_client::protocol::{wl_output, wl_output::WlOutput, wl_seat,
    wl_seat::WlSeat};
use wayland_client::{Attached, DispatchData, Display, GlobalEvent, GlobalManager, Main};

#[derive(Debug)]
struct Flags {
    tags: bool,
    title: bool,
    urgency: bool,
    viewstag: bool,
    layout: bool,
    mode: bool,
    seat: Option<String>,
}

impl Flags {
    fn default() -> Flags {
        Flags {
            tags: false,
            title: false,
            urgency: false,
            viewstag: false,
            layout: false,
            mode: false,
            seat: None,
        }
    }
}

#[derive(Clone)]
struct Tags(u32);

struct Env {
    flags: Flags,
    output_names: HashMap<u32, String>,
    layout: Option<String>,
    title: Option<String>,
    mode: Option<String>,
    tags: Option<HashMap<u32, Tags>>,
    urgency: Option<HashMap<u32, Tags>>,
    viewstag: Option<HashMap<u32, Vec<u32>>>,
    status_manager: Option<Main<ZriverStatusManagerV1>>,
}

impl Env {
    fn translate_output_map<V: Clone>(
        &self, input: &HashMap<u32, V>
    ) -> HashMap<String, V> {
        let mut output = HashMap::<String, V>::new();
        output.reserve(input.len());
        for (id, val) in input {
            if let Some(name) = self.output_names.get(&id) {
                output.insert(name.clone(), val.clone());
            }
        }
        output
    }
}

impl Serialize for Tags {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut seq = serializer.serialize_seq(Some(self.0.count_ones() as usize))?;
        for i in 0..32 {
            if self.0 >> i & 1 == 1 {
                seq.serialize_element(&(i + 1))?;
            }
        }
        seq.end()
    }
}

impl Serialize for Env {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where S: Serializer {
        let mut state = serializer.serialize_struct("Env", 6)?;
        if let Some(layout) = self.layout.as_ref() {
            state.serialize_field("layout", layout)?;
        } else {
            state.skip_field("layout")?;
        }
        if let Some(title) = self.title.as_ref() {
            state.serialize_field("title", title)?;
        } else {
            state.skip_field("title")?;
        }
        if let Some(mode) = self.mode.as_ref() {
            state.serialize_field("mode", mode)?;
        } else {
            state.skip_field("mode")?;
        }
        if let Some(tags) = self.tags.as_ref() {
            state.serialize_field("tags", &self.translate_output_map(tags))?;
        } else {
            state.skip_field("tags")?;
        }
        if let Some(urgency) = self.urgency.as_ref() {
            state.serialize_field("urgency", &self.translate_output_map(urgency))?;
        } else {
            state.skip_field("urgency")?;
        }
        if let Some(viewstag) = self.viewstag.as_ref() {
            state.serialize_field("viewstag", &self.translate_output_map(viewstag))?;
        } else {
            state.skip_field("viewstag")?;
        }
        state.end()
    }
}

impl Env {
    fn new() -> Env {
        let flags = configuration();
        Env {
            title: None,
            layout: None,
            mode: None,
            tags: flags.tags.then(Default::default),
            urgency: flags.urgency.then(Default::default),
            viewstag: flags.viewstag.then(Default::default),
            status_manager: None,
            flags,
            output_names: HashMap::default(),
        }
    }

    fn fmt(&self) {
        if self.title.is_some()
            || self.tags.is_some()
            || self.mode.is_some()
            || self.urgency.is_some()
            || self.viewstag.is_some()
            || self.layout.is_some()
        {
            println!("{}", serde_json::to_string(self).unwrap());
        }
    }
}

fn handle_zriver_status_manager_v1(
    status_manager: Main<ZriverStatusManagerV1>, mut env: DispatchData
) {
    if let Some(env) = env.get::<Env>() {
        env.status_manager = Some(status_manager);
    }
}

fn handle_river_seat_status(
    _status: Main<ZriverSeatStatusV1>, event: zriver_seat_status_v1::Event,
    mut env: DispatchData
) {
    match event {
        zriver_seat_status_v1::Event::FocusedView { title } => {
            if let Some(env) = env.get::<Env>() {
                if env.flags.title {
                    env.title = Some(title);
                }
            }
        }
        zriver_seat_status_v1::Event::Mode { name } => {
            if let Some(env) = env.get::<Env>() {
                if env.flags.mode {
                    env.mode = Some(name);
                }
            }
        }
        _ => {}
    }
}

fn handle_seat_event(
    seat: Main<WlSeat>, event: wl_seat::Event, mut env: DispatchData
) {
    match event {
        wl_seat::Event::Name { name } => {
            if let Some(env) = env.get::<Env>() {
                if (env.flags.title || env.flags.mode) && (env.flags.seat.is_none()
                    || name.eq(env.flags.seat.as_ref().unwrap())) {
                        if let Some(status_manager) = &env.status_manager {
                            let seat_status =
                                status_manager.get_river_seat_status(&seat);
                            seat_status.quick_assign(handle_river_seat_status);
                        }
                    }
            }
        }
        _ => {}
    }
}

fn handle_river_output_status(
    output: &Main<WlOutput>, event: zriver_output_status_v1::Event,
    mut env: DispatchData
) {
    if let Some(env) = env.get::<Env>() {
        let output_id = output.as_ref().id();
        match event {
            zriver_output_status_v1::Event::FocusedTags {
                tags: focused_tags,
            } => {
                if let Some(tags) = &mut env.tags {
                    if let Some(inner_value) = tags.get_mut(&output_id) {
                        (*inner_value) = Tags(focused_tags);
                    } else {
                        tags.insert(output_id, Tags(focused_tags));
                    }
                }
            }
            zriver_output_status_v1::Event::ViewTags {
                tags,
            } => {
                if let Some(viewstag) = &mut env.viewstag {
                    let tags: Vec<u32> = tags[0..]
                        .chunks(4)
                        .map(|s| {
                            let buf = [s[0], s[1], s[2], s[3]];
                            let tagmask =
                                u32::from_le_bytes(buf);
                            for i in 0..32 {
                                if 1 << i == tagmask {
                                    return 1 + i;
                                }
                            }
                            0
                        })
                        .collect();
                    if let Some(inner_value) = viewstag.get_mut(&output_id) {
                        (*inner_value) = tags;
                    } else {
                        viewstag.insert(output_id, tags);
                    }
                }
            }
            zriver_output_status_v1::Event::UrgentTags {
                tags,
            } => {
                if let Some(urgency) = &mut env.urgency {
                    if let Some(inner_value) = urgency.get_mut(&output_id) {
                        (*inner_value) = Tags(tags);
                    } else {
                        urgency .insert(output_id, Tags(tags));
                    }
                }
            }
            zriver_output_status_v1::Event::LayoutName {
                name,
            } => {
                env.layout = Some(name);
            }
            _ => {}
        }
    }
    
}

fn handle_output_event(
    output: Main<WlOutput>, event: wl_output::Event, mut env: DispatchData
) {
    match event {
        wl_output::Event::Geometry {
            x: _,
            y: _,
            physical_width: _,
            physical_height: _,
            subpixel: _,
            make: _,
            model: _,
            transform: _,
        } => {
            if let Some(env) = env.get::<Env>() {
                if let Some(status_manager) = &env.status_manager {
                    let output_status =
                        status_manager.get_river_output_status(&output);
                    output_status.quick_assign(move |_, event, env| {
                        handle_river_output_status(&output, event, env);
                    });
                }
            }
        }
        wl_output::Event::Name {
            name,
        } => {
            if let Some(env) = env.get::<Env>() {
                env.output_names.insert(output.as_ref().id(), name);
            }
        }
        _ => {}
    }
}

fn global_manager_callback(
    event: GlobalEvent, registry: Attached<WlRegistry>, mut env: DispatchData
) {
    let mut global_filter = wayland_client::global_filter!(
        [ZriverStatusManagerV1, 1, handle_zriver_status_manager_v1],
        [WlSeat, 7, |seat: Main<WlSeat>, _env: DispatchData| {
            seat.quick_assign(handle_seat_event);
        }],
        [WlOutput, 3, |output: Main<WlOutput>, _env: DispatchData| {
            output.quick_assign(handle_output_event);
        }]
    );

    match event {
        GlobalEvent::New { id: _, interface: _, version: _ } => {
            global_filter(event, registry, env);
        }
        GlobalEvent::Removed { id, interface: _ } => {
            if let Some(env) = env.get::<Env>() {
                env.output_names.remove(&id);
            }
        }
    }
}

fn main() {
    let mut env = Env::new();

    let display = Display::connect_to_env().unwrap();
    let mut event_queue = display.create_event_queue();
    let attached_display = (*display).clone().attach(event_queue.token());

    GlobalManager::new_with_cb(
        &attached_display,
        global_manager_callback,
    );

    loop {
        event_queue
            .dispatch(&mut env, |event, object, _| {
                panic!(
                    "[callop] Encountered an orphan event: {}@{}: {}",
                    event.interface,
                    object.as_ref().id(),
                    event.name
                );
            })
            .unwrap();
        env.fmt();
    }
}

fn configuration() -> Flags {
    let mut default = Flags::default();
    let mut args = std::env::args();

    loop {
        match args.next() {
            Some(flag) => match flag.as_str() {
                "--seat" | "-s" => default.seat = args.next(),
                "--urgency" | "-u" => default.urgency = true,
                "--title" | "-w" => default.title = true,
                "--mode" | "-m" => default.mode = true,
                "--tags" | "-t" => default.tags = true,
                "--layout" | "-l" => default.layout = true,
                "--views-tag" | "-vt" => default.viewstag = true,
                "--help" | "-h" => {
                    print!("Usage: ristate [option]\n\n");
                    print!("  --tag | -t                the focused tag\n");
                    print!("  --title | -w              the title of the focused view\n");
                    print!("  --mode | -m               the current input mode\n");
                    print!("  --urgency | -u            urgent tag\n");
                    print!("  --views-tag | -vt         the tag of all views\n");
                    print!("  --seat | -s <string>      select the seat\n");
                    print!("  --layout | -l <string>    display the name of the layout\n");
                    std::process::exit(0);
                }
                _ => {}
            },
            None => break,
        }
    }
    default
}