path: root/rust/src/main.rs

use std::fs::{File, OpenOptions};
use std::collections::{HashMap, HashSet};
use std::io::{self, Write};
use std::sync::{Arc, Mutex};
use std::convert::TryFrom;
use argparse::{ArgumentParser, Store, StoreOption};
use futures::channel::mpsc;
use futures::stream::StreamExt;
use once_cell::sync::Lazy;
use tokio::{runtime, task};
use tokio::sync::Mutex as AsyncMutex;
use termion::{clear, color, cursor, style};
use termion::event::{Event, Key, MouseEvent, MouseButton};
use termion::raw::{IntoRawMode, RawTerminal};
use termion::input::{TermRead, MouseTerminal};
use termion::screen::AlternateScreen;
use tomsg_rs::{connection, Command, Connection, Id, Line, Message, PushMessage, Reply, Word};
use unicode_width::UnicodeWidthChar;
use crate::editor::Editor;
use crate::error::IntoIOError;

mod auth;
mod bel;
mod editor;
mod error;
mod getpass;

static DEBUG_FILE: Lazy<Mutex<Option<File>>> = Lazy::new(|| {
    Mutex::new(
        OpenOptions::new().write(true).truncate(true).create(true)
            .open("debug.txt").ok())
});

fn debug_write(s: &str) {
    let file = DEBUG_FILE.lock().unwrap();
    file.as_ref()
        .map(|mut file| writeln!(file, "{}", s).unwrap());
}

macro_rules! debug {
    ($($args:expr),*) => {
        debug_write(&format!($($args),*));
    }
}

fn draw_hori_bar<W: Write>(mut w: W, y: u16, x1: u16, x2: u16) -> io::Result<()> {
    for x in x1..x2 {
        write!(w, "{}─", cursor::Goto(x + 1, y + 1))?;
    }
    Ok(())
}

fn draw_vert_bar<W: Write>(mut w: W, x: u16, y1: u16, y2: u16) -> io::Result<()> {
    for y in y1..y2 {
        write!(w, "{}│", cursor::Goto(x + 1, y + 1))?;
    }
    Ok(())
}

fn wrap_text(s: &str, maxwid: u16) -> Vec<String> {
    let mut iter = s.chars();
    let mut lines = Vec::new();
    let mut nextline = String::new();
    let mut cursor = 0;
    loop {
        let c = match iter.next() {
            Some(c) => c,
            None => break,
        };
        if let Some(wid) = c.width().map(|w| u16::try_from(w).unwrap()) {
            if cursor + wid > maxwid {
                lines.push(nextline);
                nextline = String::new();
                cursor = 0;
            }
            nextline.push(c);
            cursor += wid;
        }
    }
    if nextline.len() > 0 {
        lines.push(nextline);
    }
    lines
}

fn contains_mention(msg: &Line, name: &Word) -> bool {
    let is_word_bound = |c: Option<char>| match c {
        Some(c) => !c.is_alphanumeric(),
        None => true,
    };

    if let Some(index) = msg.find(name.as_str()) {
        let index2 = index + name.len();
        is_word_bound(*&msg[..index].chars().next_back()) &&
            is_word_bound(*&msg[index2..].chars().next_back())
    } else {
        false
    }
}

async fn send_command(conn: &Connection, cmd: Command) -> io::Result<Reply> {
    match conn.send_command(cmd).await {
        Ok(Ok(reply)) => Ok(reply),
        Ok(Err(connection::CloseReason::EOF)) => Err("EOF on server connection".ioerr()),
        Ok(Err(connection::CloseReason::Err(err))) => Err(format!("Error from server: {}", err).ioerr()),
        Err(err) => Err(err),
    }
}

static STATUS_ROOM: Lazy<Word> = Lazy::new(|| {
    Word::try_from("*status".to_string()).unwrap()
});

type Stdout = MouseTerminal<AlternateScreen<RawTerminal<io::Stdout>>>;
// type Stdout = MouseTerminal<io::Stdout>;

#[derive(Debug, Copy, Clone)]
struct Layout {
    rlsepx: u16,  // x-coordinate of room list / chat log separator bar
    nlsepx: u16,  // x-coordinate of chat log / nick list separator bar
}

impl Default for Layout {
    fn default() -> Self {
        Self { rlsepx: 5, nlsepx: 5 }
    }
}

impl Layout {
    fn compute(size: (u16, u16), state: &State) -> Self {
        let rlsepx = u16::try_from(state.roomlist.iter().map(|s| s.len()).max().unwrap_or(5)).unwrap();
        let room = state.rooms.get(&state.currentroom);
        let maxnicklen = room.map(|r| r.members.iter().map(|s| s.len()).max().unwrap_or(5)).unwrap_or(5);
        let nlsepx = size.0 - u16::min(size.0, u16::try_from(maxnicklen).unwrap() + 1);
        Self { rlsepx, nlsepx }
    }
}

struct App {
    conn: Connection,
    username: Word,
    state: AsyncMutex<State>,
}

struct State {
    stdout: Stdout,
    termsize: (u16, u16),
    layout: Layout,
    roomlist: Vec<Word>,
    rooms: HashMap<Word, RoomData>,
    currentroom: Word,
    msgs: HashMap<Id, Message>,
}

#[derive(Debug, Default)]
struct RoomData {
    members: Vec<Word>,
    history: Vec<HItem>,
    maxnicklen: usize,
    editor: Editor,
    highlight: bool,
}

#[derive(Debug)]
enum HItem {
    Message(Id, bool),
    Service(String, bool),
}

impl App {
    fn new(conn: Connection, stdout: Stdout, username: Word) -> Self {
        let termsize = (0, 0);

        let mut init_rooms_map = HashMap::new();
        init_rooms_map.insert(STATUS_ROOM.clone(), RoomData::default());

        let state = State {
            stdout,
            termsize,
            layout: Layout::default(),
            roomlist: vec![STATUS_ROOM.clone()],
            rooms: init_rooms_map,
            currentroom: STATUS_ROOM.clone(),
            msgs: HashMap::new(),
        };
        App {
            conn,
            username,
            state: AsyncMutex::new(state)
        }
    }

    async fn fetch_rooms_list(self: &Arc<Self>) -> io::Result<()> {
        debug!("fetch_rooms_list");
        if let Reply::List(rooms) = send_command(&self.conn, Command::ListRooms).await? {
            debug!("fetched room list: {:?}", rooms);
            let newrooms = rooms.iter().collect::<HashSet<_>>();

            {
                let mut state = self.state.lock().await;
                let toremove =
                        state.roomlist.iter()
                                      .filter(|&r| *r != *STATUS_ROOM && !newrooms.contains(r))
                                      .cloned()
                                      .collect::<HashSet<_>>();
                state.roomlist.retain(|r| !toremove.contains(r));
                state.rooms.retain(|r, _| !toremove.contains(r));
            }

            self.full_redraw().await?;

            for room in &rooms {
                self.enter_room(room).await?;
            }
        }
        Ok(())
    }

    async fn enter_room(self: &Arc<Self>, room: &Word) -> io::Result<()> {
        debug!("enter_room({})", room);

        {
            let mut state = self.state.lock().await;
            if state.rooms.insert(room.clone(), RoomData::default()).is_none() {
                state.roomlist.push(room.clone());

                if state.currentroom == *STATUS_ROOM {
                    state.currentroom = room.clone();
                }
            }
        }

        self.full_redraw().await?;

        let app = self.clone();
        let room2 = room.clone();
        task::spawn(async move {
            if let Ok(Reply::List(members)) =
                    send_command(&app.conn, Command::ListMembers { roomname: room2.clone() })
                        .await {
                app.put_members(&room2, members).await.unwrap();
            }
        });
        let app = self.clone();
        let room2 = room.clone();
        task::spawn(async move {
            if let Ok(Reply::History(hist)) =
                    send_command(&app.conn, Command::History { roomname: room2.clone(), count: 20 })
                        .await {
                app.put_history(&room2, hist).await.unwrap();
            }
        });

        Ok(())
    }

    async fn leave_room(self: &Arc<Self>, room: &Word) -> io::Result<()> {
        debug!("leave_room({})", room);

        {
            let mut state = self.state.lock().await;
            if state.currentroom == *room {
                state.currentroom = STATUS_ROOM.clone();
            }
            state.roomlist.retain(|r| r != room);
            state.rooms.remove(room);
        }

        self.full_redraw().await
    }

    async fn put_members(self: &Arc<Self>, room: &Word, ms: Vec<Word>) -> io::Result<()> {
        debug!("put_members {} #{}", room, ms.len());
        let mut state = self.state.lock().await;
        state.rooms.get_mut(room).unwrap().members = ms;
        state.full_redraw()
    }

    async fn add_member(self: &Arc<Self>, room: &Word, user: Word) -> io::Result<()> {
        debug!("add_member {} {}", room, user);
        let mut state = self.state.lock().await;
        state.rooms.get_mut(room).unwrap().members.push(user);
        state.full_redraw()
    }

    async fn remove_member(self: &Arc<Self>, room: &Word, user: Word) -> io::Result<()> {
        debug!("remove_member {} {}", room, user);
        let mut state = self.state.lock().await;
        state.rooms.get_mut(room).unwrap().members.retain(|u| *u != user);
        state.full_redraw()
    }

    async fn put_history(self: &Arc<Self>, room: &Word, hist: Vec<Message>) -> io::Result<()> {
        debug!("put_history {} #{}", room, hist.len());
        let mut state = self.state.lock().await;
        let ids = hist.iter().map(|m| m.id).collect::<Vec<_>>();
        let replies = hist.iter().filter_map(|m| m.reply_on).collect::<Vec<_>>();

        for msg in hist {
            state.put_message(msg)?;
        }

        for replyid in replies {
            if state.msgs.get(&replyid).is_none() {
                let app = self.clone();
                task::spawn(async move {
                    if let Ok(Reply::Message(msg2)) =
                            send_command(&app.conn, Command::GetMessage(replyid)).await {
                        app.state.lock().await.put_message(msg2).unwrap();
                    }
                });
            }
        }

        for id in ids {
            state.append_history(room, HItem::Message(id, false))?;
        }

        Ok(())
    }

    async fn append_history_item(self: &Arc<Self>, room: &Word, item: HItem) -> io::Result<()> {
        let mut state = self.state.lock().await;
        state.append_history(room, item)
    }

    async fn add_message(self: &Arc<Self>, msg: Message, highlight: bool) -> io::Result<()> {
        debug!("add_message {{id={} room={} user={}}}", msg.id, msg.roomname, msg.username);
        let mut state = self.state.lock().await;

        let roomname = msg.roomname.clone();
        let msgid = msg.id;
        let replyid = msg.reply_on;
        state.put_message(msg)?;

        if let Some(replyid) = replyid {
            if state.msgs.get(&replyid).is_none() {
                let app = self.clone();
                task::spawn(async move {
                    if let Ok(Reply::Message(msg2)) =
                            send_command(&app.conn, Command::GetMessage(replyid)).await {
                        app.state.lock().await.put_message(msg2).unwrap();
                    }
                });
            }
        }

        state.append_history(&roomname, HItem::Message(msgid, highlight))?;

        Ok(())
    }

    async fn on_push(self: &Arc<Self>, pm: PushMessage) {
        match pm {
            PushMessage::Invite { roomname: room, inviter: _ } => {
                self.enter_room(&room).await.unwrap();
            }
            PushMessage::Join { roomname: room, username } => {
                self.add_member(&room, username.clone()).await.unwrap();
                self.append_history_item(&room, HItem::Service(format!("Join: <{}>", username), false)).await.unwrap();
            }
            PushMessage::Leave { roomname: room, username: user } => {
                if user == self.username {
                    self.leave_room(&room).await.unwrap();
                } else {
                    self.remove_member(&room, user.clone()).await.unwrap();
                    self.append_history_item(&room, HItem::Service(format!("Leave: <{}>", user), false)).await.unwrap();
                }
            }
            PushMessage::Message(msg) => {
                let mention = contains_mention(&msg.message, &self.username);
                if self.decide_push_notify(&msg.roomname, mention).await {
                    bel::bel();
                    self.add_room_highlight(&msg.roomname).await.unwrap();
                }
                self.add_message(msg, mention).await.unwrap();
            }
            PushMessage::Online{ sessions: _, username: _ } => {}
        }
    }

    async fn switch_buffer_rel(self: &Arc<Self>, delta: isize) -> io::Result<()> {
        let mut state = self.state.lock().await;
        let current_index =
                state.roomlist.iter().enumerate()
                              .filter(|&(_, r)| *r == state.currentroom)
                              .map(|(i, _)| i)
                              .next().unwrap_or(0);
        let index = (current_index as isize + delta)
                        .max(0).min(state.roomlist.len() as isize - 1);
        state.switch_buffer(index as usize)
    }

    // Returns whether application should quit
    async fn on_key(self: &Arc<Self>, key: Key) -> io::Result<bool> {
        match key {
            Key::Ctrl('c') => { return Ok(true); }
            Key::Ctrl('l') => self.full_redraw().await?,
            Key::F(5) => self.switch_buffer_rel(-1).await?,
            Key::F(6) => self.switch_buffer_rel(1).await?,
            _ => {
                let submitted = {
                    let mut state = self.state.lock().await;
                    let state = &mut *state;
                    let data = state.rooms.get_mut(&state.currentroom).unwrap();
                    data.editor.keypress(key)
                };
                if let Some(submitted) = submitted {
                    if self.handle_input(&submitted).await? {
                        return Ok(true);
                    }
                }
                self.full_redraw().await?;
            }
        }
        Ok(false)
    }

    async fn on_mouse(self: &Arc<Self>, x: u16, y: u16) -> io::Result<()> {
        let mut state = self.state.lock().await;
        if x < state.layout.rlsepx && usize::from(y) < state.roomlist.len() {
            state.switch_buffer(usize::from(y))?;
        }

        Ok(())
    }

    // Returns whether application should quit
    async fn handle_input(self: &Arc<Self>, line: &str) -> io::Result<bool> {
        let parsed = if line.starts_with("//") {
            Err(&line[1..])
        } else if line.starts_with("/") {
            match line.find(' ') {
                Some(idx) => Ok((&line[1..idx], &line[idx+1..])),
                None => Ok((&line[1..], ""))
            }
        } else {
            Err(line)
        };

        let room = {
            let mut state = self.state.lock().await;
            let room = state.currentroom.clone();
            state.user_room_activity(&room);
            room
        };

        match parsed {
            Err(msg) => {
                if room == *STATUS_ROOM {
                    bel::bel();
                    self.add_room_highlight(&STATUS_ROOM).await.unwrap();
                    self.append_history_item(&STATUS_ROOM,
                        HItem::Service("Cannot send messages in this magic status room".to_string(), true)
                    ).await.unwrap();
                } else {
                    let msg = msg.to_string();
                    let app = self.clone();
                    task::spawn(async move {
                        let line = Line::try_from(msg).unwrap();
                        if let Ok(Reply::Number(id)) = send_command(&app.conn, Command::Send {
                            roomname: room.clone(), reply_on: None, message: line.clone()
                        }).await {
                            app.add_message(Message {
                                id: Id::try_from(id).unwrap(),
                                reply_on: None,
                                roomname: room,
                                username: app.username.clone(),
                                timestamp: std::time::SystemTime::now(),
                                message: line,
                            }, false).await.unwrap();
                        }
                    });
                }
            }

            Ok(("quit", _)) | Ok(("exit", _)) => {
                return Ok(true);
            }

            Ok((cmd, _)) => {
                bel::bel();
                self.add_room_highlight(&STATUS_ROOM).await.unwrap();
                self.append_history_item(&STATUS_ROOM, HItem::Service(format!("Unknown command: '{}'", cmd), true)).await?;
            }
        }

        Ok(false)
    }

    async fn decide_push_notify(self: &Arc<Self>, room: &Word, mention: bool) -> bool {
        let state = self.state.lock().await;

        let private_chat = state.rooms.get(room).unwrap().members.len() <= 2;
        if private_chat {
            true
        } else {
            mention
        }
    }

    async fn add_room_highlight(self: &Arc<Self>, room: &Word) -> io::Result<()> {
        let mut state = self.state.lock().await;
        state.add_room_highlight(room);
        state.full_redraw()
    }

    async fn full_redraw(self: &Arc<Self>) -> io::Result<()> {
        self.state.lock().await.full_redraw()
    }
}

impl State {
    fn full_redraw(&mut self) -> io::Result<()> {
        self.termsize = termion::terminal_size()?;
        self.layout = Layout::compute(self.termsize, &self);
        print!("{}", clear::All);
        draw_vert_bar(&mut self.stdout, self.layout.rlsepx, 0, self.termsize.1)?;
        draw_vert_bar(&mut self.stdout, self.layout.nlsepx, 0, self.termsize.1 - 2)?;
        draw_hori_bar(&mut self.stdout, self.termsize.1 - 2, self.layout.rlsepx + 1, self.termsize.0)?;

        for (i, room) in self.roomlist.iter().enumerate() {
            print!("{}", cursor::Goto(1, u16::try_from(i).unwrap() + 1));
            if self.rooms.get(room).unwrap().highlight {
                print!("{}{}", color::Fg(color::Yellow), style::Bold);
            }
            if room == &self.currentroom {
                print!("{}", color::Bg(color::Blue));
            }
            print!("{}{}", room, style::Reset);
        }

        if self.currentroom.len() > 0 {
            let data = self.rooms.get_mut(&self.currentroom).unwrap();

            for (i, user) in data.members.iter().enumerate() {
                print!("{}", cursor::Goto(self.layout.nlsepx + 2, u16::try_from(i).unwrap() + 1));
                print!("{}", user);
            }

            let nicksepx = self.layout.rlsepx + u16::try_from(data.maxnicklen).unwrap() + 2;

            draw_vert_bar(&mut self.stdout, nicksepx, 0, self.termsize.1 - 2)?;

            let wid = self.layout.nlsepx - nicksepx - 2;
            let mut y = self.termsize.1 - 3;
            for item in data.history.iter().rev() {
                let (prefix, text_to_format, highlight) = match item {
                    HItem::Message(id, highlight) => {
                        let msg = self.msgs.get(&id).unwrap();
                        (format!("<{}>", msg.username), &*msg.message, *highlight)
                    }

                    HItem::Service(msg, highlight) => {
                        (String::from(" --"), msg.as_str(), *highlight)
                    }
                };

                let lines = wrap_text(text_to_format, wid);
                let mut done = false;
                for (i, line) in lines.iter().enumerate().rev() {
                    if i == 0 {
                        print!("{}", cursor::Goto(nicksepx - u16::try_from(prefix.len()).unwrap(), y + 1));
                        if highlight { print!("{}", color::Bg(color::Cyan)); }
                        print!("{}", prefix);
                        if highlight { print!("{}", color::Bg(color::Reset)); }
                    }
                    print!("{}{}", cursor::Goto(nicksepx + 3, y + 1), line);
                    if y == 0 { done = true; break; }
                    y -= 1;
                }
                if done { break; }
            }

            data.editor.set_wid(usize::from(self.termsize.0 - self.layout.rlsepx - 1));

            let (editor_str, editor_cursor) = data.editor.displayed();
            let editor_cursor = u16::try_from(editor_cursor).unwrap();
            print!("{}{}{}",
                   cursor::Goto(self.layout.rlsepx + 2, self.termsize.1),
                   editor_str,
                   cursor::Goto(self.layout.rlsepx + 2 + editor_cursor, self.termsize.1));
        }

        self.stdout.flush()
    }

    fn put_message(&mut self, msg: Message) -> io::Result<()> {
        self.msgs.insert(msg.id, msg);
        Ok(())
    }

    fn append_history(&mut self, room: &Word, item: HItem) -> io::Result<()> {
        let mut data = self.rooms.get_mut(room).unwrap();
        match &item {
            HItem::Message(id, _) => {
                let msg = self.msgs.get(&id).unwrap();
                data.maxnicklen = data.maxnicklen.max(msg.username.len() + 2);
            }
            HItem::Service(_, _) => {
                data.maxnicklen = data.maxnicklen.max(3);
            }
        }
        self.rooms.get_mut(room).unwrap().history.push(item);
        self.full_redraw()
    }

    fn switch_buffer(&mut self, index: usize) -> io::Result<()> {
        let room = self.roomlist[index].clone();
        self.user_room_activity(&room);
        self.currentroom = room;
        self.full_redraw()
    }

    fn add_room_highlight(&mut self, room: &Word) {
        self.rooms.get_mut(room).unwrap().highlight = true;
    }

    fn user_room_activity(&mut self, room: &Word) {
        self.rooms.get_mut(room).unwrap().highlight = false;
    }
}

async fn pushchan_thread(mut chan: mpsc::Receiver<PushMessage>, app: Arc<App>) {
    loop {
        match chan.next().await {
            None => break,
            Some(pm) => app.on_push(pm).await,
        }
    }
}

struct Options {
    server: String,
    username: Option<String>,
}

impl Default for Options {
    fn default() -> Self {
        Self { server: String::from("127.0.0.1:29538"), username: None }
    }
}

fn parse_options() -> Options {
    let mut options = Options::default();
    let defserver = options.server.clone();
    let server_descr = format!("Address (host:port) of tomsg server (default: {})", defserver);

    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Simple client for tomsg with auto-generated credentials.");

        ap.refer(&mut options.server)
            .add_option(&["-s", "--server"], Store, &server_descr)
            .metavar("ADDR");
        ap.refer(&mut options.username)
            .add_option(&["-u", "--user"], StoreOption,
                        "Username to connect as (if given, will prompt for password on stdin) (default: current system username)")
            .metavar("USER");

        ap.parse_args_or_exit();
    }
    options
}

async fn async_main() -> io::Result<()> {
    let options = parse_options();

    let (conn, pushchan) = Connection::connect(connection::Type::Plain, options.server).await?;

    let (user, pass) = auth::get_auth_info()?;

    let (user, pass) = match options.username {
        None => (Word::try_from(user.to_string()).unwrap()
                ,Line::try_from(pass.to_string()).unwrap()),
        Some(user) => {
            let pass = getpass::get_pass(&format!("Password for user '{}': ", user))?;
            if pass.len() == 0 {
                println!("Cancelled.");
                std::process::exit(1);
            }
            (Word::try_from(user).unwrap(), Line::try_from(pass).unwrap())
        }
    };

    send_command(&conn, Command::Register { username: user.clone(), password: pass.clone() }).await?;

    match send_command(&conn, Command::Login { username: user.clone(), password: pass }).await? {
        Reply::Ok => {},
        _ => {
            eprintln!("Failed to login!");
            return Ok(());
        },
    }

    debug!("initializing");

    let stdout = MouseTerminal::from(AlternateScreen::from(io::stdout().into_raw_mode()?));
    // let stdout = MouseTerminal::from(io::stdout());

    let app = Arc::new(App::new(conn, stdout, user));

    task::spawn(pushchan_thread(pushchan, app.clone()));

    app.full_redraw().await?;

    let app_clone = app.clone();
    task::spawn(async move {
        app_clone.fetch_rooms_list().await.unwrap();
    });

    for event in io::stdin().events() {
        match event? {
            Event::Key(key) => {
                if app.on_key(key).await? {
                    break;
                }
            }

            Event::Mouse(MouseEvent::Press(MouseButton::Left, x, y)) => {
                app.on_mouse(x - 1, y - 1).await?;
            }

            _ => {}
        }
    }

    Ok(())
}

fn main() -> io::Result<()> {
    runtime::Runtime::new()?.block_on(async_main())
}