extern crate clap;
extern crate termion;

mod display;
mod os;
mod postgresql;

use clap::{App, Arg};
use std::env;
use std::io::{stdout, Write};
use std::thread;
use std::time::Duration;
use termion::input::TermRead;
use termion::raw::IntoRawMode;

pub struct PGTOP {
    pub new: usize,
    pub old: usize,
    pub cpu: os::CPUSTAT,
}

fn move_to_end(line: u16) {
    let mut stdout = stdout().into_raw_mode().unwrap();
    write!(
        stdout,
        "{}{}",
        termion::cursor::Goto(1, line),
        termion::cursor::Show,
    )
    .unwrap();
}

fn main() {
    let matches = App::new("pg_top")
        .version("5.0.0")
        .about("'top' for PostgreSQL")
        .arg(
            Arg::with_name("DELAY")
                .help("set delay between screen updates")
                .default_value("1")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("COUNT")
                .help("change number of displays to show")
                .takes_value(true),
        )
        .get_matches();

    let delay = matches.value_of("DELAY").unwrap().parse::<u64>().unwrap();

    let infinite_loop = if matches.value_of("COUNT").is_none() {
        true
    } else {
        false
    };

    let mut count = if matches.value_of("COUNT").is_none() {
        1
    } else {
        matches.value_of("COUNT").unwrap().parse::<u64>().unwrap()
    };

    let mut conninfo = String::new();

    // Host is required.
    let pghost = if env::var("PGHOST").is_ok() {
        env::var("PGHOST").unwrap()
    } else {
        String::from("/tmp,/var/run/postgresql")
    };
    conninfo.push_str("host=");
    conninfo.push_str(&pghost);

    // User is required.
    let pguser = if env::var("PGUSER").is_ok() {
        env::var("PGUSER").unwrap()
    } else {
        env::var("USER").unwrap()
    };
    conninfo.push_str(" user=");
    conninfo.push_str(&pguser);

    // Database is not required and will default to PGUSER.
    if env::var("PGDATABASE").is_ok() {
        conninfo.push_str(" database=");
        env::var("PGDATABASE").unwrap();
    }

    // Port is not required.
    if env::var("PGPORT").is_ok() {
        conninfo.push_str(" port=");
        env::var("PGPORT").unwrap();
    }

    let mut pg_top = PGTOP {
        new: 1,
        old: 0,
        cpu: os::CPUSTAT {
            stat: [[0; 5], [0; 5]],
        },
    };

    let mut stdin = termion::async_stdin().keys();

    let mut stdout = stdout().into_raw_mode().unwrap();
    write!(stdout, "{}{}", termion::clear::All, termion::cursor::Hide).unwrap();

    loop {
        let mut client = postgresql::connect(&conninfo);
        let mut processes = postgresql::get_processes(&mut client);
        let terminal_size = termion::terminal_size().unwrap();

        /* Process any input. */

        let input = stdin.next();
        if let Some(Ok(key)) = input {
            match key {
                termion::event::Key::Char('q') => {
                    move_to_end(terminal_size.1);
                    break;
                }
                _ => { /* */ }
            }
        }

        /* Update the display. */

        display::loadavg(terminal_size.0 as usize);
        processes = display::process_summary(processes, terminal_size.0 as usize);
        pg_top = display::stat_cpu(pg_top, terminal_size.0 as usize);
        display::processes(processes, terminal_size.0 as usize);
        write!(stdout, "{}", termion::clear::AfterCursor,).unwrap();
        stdout.flush().unwrap();

        if infinite_loop == false {
            count = count - 1;
        }
        if count == 0 {
            move_to_end(terminal_size.1);
            break;
        }
        thread::sleep(Duration::from_secs(delay));
    }
}