// 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

use std;
use messaging;
use messaging::itc;
use metadata;
use constants;
use rand;
use ring;
use rusqlite;
use ring::rand::SecureRandom;
use untrusted;

use chashmap::CHashMap;

use std::collections::HashSet;

use std::io::Write;

use std::sync::Arc;

use std::net::{IpAddr, SocketAddr};

use std::time::{Instant, Duration};

pub struct MessageInput {
    pub message: messaging::message::Message,
    pub ip: std::net::IpAddr,
    pub kex_state: Arc<CHashMap<(Vec<u8>, IpAddr), itc::KexState>>,
    pub registered_peers: Arc<CHashMap<IpAddr, ()>>,
    pub advertisements: Arc<CHashMap<itc::Advert, ()>>,
    pub confirmed_servers: Arc<CHashMap<IpAddr, ()>>,
    pub current_keys: Arc<CHashMap<Vec<u8>, itc::Key>>,
    pub ignored_hashes: Arc<CHashMap<Vec<u8>, Instant>>,
    pub active_uploads: Arc<std::sync::atomic::AtomicUsize>,
}

pub fn message(input: MessageInput) {
    let MessageInput {
        message,
        ip,
        kex_state,
        registered_peers,
        advertisements,
        confirmed_servers,
        current_keys,
        ignored_hashes,
        active_uploads,
    } = input;

    let rng = ring::rand::SystemRandom::new();
    let eae_hash = &ring::digest::SHA256;
    let eae_cipher = &ring::aead::AES_256_GCM;

    let connection = metadata::connect();

    match message {
        messaging::message::Message::Ping => {
            let message_send = messaging::message::Message::Pong;
            messaging::send(confirmed_servers, &message_send, ip)
        }

        messaging::message::Message::Pong => {
            confirmed_servers.insert(ip, ());
        }

        messaging::message::Message::IWantKey { key } => {
            for peer in registered_peers.as_ref().clone() {
                let peer = peer.0;
                let message_send = messaging::message::Message::DoYouHaveKey {
                    key: key.clone(),
                    contact_address: ip,
                };
                messaging::send(confirmed_servers.clone(), &message_send, peer)
            }
        }

        messaging::message::Message::DoYouHaveKey {
            key,
            contact_address,
        } => {
            let raw_file_path: Result<String, rusqlite::Error> =
                connection
                    .prepare("SELECT file_path \
                    FROM metadata \
                    WHERE key = ?")
                    .unwrap()
                    .query_row(&[&key], |row| row.get(0));

            if let Ok(raw_file_path) = raw_file_path {
                let file_path = std::path::Path::new(&raw_file_path);

                if file_path.exists() {
                    let mut query = connection
                        .prepare("SELECT hash \
                        FROM parts \
                        INNER JOIN metadata ON metadata.id = parts.metadata_id \
                        WHERE metadata.key = ?")
                        .unwrap();

                    let hashes = query
                        .query_map(&[&key], |row| {
                            let output: Vec<u8> = row.get(0);
                            output
                        })
                        .unwrap();

                    let randomized_hashes = {
                        let mut randomized_hashes: Vec<Vec<u8>> =
                            hashes.map(|x| x.unwrap()).collect();
                        use rand::Rng;
                        let mut rng = rand::thread_rng();
                        rng.shuffle(&mut randomized_hashes);
                        randomized_hashes
                    };

                    for hash in randomized_hashes {
                        let message = messaging::message::Message::DoYouWantHash {
                            hash: hash.to_vec(),
                            public_key: key.clone(),
                        };
                        messaging::send(registered_peers.clone(), &message, contact_address);
                    }
                } else {
                    warn!("File '{}' not found.", raw_file_path);
                }
            }
        }

        messaging::message::Message::DoYouWantHash { hash, public_key } => {
            let advert = itc::Advert {
                ip,
                hash,
                public_key,
            };
            advertisements.insert(advert, ());
        }

        messaging::message::Message::IWantHash { hash, public_key } => {
            while !(active_uploads.load(std::sync::atomic::Ordering::SeqCst) < constants::MAX_ULS) {
                std::thread::sleep(std::time::Duration::from_millis(50));
            }

            active_uploads.fetch_add(1, std::sync::atomic::Ordering::SeqCst);

            let file_metadata = connection
                .prepare("SELECT file_path, nonce \
                FROM metadata \
                INNER JOIN parts ON parts.metadata_id = metadata.id \
                WHERE parts.hash = ?")
                .unwrap()
                .query_row(&[&hash], |row| {
                    itc::Metadata {
                        file_path: row.get(0),
                        nonce: row.get(1),
                    }
                });

            if let Ok(file_metadata) = file_metadata {
                let part = connection
                    .prepare("SELECT signature, part_number \
                    FROM parts \
                    WHERE hash = ?")
                    .unwrap()
                    .query_row(&[&hash], |row| {
                        let part_number: i64 = row.get(1);
                        metadata::Part {
                            signature: row.get(0),
                            part_number: part_number as u64,
                        }
                    })
                    .unwrap();

                let part_num = part.part_number;

                if let Ok(plaintext) = metadata::read_part(&file_metadata.file_path, part_num) {

                    let disk_hash = metadata::hash(&plaintext, &file_metadata.nonce, part_num);

                    if ring::constant_time::verify_slices_are_equal(&hash, &disk_hash).is_ok() {
                        let salt = {
                            let mut salt = vec!(0u8; ring::hmac::recommended_key_len(eae_hash));
                            let _ = ring::hmac::SigningKey::generate_serializable(eae_hash,
                                                                                  &rng,
                                                                                  &mut salt);
                            salt
                        };

                        let (my_public_key, nonce, ciphertext) = {
                            let my_private_key =
                            ring::agreement::EphemeralPrivateKey::generate(&ring::agreement::X25519,
                                                                           &rng)
                                .expect("Couldn't generate Ephemeral Key pair.");

                            let mut my_public_key_buf = [0u8; ring::agreement::PUBLIC_KEY_MAX_LEN];
                            let my_public_key;
                            {
                                my_public_key = &mut my_public_key_buf[..my_private_key
                                                           .public_key_len()];
                                my_private_key.compute_public_key(my_public_key).unwrap();
                            }
                            let peer_public_key = untrusted::Input::from(&public_key);

                            let nonce = {
                                let mut nonce = [0u8; 12];
                                rng.fill(&mut nonce).expect("Couldn't generate nonce");
                                nonce
                            };

                            let ciphertext =
                                ring::agreement::agree_ephemeral(my_private_key,
                                                                 &ring::agreement::X25519,
                                                                 peer_public_key,
                                                                 ring::error::Unspecified,
                                                                 |key_material| {
                                    let key = {
                                        let salt = ring::hmac::SigningKey::new(eae_hash, &salt);
                                        let mut key_bytes =
                                            vec!(0u8; ring::hmac::recommended_key_len(eae_hash));
                                        ring::hkdf::extract_and_expand(&salt,
                                                                       key_material,
                                                                       &file_metadata.nonce,
                                                                       &mut key_bytes);
                                        ring::aead::SealingKey::new(eae_cipher, &key_bytes)?
                                    };

                                    let tag_len = key.algorithm().tag_len();
                                    let plaintext_len = plaintext.len();
                                    let mut in_out = plaintext;
                                    in_out.resize(plaintext_len + tag_len, 0);

                                    let output_size = ring::aead::seal_in_place(&key,
                                                                                &nonce,
                                                                                &[],
                                                                                &mut in_out,
                                                                                tag_len)?;
                                    in_out.truncate(output_size);

                                    Ok(in_out)
                                })
                                        .unwrap();

                            (my_public_key.to_vec(), nonce, ciphertext)
                        };

                        let message = messaging::message::Message::HereYouGo {
                            public_key: my_public_key,
                            decryption_salt: salt,
                            nonce,
                            hash,
                            part: part.clone(),
                            ciphertext,
                        };
                        messaging::send(registered_peers, &message, ip)
                    } else {
                        warn!("Detected change in file '{}'.", &file_metadata.file_path);
                    }
                } else {
                    warn!("Tried to read part that doesn't exit from file '{}'.",
                          file_metadata.file_path);
                }
            } else {
                warn!("Peer asked us for a part we don't have.");
            }

            active_uploads.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
        }

        messaging::message::Message::HereYouGo {
            public_key,
            decryption_salt,
            nonce,
            hash,
            part,
            ciphertext,
        } => {
            let current_kex = kex_state.remove(&(hash.clone(), ip));

            if let Some(itc::KexState {
                            my_private_key,
                            link_public_key,
                        }) = current_kex {
                let peer_public_key = untrusted::Input::from(&public_key);

                if let Some(key) = current_keys.get(&link_public_key) {
                    let link_nonce = key.nonce.clone();

                    let data = ring::agreement::agree_ephemeral(my_private_key,
                                                                &ring::agreement::X25519,
                                                                peer_public_key,
                                                                ring::error::Unspecified,
                                                                |key_material| {
                        let key = {
                            let salt = ring::hmac::SigningKey::new(eae_hash, &decryption_salt);
                            let mut key_bytes =
                                vec!(0u8; ring::hmac::recommended_key_len(eae_hash));
                            ring::hkdf::extract_and_expand(&salt,
                                                           key_material,
                                                           &link_nonce,
                                                           &mut key_bytes);

                            ring::aead::OpeningKey::new(eae_cipher, &key_bytes)?
                        };

                        let mut in_out = ciphertext;

                        Ok(ring::aead::open_in_place(&key, &nonce, &[], 0, &mut in_out)?
                               .to_vec())
                    });
                    if let Ok(data) = data {
                        if data.len() as u64 <= metadata::PART_SIZE {

                            let message = untrusted::Input::from(&hash);
                            let untrusted_public_key = untrusted::Input::from(&link_public_key);
                            let signature = untrusted::Input::from(&part.signature);

                            if ring::signature::verify(&ring::signature::ED25519,
                                                       untrusted_public_key,
                                                       message,
                                                       signature)
                                       .is_ok() {

                                let my_hash = metadata::hash(&data, &link_nonce, part.part_number);

                                if ring::constant_time::verify_slices_are_equal(&hash, &my_hash)
                                       .is_ok() {

                                    let result = metadata::write_part(&key.file_path,
                                                                      part.part_number,
                                                                      &data);
                                    if result.is_ok() {

                                        connection
                                            .execute("INSERT OR IGNORE INTO metadata (\
                                            num_parts, file_path, key, nonce\
                                            ) VALUES ($1, $2, $3, $4)",
                                                     &[&(key.num_parts as i64),
                                                       &key.file_path,
                                                       &link_public_key,
                                                       &link_nonce.to_vec()])
                                            .unwrap();

                                        let metadata_id: i64 =
                                            connection.
                                                prepare("SELECT id \
                                                FROM metadata \
                                                WHERE key = ?")
                                                .unwrap()
                                                .query_row(&[&link_public_key], |row| row.get(0))
                                                .unwrap();

                                        connection
                                            .execute("INSERT OR IGNORE INTO parts (\
                                            hash,\
                                            signature,\
                                            part_number,\
                                            metadata_id)\
                                            VALUES ($1, $2, $3, $4)",
                                                     &[&my_hash,
                                                       &part.signature,
                                                       &(part.part_number as i64),
                                                       &metadata_id])
                                            .unwrap();

                                        // Remove key from current_keys if no parts are left.
                                        {
                                            let num_parts_disk: i64 =
                                                connection
                                                    .prepare("SELECT COUNT(*) \
                                                    FROM parts \
                                                    INNER JOIN metadata \
                                                    ON metadata.id = parts.metadata_id \
                                                    WHERE metadata.key = ?")
                                                    .unwrap()
                                                    .query_row(&[&link_public_key],
                                                               |row| row.get(0))
                                                    .unwrap();

                                            if num_parts_disk as u64 == key.num_parts {
                                                current_keys.remove(&link_public_key);
                                            }
                                        }
                                    } else {
                                        warn!("Failed to write part to file '{}'.", &key.file_path);
                                    }
                                } else {
                                    warn!("Failed to verify checksum of part from peer at '{}'.",
                                          ip);
                                }
                            } else {
                                warn!("Failed to verify origin of hash from peer at '{}'.", ip);
                            }
                        } else {
                            warn!("Received an over-sized part from peer at '{}'.", ip);
                        }
                    } else {
                        warn!("Failed to decrypt part from peer at '{}'.", ip);
                    }
                }
            } else {
                warn!("Peer at '{}' sent us a part we didn't ask for.", ip);
            }
            ignored_hashes.remove(&hash);
        }
    }
}

pub fn advert_selector(active_downloads: Arc<std::sync::atomic::AtomicUsize>,
                       kex_state: Arc<CHashMap<(Vec<u8>, IpAddr), itc::KexState>>,
                       registered_peers: Arc<CHashMap<IpAddr, ()>>,
                       advertisements: Arc<CHashMap<itc::Advert, ()>>,
                       ignored_hashes: Arc<CHashMap<Vec<u8>, Instant>>) {
    use rand::distributions::{IndependentSample, Range};
    let mut thread_rng = rand::thread_rng();
    let rng = ring::rand::SystemRandom::new();
    let connection = metadata::connect();

    loop {
        std::thread::sleep(std::time::Duration::from_secs(1));

        for entry in ignored_hashes.as_ref().clone() {
            if Instant::now().duration_since(entry.1) > Duration::from_secs(constants::TIMEOUT) {
                ignored_hashes.remove(&entry.0);
            }
        }

        if active_downloads.load(std::sync::atomic::Ordering::SeqCst) < constants::MAX_DLS {
            let adverts = advertisements.as_ref().clone();
            let good_adverts: CHashMap<itc::Advert, ()> = adverts
                .into_iter()
                .skip_while(|advert| {
                    let on_disk = connection
                        .prepare("SELECT part_number \
                        FROM parts \
                        WHERE hash = ?")
                        .unwrap()
                        .query_row(&[&advert.0.hash], |_| {})
                        .is_ok();
                    ignored_hashes.contains_key(&advert.0.hash) || on_disk
                })
                .collect();
            if !good_adverts.is_empty() {
                let distribution = Range::new(0, good_adverts.len());
                let advert = good_adverts
                    .into_iter()
                    .nth(distribution.ind_sample(&mut thread_rng));
                if let Some(advert) = advert {
                    let advert = advert.0;

                    advertisements.remove(&advert);
                    ignored_hashes.insert(advert.hash.clone(), Instant::now());

                    let private_key =
                        ring::agreement::EphemeralPrivateKey::generate(&ring::agreement::X25519,
                                                                       &rng)
                                .expect("Couldn't generate Ephemeral Key pair.");

                    let mut public_key_buf = [0u8; ring::agreement::PUBLIC_KEY_MAX_LEN];
                    let public_key;
                    {
                        public_key = &mut public_key_buf[..private_key.public_key_len()];
                        private_key.compute_public_key(public_key).unwrap();
                    }

                    kex_state.insert((advert.hash.clone(), advert.ip),
                                     itc::KexState {
                                         my_private_key: private_key,
                                         link_public_key: advert.public_key.clone(),
                                     });

                    let message = messaging::message::Message::IWantHash {
                        hash: advert.hash.clone(),
                        public_key: public_key.to_vec(),
                    };
                    messaging::send(registered_peers.clone(), &message, advert.ip);
                }
            }
        }
    }
}

pub fn advert_requester(servers: &HashSet<SocketAddr>,
                        advertisements: Arc<CHashMap<itc::Advert, ()>>,
                        current_keys: Arc<CHashMap<Vec<u8>, itc::Key>>) {
    loop {
        std::thread::sleep(std::time::Duration::from_secs(constants::ADVERTS_INTERVAL));
        let default: usize = 0;
        let mut adverts_per_key = std::collections::HashMap::<Vec<u8>, usize>::new();
        for advert in advertisements.as_ref().clone() {
            let advert: messaging::itc::Advert = advert.0;
            if !adverts_per_key.contains_key(&advert.hash) {
                adverts_per_key.insert(advert.hash.clone(), default);
            }
            *adverts_per_key.get_mut(&advert.hash).unwrap() += 1;
        }
        for key in current_keys.as_ref().clone() {
            let count = adverts_per_key.get(&key.0).unwrap_or(&default);
            if count < &constants::MAX_DLS {
                for server in servers {
                    let connection = std::net::TcpStream::connect(server);
                    if let Ok(mut connection) = connection {
                        let message = messaging::message::Message::IWantKey { key: key.0.clone() };
                        let _ = connection.write_all(&messaging::encode(&message));
                    } else {
                        warn!("Server at '{}' is unresponsive.", server);
                    }
                }
            }
        }
    }
}
