Kademlia

use chrono::Local;
use clap::{Arg, App};
use colored::*;
use std::fs;
use std::path::Path;
use std::io::{stdout, Write};
use libp2p::identity::{ed25519, Keypair, PublicKey};
use libp2p::Multiaddr;

use async_std::{io, task};
use futures::prelude::*;
use libp2p::kad::record::store::MemoryStore;
use libp2p::kad::{
    record::Key, AddProviderOk, Kademlia, KademliaEvent, PeerRecord, PutRecordOk, QueryResult,
    Quorum, Record,
};
use libp2p::{
    build_development_transport, identity,
    mdns::{Mdns, MdnsEvent},
    swarm::NetworkBehaviourEventProcess,
    NetworkBehaviour, PeerId, Swarm,
};
use std::{
    error::Error,
    task::{Context, Poll},
};

fn main() -> Result<(), Box<dyn Error>> {
    // コマンドラインパラメータの解析
    let matches = App::new("Kadmlia-based Simple KVS").version("1.0").author("TAKAMI Torao <koiroha@gmail.com>")
        .about("Simple key-value store implementation using the Kademlia DHT of libp2p.")
        .arg(Arg::with_name("ID")
            .help("numeric node id")
            .required(true))
        .arg(Arg::with_name("daemon")
            .short("d")
            .long("daemon")
            .help("use when no input from stdin is accepted.")
            .required(false)
            .takes_value(false))
        .arg(Arg::with_name("bootstrap")
            .short("b")
            .long("bootstrap")
            .value_name("PEER_ID=NODE_ADDRESS")
            .help("bootstrap node address (e.g., 12D3KooWCzsfEkYrFf2djHem3qdagCjV4M6EXkcPUQzP9Xq76EPV=/ip4/127.0.0.1/tcp/6200)")
            .takes_value(true))
        .get_matches();

    // コマンドラインオプションからノード ID を取得し鍵ペアをロード
    let id:u32 = matches.value_of("ID").map(|s| s.parse().unwrap()).unwrap_or(0);
    let key_pair = load_keypair(id);

    let peer_id = PeerId::from(PublicKey::Ed25519(key_pair.public()));
    let kv_store = MemoryStore::new(peer_id.clone());
    let mut kademlia = Kademlia::new(peer_id.clone(), kv_store);
    log(id, format!("Peer ID: {}", peer_id));
    log(id, format!("Protocol Name: {}", String::from_utf8(kademlia.protocol_name().to_vec()).unwrap()));

    let bind_address = address(id);
    log(id, format!("Bind Address: {}", bind_address.to_string()));
    let transport = libp2p::build_development_transport(Keypair::Ed25519(key_pair)).unwrap();
    let behaviour = DHT{ id, kademlia };
    let mut swarm = Swarm::new(transport, behaviour, peer_id);
    Swarm::listen_on(&mut swarm, bind_address).unwrap();

    env_logger::init();

    // Create a random key for ourselves.
    let local_key = identity::Keypair::generate_ed25519();
    let local_peer_id = PeerId::from(local_key.public());

    // Set up a an encrypted DNS-enabled TCP Transport over the Mplex protocol.
    let transport = build_development_transport(local_key)?;

    // Create a swarm to manage peers and events.
    let mut swarm = {
        // Create a Kademlia behaviour.
        let store = MemoryStore::new(local_peer_id.clone());
        let kademlia = Kademlia::new(local_peer_id.clone(), store);
        let mdns = Mdns::new()?;
        let behaviour = DHT { kademlia, mdns };
        Swarm::new(transport, behaviour, local_peer_id)
    };

    // Read full lines from stdin
    let mut stdin = io::BufReader::new(io::stdin()).lines();

    // Listen on all interfaces and whatever port the OS assigns.
    Swarm::listen_on(&mut swarm, "/ip4/0.0.0.0/tcp/0".parse()?)?;

    // Kick it off.
    let mut listening = false;
    task::block_on(future::poll_fn(move |cx: &mut Context<'_>| {
        loop {
            match stdin.try_poll_next_unpin(cx)? {
                Poll::Ready(Some(line)) => handle_input_line(&mut swarm.kademlia, line),
                Poll::Ready(None) => panic!("Stdin closed"),
                Poll::Pending => break,
            }
        }
        loop {
            match swarm.poll_next_unpin(cx) {
                Poll::Ready(Some(event)) => println!("{:?}", event),
                Poll::Ready(None) => return Poll::Ready(Ok(())),
                Poll::Pending => {
                    if !listening {
                        if let Some(a) = Swarm::listeners(&swarm).next() {
                            println!("Listening on {:?}", a);
                            listening = true;
                        }
                    }
                    break;
                }
            }
        }
        Poll::Pending
    }))
}


#[derive(NetworkBehaviour)]
struct DHT {
    kademlia: Kademlia<MemoryStore>,
    mdns: Mdns,
}

impl NetworkBehaviourEventProcess<MdnsEvent> for DHT {
    // Called when `mdns` produces an event.
    fn inject_event(&mut self, event: MdnsEvent) {
        if let MdnsEvent::Discovered(list) = event {
            for (peer_id, multiaddr) in list {
                self.kademlia.add_address(&peer_id, multiaddr);
            }
        }
    }
}

impl NetworkBehaviourEventProcess<KademliaEvent> for DHT {
    // Called when `kademlia` produces an event.
    fn inject_event(&mut self, message: KademliaEvent) {
        match message {
            KademliaEvent::QueryResult { result, .. } => match result {
                QueryResult::GetProviders(Ok(ok)) => {
                    for peer in ok.providers {
                        println!(
                            "Peer {:?} provides key {:?}",
                            peer,
                            std::str::from_utf8(ok.key.as_ref()).unwrap()
                        );
                    }
                }
                QueryResult::GetProviders(Err(err)) => {
                    eprintln!("Failed to get providers: {:?}", err);
                }
                QueryResult::GetRecord(Ok(ok)) => {
                    for PeerRecord {
                        record: Record { key, value, .. },
                        ..
                    } in ok.records
                    {
                        println!(
                            "Got record {:?} {:?}",
                            std::str::from_utf8(key.as_ref()).unwrap(),
                            std::str::from_utf8(&value).unwrap(),
                        );
                    }
                }
                QueryResult::GetRecord(Err(err)) => {
                    eprintln!("Failed to get record: {:?}", err);
                }
                QueryResult::PutRecord(Ok(PutRecordOk { key })) => {
                    println!(
                        "Successfully put record {:?}",
                        std::str::from_utf8(key.as_ref()).unwrap()
                    );
                }
                QueryResult::PutRecord(Err(err)) => {
                    eprintln!("Failed to put record: {:?}", err);
                }
                QueryResult::StartProviding(Ok(AddProviderOk { key })) => {
                    println!(
                        "Successfully put provider record {:?}",
                        std::str::from_utf8(key.as_ref()).unwrap()
                    );
                }
                QueryResult::StartProviding(Err(err)) => {
                    eprintln!("Failed to put provider record: {:?}", err);
                }
                _ => {}
            },
            _ => {}
        }
    }
}

fn handle_input_line(kademlia: &mut Kademlia<MemoryStore>, line: String) {
    let mut args = line.split(" ");

    match args.next() {
        Some("GET") => {
            let key = {
                match args.next() {
                    Some(key) => Key::new(&key),
                    None => {
                        eprintln!("Expected key");
                        return;
                    }
                }
            };
            kademlia.get_record(&key, Quorum::One);
        }
        Some("GET_PROVIDERS") => {
            let key = {
                match args.next() {
                    Some(key) => Key::new(&key),
                    None => {
                        eprintln!("Expected key");
                        return;
                    }
                }
            };
            kademlia.get_providers(key);
        }
        Some("PUT") => {
            let key = {
                match args.next() {
                    Some(key) => Key::new(&key),
                    None => {
                        eprintln!("Expected key");
                        return;
                    }
                }
            };
            let value = {
                match args.next() {
                    Some(value) => value.as_bytes().to_vec(),
                    None => {
                        eprintln!("Expected value");
                        return;
                    }
                }
            };
            let record = Record {
                key,
                value,
                publisher: None,
                expires: None,
            };
            kademlia
                .put_record(record, Quorum::One)
                .expect("Failed to store record locally.");
        }
        Some("PUT_PROVIDER") => {
            let key = {
                match args.next() {
                    Some(key) => Key::new(&key),
                    None => {
                        eprintln!("Expected key");
                        return;
                    }
                }
            };

            kademlia
                .start_providing(key)
                .expect("Failed to start providing key");
        }
        _ => {
            eprintln!("expected GET, GET_PROVIDERS, PUT or PUT_PROVIDER");
        }
    }
}

fn load_keypair(id:u32) -> ed25519::Keypair {
    let path = format!("node/{}/", id);
    let dir = Path::new(&path);
    if !dir.is_dir() {
        fs::create_dir_all(dir).unwrap();
    }
    let file = dir.join("key");
    if file.is_file() {
        let mut binary = fs::read(file).unwrap();
        ed25519::Keypair::decode(&mut binary).unwrap()
    } else {
        let key_pair = ed25519::Keypair::generate();
        let binary = key_pair.encode();
        fs::write(file, binary.to_vec()).unwrap();
        key_pair
    }
}

fn address(id:u32) -> Multiaddr {
    let bind_address = format!("/ip4/127.0.0.1/tcp/{}", 6200 + id);
    bind_address.parse().unwrap()
}

fn prompt<T: std::str::FromStr>(id:u32, msg:&str) -> T {
    let color = color(id);
    print!("{}> ", msg.color(color));
    stdout().flush().unwrap();
    let mut s = String::new();
    std::io::stdin().read_line(&mut s).unwrap();
    s.trim().parse().ok().unwrap()
}

fn color(id:u32) -> Color {
    let colors = [Color::BrightRed, Color::BrightGreen, Color::BrightYellow, Color::BrightBlue, Color::BrightMagenta,
        Color::BrightCyan];
    colors[id as usize % colors.len()]
}

fn log(id:u32, msg:String) {
    let color = color(id);
    let tm = Local::now().format("%Y-%m-%d %H:%M:%S%.3f").to_string();
    println!("{} {}", tm.bright_black(), format!("Node{} {}", id, msg).color(color));
}