// Rshare - a file sharing program written in Rust.
// Copyright (C) 2017 Cooper Paul EdenDay
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
//
// Electronic mail: cedenday@protonmail.com
// Paper mail: 100 Front Street, #31, Bath, Maine, 04530, USA

#![forbid(unsafe_code)]

extern crate clap;
extern crate rustyline;
extern crate shlex;
extern crate exitcode;
extern crate app_dirs;
extern crate semver;
extern crate byteorder;
extern crate zopfli;
extern crate inflate;
extern crate crc;
extern crate base64;
extern crate rand;
extern crate ring;
extern crate untrusted;
extern crate rayon;
extern crate chashmap;
extern crate rusqlite;

#[macro_use]
extern crate serde;
extern crate bincode;

#[macro_use]
extern crate log;

mod cli;
mod logger;
mod constants;
mod server;
mod messaging;
mod client;
mod metadata;

use std::io::Write;
use std::str::FromStr;

use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::thread;

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener};

use chashmap::CHashMap;

use std::time::Instant;

use constants::PORT;

const NAME: &'static str = env!("CARGO_PKG_NAME");
const DESCRIPTION: &'static str = env!("CARGO_PKG_DESCRIPTION");
const AUTHORS: &'static str = env!("CARGO_PKG_AUTHORS");
const VERSION: &'static str = env!("CARGO_PKG_VERSION");

const APP_DATA: app_dirs::AppInfo = app_dirs::AppInfo {
    name: NAME,
    author: AUTHORS,
};

fn main() {
    let _ = logger::init();
    info!("Program started.");

    let rng = ring::rand::SystemRandom::new();

    let matches = cli::build_cli().get_matches();

    // Generate CLI tab completions.
    if let Some(sub_matches) = matches.subcommand_matches("completions") {
        let shell = sub_matches.value_of("SHELL").unwrap();
        let out = std::io::stdout();
        cli::build_cli()
            .gen_completions_to(NAME, shell.parse::<clap::Shell>().unwrap(), &mut out.lock());
        std::process::exit(exitcode::OK);
    }

    // Make sure that the metadata database is initialized.
    metadata::connect();

    let bind = {
        let bind_text = matches.value_of("bind").unwrap_or("::");
        let bind_ip = IpAddr::from(Ipv6Addr::from_str(bind_text).unwrap_or_else(|_| {
            println!("Could not parse IP '{}'.", bind_text);
            std::process::exit(exitcode::USAGE);
        }));
        SocketAddr::new(bind_ip, PORT)
    };

    let listener = TcpListener::bind(bind).unwrap_or_else(|_| {
        println!("Couldn't bind to port '{}', is it already in use?", PORT);
        std::process::exit(exitcode::USAGE);
    });

    // Parse servers from CLI input.
    let mut servers = std::collections::HashSet::new();
    if matches.is_present("server") {
        for ip_string in matches.values_of("server").unwrap() {
            if let Ok(ip) = Ipv6Addr::from_str(ip_string) {
                let ip = IpAddr::from(ip);
                let socket = std::net::SocketAddr::new(ip, PORT);
                servers.insert(socket);
            } else if let Ok(ip) = Ipv4Addr::from_str(ip_string) {
                let ip = IpAddr::from(ip.to_ipv6_mapped());
                let socket = std::net::SocketAddr::new(ip, PORT);
                servers.insert(socket);
            } else {
                println!("Couldn't parse server ip '{}' and it was not added to the list.",
                         ip_string);
            }
        }
    }

    // Stores the confirmed servers.
    let confirmed_servers = Arc::new(CHashMap::<IpAddr, ()>::new());

    // Stores the keys we are trying to download.
    let current_keys = Arc::new(CHashMap::<Vec<u8>, messaging::itc::Key>::new());

    // Stores the received advertisements for parts.
    let advertisements = Arc::new(CHashMap::<messaging::itc::Advert, ()>::new());

    info!("Starting server component.");
    {
        let confirmed_servers = confirmed_servers.clone();
        let current_keys = current_keys.clone();
        let advertisements = advertisements.clone();
        thread::spawn(move || {
            // Store the IPs that we connect to and the ones that connect to us.
            let registered_peers = Arc::new(CHashMap::<IpAddr, ()>::new());
            registered_peers.insert(IpAddr::from(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), ());

            // Store the ignored hashes to stop redundant concurrent downloading.
            let ignored_hashes = Arc::new(CHashMap::<Vec<u8>, Instant>::new());

            //Stores KEX state for decryption.
            let kex_state =
                Arc::new(CHashMap::<(Vec<u8>, IpAddr), messaging::itc::KexState>::new());

            // Try to stop DoS attacks by rate-limiting transfers.
            let active_downloads = Arc::new(AtomicUsize::new(0));
            let active_uploads = Arc::new(AtomicUsize::new(0));

            //Start the advertisement selector.
            {
                let registered_peers = registered_peers.clone();
                let advertisements = advertisements.clone();
                let ignored_hashes = ignored_hashes.clone();
                let kex_state = kex_state.clone();
                let active_downloads = active_downloads.clone();

                thread::spawn(move || {
                                  server::advert_selector(active_downloads,
                                                          kex_state,
                                                          registered_peers,
                                                          advertisements,
                                                          ignored_hashes)
                              });
            }

            // Main server loop to accept connections.
            loop {
                let confirmed_servers = confirmed_servers.clone();

                let registered_peers = registered_peers.clone();
                let advertisements = advertisements.clone();
                let current_keys = current_keys.clone();
                let ignored_hashes = ignored_hashes.clone();
                let kex_state = kex_state.clone();

                let active_downloads = active_downloads.clone();
                let active_uploads = active_uploads.clone();

                // If the connection succeeds spawn a new thread to handle it.
                if let Ok((connection, port)) = listener.accept() {
                    thread::spawn(move || {
                        let ip = port.ip();

                        let timeout = std::time::Duration::from_secs(constants::TIMEOUT);
                        connection.set_read_timeout(Some(timeout)).unwrap();

                        active_downloads.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
                        let wrapper = messaging::wrapper::decode(connection);
                        active_downloads.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);

                        match wrapper {
                            Ok(wrapper) => {
                                registered_peers.insert(ip, ());

                                let message = messaging::message::decode(&wrapper);
                                match message {
                                    Ok(message) => {
                                        server::message(server::MessageInput {
                                                            message,
                                                            ip,
                                                            kex_state,
                                                            registered_peers,
                                                            advertisements,
                                                            confirmed_servers,
                                                            current_keys,
                                                            ignored_hashes,
                                                            active_uploads,
                                                        });
                                    }
                                    Err(error) => {
                                        warn!("Error '{}' while parsing message from \
                                    peer at '{}'.",
                                              error,
                                              ip)
                                    }
                                }
                            }
                            Err(error) => {
                                warn!("Error '{}' while parsing wrapper from peer at '{}'.",
                                      error,
                                      ip)
                            }
                        }
                    });
                } else {
                    warn!("A connection to this server failed.");
                }
            }
        });
    }
    info!("Started server component.");

    // Make sure the servers can connect back.
    {
        let timeout = std::time::Duration::from_millis(500);
        for server in servers.clone() {
            if let Ok(mut connection) = std::net::TcpStream::connect(server) {
                let _ = connection.set_write_timeout(Some(timeout));
                let message = messaging::encode(&messaging::message::Message::Ping);
                let _ = connection.write_all(&message);
            } else {
                println!("Couldn't connect to server '{}', removing it from the list.",
                         server.ip());
                servers.remove(&server);
            }
        }

        thread::sleep(timeout);

        for server in &servers {
            if !confirmed_servers.contains_key(&server.ip()) {
                println!("Could connect to server '{}' but it couldn't connect back.\n\
                Hint: check if TCP port '{}' is allowed through the firewall(s) in front of\n\
                this program and if the program is updated to the latest version.",
                         server.ip(),
                         PORT);
            }
        }
    }

    // Start the advertisement requester.
    {
        let servers = servers.clone();
        let advertisements = advertisements.clone();
        let current_keys = current_keys.clone();
        thread::spawn(move || server::advert_requester(&servers, advertisements, current_keys));
    }

    if !matches.is_present("no-client") {
        // Start the client.
        client::client(&rng,
                       current_keys,
                       matches.is_present("quiet"),
                       servers.is_empty());
    } else {
        // Park the thread.
        loop {
            std::thread::park();
        }
    }
}
