// Copyright 2015, 2016 Ethcore (UK) Ltd.
// This file is part of Parity.

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

// Parity 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 Parity.  If not, see <http://www.gnu.org/licenses/>.

//! Peer status and capabilities.

use rlp::{DecoderError, RlpDecodable, RlpEncodable, RlpStream, Stream, UntrustedRlp, View};
use util::{H256, U256};

use super::buffer_flow::FlowParams;

// recognized handshake/announcement keys.
// unknown keys are to be skipped, known keys have a defined order.
// their string values are defined in the LES spec.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Ord, PartialOrd)]
enum Key {
	ProtocolVersion,
	NetworkId,
	HeadTD,
	HeadHash,
	HeadNum,
	GenesisHash,
	ServeHeaders,
	ServeChainSince,
	ServeStateSince,
	TxRelay,
	BufferLimit,
	BufferCostTable,
	BufferRechargeRate,
}

impl Key {
	// get the string value of this key.
	fn as_str(&self) -> &'static str {
		match *self {
			Key::ProtocolVersion => "protocolVersion",
			Key::NetworkId => "networkId",
			Key::HeadTD => "headTd",
			Key::HeadHash => "headHash",
			Key::HeadNum => "headNum",
			Key::GenesisHash => "genesisHash",
			Key::ServeHeaders => "serveHeaders",
			Key::ServeChainSince => "serveChainSince",
			Key::ServeStateSince => "serveStateSince",
			Key::TxRelay => "txRelay",
			Key::BufferLimit => "flowControl/BL",
			Key::BufferCostTable => "flowControl/MRC",
			Key::BufferRechargeRate => "flowControl/MRR",
		}
	}

	// try to parse the key value from a string.
	fn from_str(s: &str) -> Option<Self> {
		match s {
			"protocolVersion" => Some(Key::ProtocolVersion),
			"networkId" => Some(Key::NetworkId),
			"headTd" => Some(Key::HeadTD),
			"headHash" => Some(Key::HeadHash),
			"headNum" => Some(Key::HeadNum),
			"genesisHash" => Some(Key::GenesisHash),
			"serveHeaders" => Some(Key::ServeHeaders),
			"serveChainSince" => Some(Key::ServeChainSince),
			"serveStateSince" => Some(Key::ServeStateSince),
			"txRelay" => Some(Key::TxRelay),
			"flowControl/BL" => Some(Key::BufferLimit),
			"flowControl/MRC" => Some(Key::BufferCostTable),
			"flowControl/MRR" => Some(Key::BufferRechargeRate),
			_ => None
		}
	}
}

/// Network ID structure.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u32)]
pub enum NetworkId {
	/// ID for the mainnet
	Mainnet = 1,
	/// ID for the testnet
	Testnet = 0,
}

impl NetworkId {
	fn from_raw(raw: u32) -> Option<Self> {
		match raw {
			0 => Some(NetworkId::Testnet),
			1 => Some(NetworkId::Mainnet),
			_ => None,
		}
	}
}

// helper for decoding key-value pairs in the handshake or an announcement.
struct Parser<'a> {
	pos: usize,
	rlp: UntrustedRlp<'a>,
}

impl<'a> Parser<'a> {
	// expect a specific next key, and decode the value.
	// error on unexpected key or invalid value.
	fn expect<T: RlpDecodable>(&mut self, key: Key) -> Result<T, DecoderError> {
		self.expect_raw(key).and_then(|item| item.as_val())
	}

	// expect a specific next key, and get the value's RLP.
	// if the key isn't found, the position isn't advanced.
	fn expect_raw(&mut self, key: Key) -> Result<UntrustedRlp<'a>, DecoderError> {
		let pre_pos = self.pos;
		if let Some((k, val)) = try!(self.get_next()) {
			if k == key { return Ok(val) }
		}

		self.pos = pre_pos;
		Err(DecoderError::Custom("Missing expected key"))
	}

	// get the next key and value RLP.
	fn get_next(&mut self) -> Result<Option<(Key, UntrustedRlp<'a>)>, DecoderError> {
		while self.pos < self.rlp.item_count() {
			let pair = try!(self.rlp.at(self.pos));
			let k: String = try!(pair.val_at(0));

			self.pos += 1;
			match Key::from_str(&k) {
				Some(key) => return Ok(Some((key , try!(pair.at(1))))),
				None => continue,
			}
		}

		Ok(None)
	}
}

// Helper for encoding a key-value pair
fn encode_pair<T: RlpEncodable>(key: Key, val: &T) -> Vec<u8> {
	let mut s = RlpStream::new_list(2);
	s.append(&key.as_str()).append(val);
	s.out()
}

// Helper for encoding a flag.
fn encode_flag(key: Key) -> Vec<u8> {
	let mut s = RlpStream::new_list(2);
	s.append(&key.as_str()).append_empty_data();
	s.out()
}

/// A peer status message.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Status {
	/// Protocol version.
	pub protocol_version: u32,
	/// Network id of this peer.
	pub network_id: NetworkId,
	/// Total difficulty of the head of the chain.
	pub head_td: U256,
	/// Hash of the best block.
	pub head_hash: H256,
	/// Number of the best block.
	pub head_num: u64,
	/// Genesis hash
	pub genesis_hash: H256,
	/// Last announced chain head and reorg depth to common ancestor.
	pub last_head: Option<(H256, u64)>,
}

/// Peer capabilities.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Capabilities {
	/// Whether this peer can serve headers
	pub serve_headers: bool,
	/// Earliest block number it can serve block/receipt requests for.
	pub serve_chain_since: Option<u64>,
	/// Earliest block number it can serve state requests for.
	pub serve_state_since: Option<u64>,
	/// Whether it can relay transactions to the eth network.
	pub tx_relay: bool,
}

impl Default for Capabilities {
	fn default() -> Self {
		Capabilities {
			serve_headers: true,
			serve_chain_since: None,
			serve_state_since: None,
			tx_relay: false,
		}
	}
}

/// Attempt to parse a handshake message into its three parts:
///   - chain status
///   - serving capabilities
///   - buffer flow parameters
pub fn parse_handshake(rlp: UntrustedRlp) -> Result<(Status, Capabilities, FlowParams), DecoderError> {
	let mut parser = Parser {
		pos: 0,
		rlp: rlp,
	};

	let status = Status {
		protocol_version: try!(parser.expect(Key::ProtocolVersion)),
		network_id: try!(parser.expect(Key::NetworkId)
			.and_then(|id: u32| NetworkId::from_raw(id).ok_or(DecoderError::Custom("Invalid network ID")))),
		head_td: try!(parser.expect(Key::HeadTD)),
		head_hash: try!(parser.expect(Key::HeadHash)),
		head_num: try!(parser.expect(Key::HeadNum)),
		genesis_hash: try!(parser.expect(Key::GenesisHash)),
		last_head: None,
	};

	let capabilities = Capabilities {
		serve_headers: parser.expect_raw(Key::ServeHeaders).is_ok(),
		serve_chain_since: parser.expect(Key::ServeChainSince).ok(),
		serve_state_since: parser.expect(Key::ServeStateSince).ok(),
		tx_relay: parser.expect_raw(Key::TxRelay).is_ok(),
	};

	let flow_params = FlowParams::new(
		try!(parser.expect(Key::BufferLimit)),
		try!(parser.expect(Key::BufferCostTable)),
		try!(parser.expect(Key::BufferRechargeRate)),
	);

	Ok((status, capabilities, flow_params))
}

/// Write a handshake, given status, capabilities, and flow parameters.
pub fn write_handshake(status: &Status, capabilities: &Capabilities, flow_params: &FlowParams) -> Vec<u8> {
	let mut pairs = Vec::new();
	pairs.push(encode_pair(Key::ProtocolVersion, &status.protocol_version));
	pairs.push(encode_pair(Key::NetworkId, &(status.network_id as u32)));
	pairs.push(encode_pair(Key::HeadTD, &status.head_td));
	pairs.push(encode_pair(Key::HeadHash, &status.head_hash));
	pairs.push(encode_pair(Key::HeadNum, &status.head_num));
	pairs.push(encode_pair(Key::GenesisHash, &status.genesis_hash));

	if capabilities.serve_headers {
		pairs.push(encode_flag(Key::ServeHeaders));
	}
	if let Some(ref serve_chain_since) = capabilities.serve_chain_since {
		pairs.push(encode_pair(Key::ServeChainSince, serve_chain_since));
	}
	if let Some(ref serve_state_since) = capabilities.serve_state_since {
		pairs.push(encode_pair(Key::ServeStateSince, serve_state_since));
	}
	if capabilities.tx_relay {
		pairs.push(encode_flag(Key::TxRelay));
	}

	pairs.push(encode_pair(Key::BufferLimit, flow_params.limit()));
	pairs.push(encode_pair(Key::BufferCostTable, flow_params.cost_table()));
	pairs.push(encode_pair(Key::BufferRechargeRate, flow_params.recharge_rate()));

	let mut stream = RlpStream::new_list(pairs.len());

	for pair in pairs {
		stream.append_raw(&pair, 1);
	}

	stream.out()
}

/// An announcement of new chain head or capabilities made by a peer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Announcement {
	/// Hash of the best block.
	pub head_hash: H256,
	/// Number of the best block.
	pub head_num: u64,
	/// Head total difficulty
	pub head_td: U256,
	/// reorg depth to common ancestor of last announced head.
	pub reorg_depth: u64,
	/// optional new header-serving capability. false means "no change"
	pub serve_headers: bool,
	/// optional new state-serving capability
	pub serve_state_since: Option<u64>,
	/// optional new chain-serving capability
	pub serve_chain_since: Option<u64>,
	/// optional new transaction-relay capability. false means "no change"
	pub tx_relay: bool,
	// TODO: changes in buffer flow?
}

/// Parse an announcement.
pub fn parse_announcement(rlp: UntrustedRlp) -> Result<Announcement, DecoderError> {
	let mut last_key = None;

	let mut announcement = Announcement {
		head_hash: try!(rlp.val_at(0)),
		head_num: try!(rlp.val_at(1)),
		head_td: try!(rlp.val_at(2)),
		reorg_depth: try!(rlp.val_at(3)),
		serve_headers: false,
		serve_state_since: None,
		serve_chain_since: None,
		tx_relay: false,
	};

	let mut parser = Parser {
		pos: 4,
		rlp: rlp,
	};

	while let Some((key, item)) = try!(parser.get_next()) {
		if Some(key) <= last_key { return Err(DecoderError::Custom("Invalid announcement key ordering")) }
		last_key = Some(key);

		match key {
			Key::ServeHeaders => announcement.serve_headers = true,
			Key::ServeStateSince => announcement.serve_state_since = Some(try!(item.as_val())),
			Key::ServeChainSince => announcement.serve_chain_since = Some(try!(item.as_val())),
			Key::TxRelay => announcement.tx_relay = true,
			_ => return Err(DecoderError::Custom("Nonsensical key in announcement")),
		}
	}

	Ok(announcement)
}

/// Write an announcement out.
pub fn write_announcement(announcement: &Announcement) -> Vec<u8> {
	let mut pairs = Vec::new();
	if announcement.serve_headers {
		pairs.push(encode_flag(Key::ServeHeaders));
	}
	if let Some(ref serve_chain_since) = announcement.serve_chain_since {
		pairs.push(encode_pair(Key::ServeChainSince, serve_chain_since));
	}
	if let Some(ref serve_state_since) = announcement.serve_state_since {
		pairs.push(encode_pair(Key::ServeStateSince, serve_state_since));
	}
	if announcement.tx_relay {
		pairs.push(encode_flag(Key::TxRelay));
	}

	let mut stream = RlpStream::new_list(4 + pairs.len());
	stream
		.append(&announcement.head_hash)
		.append(&announcement.head_num)
		.append(&announcement.head_td)
		.append(&announcement.reorg_depth);

	for item in pairs {
		stream.append_raw(&item, 1);
	}

	stream.out()
}

#[cfg(test)]
mod tests {
	use super::*;
	use super::super::buffer_flow::FlowParams;
	use util::{U256, H256, FixedHash};
	use rlp::{RlpStream, Stream ,UntrustedRlp, View};

	#[test]
	fn full_handshake() {
		let status = Status {
			protocol_version: 1,
			network_id: NetworkId::Mainnet,
			head_td: U256::default(),
			head_hash: H256::default(),
			head_num: 10,
			genesis_hash: H256::zero(),
			last_head: None,
		};

		let capabilities = Capabilities {
			serve_headers: true,
			serve_chain_since: Some(5),
			serve_state_since: Some(8),
			tx_relay: true,
		};

		let flow_params = FlowParams::new(
			1_000_000.into(),
			Default::default(),
			1000.into(),
		);

		let handshake = write_handshake(&status, &capabilities, &flow_params);

		let (read_status, read_capabilities, read_flow)
			= parse_handshake(UntrustedRlp::new(&handshake)).unwrap();

		assert_eq!(read_status, status);
		assert_eq!(read_capabilities, capabilities);
		assert_eq!(read_flow, flow_params);
	}

	#[test]
	fn partial_handshake() {
		let status = Status {
			protocol_version: 1,
			network_id: NetworkId::Mainnet,
			head_td: U256::default(),
			head_hash: H256::default(),
			head_num: 10,
			genesis_hash: H256::zero(),
			last_head: None,
		};

		let capabilities = Capabilities {
			serve_headers: false,
			serve_chain_since: Some(5),
			serve_state_since: None,
			tx_relay: true,
		};

		let flow_params = FlowParams::new(
			1_000_000.into(),
			Default::default(),
			1000.into(),
		);

		let handshake = write_handshake(&status, &capabilities, &flow_params);

		let (read_status, read_capabilities, read_flow)
			= parse_handshake(UntrustedRlp::new(&handshake)).unwrap();

		assert_eq!(read_status, status);
		assert_eq!(read_capabilities, capabilities);
		assert_eq!(read_flow, flow_params);
	}

	#[test]
	fn skip_unknown_keys() {
		let status = Status {
			protocol_version: 1,
			network_id: NetworkId::Mainnet,
			head_td: U256::default(),
			head_hash: H256::default(),
			head_num: 10,
			genesis_hash: H256::zero(),
			last_head: None,
		};

		let capabilities = Capabilities {
			serve_headers: false,
			serve_chain_since: Some(5),
			serve_state_since: None,
			tx_relay: true,
		};

		let flow_params = FlowParams::new(
			1_000_000.into(),
			Default::default(),
			1000.into(),
		);

		let handshake = write_handshake(&status, &capabilities, &flow_params);
		let interleaved = {
			let handshake = UntrustedRlp::new(&handshake);
			let mut stream = RlpStream::new_list(handshake.item_count() * 3);

			for item in handshake.iter() {
				stream.append_raw(item.as_raw(), 1);
				let (mut s1, mut s2) = (RlpStream::new_list(2), RlpStream::new_list(2));
				s1.append(&"foo").append_empty_data();
				s2.append(&"bar").append_empty_data();
				stream.append_raw(&s1.out(), 1);
				stream.append_raw(&s2.out(), 1);
			}

			stream.out()
		};

		let (read_status, read_capabilities, read_flow)
			= parse_handshake(UntrustedRlp::new(&interleaved)).unwrap();

		assert_eq!(read_status, status);
		assert_eq!(read_capabilities, capabilities);
		assert_eq!(read_flow, flow_params);
	}

	#[test]
	fn announcement_roundtrip() {
		let announcement = Announcement {
			head_hash: H256::random(),
			head_num: 100_000,
			head_td: 1_000_000.into(),
			reorg_depth: 4,
			serve_headers: false,
			serve_state_since: Some(99_000),
			serve_chain_since: Some(1),
			tx_relay: true,
		};

		let serialized = write_announcement(&announcement);
		let read = parse_announcement(UntrustedRlp::new(&serialized)).unwrap();

		assert_eq!(read, announcement);
	}

	#[test]
	fn keys_out_of_order() {
		use super::{Key, encode_pair, encode_flag};

		let mut stream = RlpStream::new_list(6);
		stream
			.append(&H256::zero())
			.append(&10u64)
			.append(&100_000u64)
			.append(&2u64)
			.append_raw(&encode_pair(Key::ServeStateSince, &44u64), 1)
			.append_raw(&encode_flag(Key::ServeHeaders), 1);

		let out = stream.drain();
		assert!(parse_announcement(UntrustedRlp::new(&out)).is_err());

		let mut stream = RlpStream::new_list(6);
		stream
			.append(&H256::zero())
			.append(&10u64)
			.append(&100_000u64)
			.append(&2u64)
			.append_raw(&encode_flag(Key::ServeHeaders), 1)
			.append_raw(&encode_pair(Key::ServeStateSince, &44u64), 1);

		let out = stream.drain();
		assert!(parse_announcement(UntrustedRlp::new(&out)).is_ok());
	}
}