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

//! Tests for the `LightProtocol` implementation.
//! These don't test of the higher level logic on top of 

use ethcore::blockchain_info::BlockChainInfo;
use ethcore::client::{BlockChainClient, EachBlockWith, TestBlockChainClient};
use ethcore::ids::BlockID;
use ethcore::transaction::SignedTransaction;
use network::PeerId;

use net::buffer_flow::FlowParams;
use net::context::IoContext;
use net::status::{Capabilities, Status, NetworkId, write_handshake};
use net::{encode_request, LightProtocol, Params, packet};
use provider::Provider;
use request::{self, Request, Headers};

use rlp::*;
use util::{Bytes, H256};

use std::sync::Arc;

// expected result from a call.
#[derive(Debug, PartialEq, Eq)]
enum Expect {
    /// Expect to have message sent to peer.
    Send(PeerId, u8, Vec<u8>),
    /// Expect this response.
    Respond(u8, Vec<u8>),
    /// Expect a punishment (disconnect/disable)
    Punish(PeerId),
    /// Expect nothing.
    Nothing,
}

impl IoContext for Expect {
    fn send(&self, peer: PeerId, packet_id: u8, packet_body: Vec<u8>) {
        assert_eq!(self, &Expect::Send(peer, packet_id, packet_body));
    }

    fn respond(&self, packet_id: u8, packet_body: Vec<u8>) {
        assert_eq!(self, &Expect::Respond(packet_id, packet_body));
    }

    fn disconnect_peer(&self, peer: PeerId) {
        assert_eq!(self, &Expect::Punish(peer));
    }

    fn disable_peer(&self, peer: PeerId) {
        assert_eq!(self, &Expect::Punish(peer));
    }
}

// can't implement directly for Arc due to cross-crate orphan rules.
struct TestProvider(Arc<TestProviderInner>);

struct TestProviderInner {
    client: TestBlockChainClient,
}

impl Provider for TestProvider {
    fn chain_info(&self) -> BlockChainInfo {
        self.0.client.chain_info()
    }

    fn reorg_depth(&self, a: &H256, b: &H256) -> Option<u64> {
        self.0.client.tree_route(a, b).map(|route| route.index as u64)
    }

    fn earliest_state(&self) -> Option<u64> {
        None
    }

    fn block_headers(&self, req: request::Headers) -> Vec<Bytes> {
		let best_num = self.0.client.chain_info().best_block_number;
		let start_num = req.block_num;

		match self.0.client.block_hash(BlockID::Number(req.block_num)) {
			Some(hash) if hash == req.block_hash => {}
			_=> {
				trace!(target: "les_provider", "unknown/non-canonical start block in header request: {:?}", (req.block_num, req.block_hash));
				return vec![]
			}
		}

		(0u64..req.max as u64)
			.map(|x: u64| x.saturating_mul(req.skip + 1))
			.take_while(|x| if req.reverse { x < &start_num } else { best_num - start_num >= *x })
			.map(|x| if req.reverse { start_num - x } else { start_num + x })
			.map(|x| self.0.client.block_header(BlockID::Number(x))) 
			.take_while(|x| x.is_some())
			.flat_map(|x| x)
			.collect()
	}

	fn block_bodies(&self, req: request::Bodies) -> Vec<Bytes> {
		req.block_hashes.into_iter()
			.map(|hash| self.0.client.block_body(BlockID::Hash(hash)))
			.map(|body| body.unwrap_or_else(|| ::rlp::EMPTY_LIST_RLP.to_vec()))
			.collect()
	}

	fn receipts(&self, req: request::Receipts) -> Vec<Bytes> {
		req.block_hashes.into_iter()
			.map(|hash| self.0.client.block_receipts(&hash))
			.map(|receipts| receipts.unwrap_or_else(|| ::rlp::EMPTY_LIST_RLP.to_vec()))
			.collect()
	}

	fn proofs(&self, req: request::StateProofs) -> Vec<Bytes> {
		req.requests.into_iter()
            .map(|req| {
                match req.key2 {
                    Some(key2) => ::util::sha3::SHA3_NULL_RLP.to_owned(),
                    None => {
                        // sort of a leaf node
                        let mut stream = RlpStream::new_list(2);
                        stream.append(&req.key1).append_empty_data();
                        stream.out()
                    }
                }
            })
            .collect()		
	}

	fn contract_code(&self, req: request::ContractCodes) -> Vec<Bytes> {
		req.code_requests.into_iter().map(|_| Vec::new()).collect()
	}

	fn header_proofs(&self, req: request::HeaderProofs) -> Vec<Bytes> {
		req.requests.into_iter().map(|_| ::rlp::EMPTY_LIST_RLP.to_vec()).collect()
	}

	fn pending_transactions(&self) -> Vec<SignedTransaction> {
		self.0.client.pending_transactions()
	}
}

fn make_flow_params() -> FlowParams {
    FlowParams::new(5_000_000.into(), Default::default(), 100_000.into())
}

fn capabilities() -> Capabilities {
    Capabilities {
        serve_headers: true,
        serve_chain_since: Some(1),
        serve_state_since: Some(1),
        tx_relay: true,
    }
}

// helper for setting up the protocol handler and provider.
fn setup(flow_params: FlowParams, capabilities: Capabilities) -> (Arc<TestProviderInner>, LightProtocol) {
    let provider = Arc::new(TestProviderInner {
        client: TestBlockChainClient::new(),
    });

    let proto = LightProtocol::new(Box::new(TestProvider(provider.clone())), Params {
        network_id: NetworkId::Testnet,
        flow_params: flow_params,
        capabilities: capabilities,
    });

    (provider, proto)
}

fn status(chain_info: BlockChainInfo) -> Status {
    Status {
        protocol_version: ::net::PROTOCOL_VERSION,
        network_id: NetworkId::Testnet,
        head_td: chain_info.total_difficulty,
        head_hash: chain_info.best_block_hash,
        head_num: chain_info.best_block_number,
        genesis_hash: chain_info.genesis_hash,
        last_head: None,
    }
}

#[test]
fn handshake_expected() {
    let flow_params = make_flow_params();
    let capabilities = capabilities();    

    let (provider, proto) = setup(flow_params.clone(), capabilities.clone());    

    let status = status(provider.client.chain_info());

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

    proto.on_connect(&1, &Expect::Send(1, packet::STATUS, packet_body));
}

#[test]
#[should_panic]
fn genesis_mismatch() {
    let flow_params = make_flow_params();
    let capabilities = capabilities();    

    let (provider, proto) = setup(flow_params.clone(), capabilities.clone());    

    let mut status = status(provider.client.chain_info());
    status.genesis_hash = H256::default();

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

    proto.on_connect(&1, &Expect::Send(1, packet::STATUS, packet_body));
}

#[test]
fn buffer_overflow() {
    let flow_params = make_flow_params();
    let capabilities = capabilities();    

    let (provider, proto) = setup(flow_params.clone(), capabilities.clone());    

    let status = status(provider.client.chain_info());

    {
        let packet_body = write_handshake(&status, &capabilities, &flow_params);    
        proto.on_connect(&1, &Expect::Send(1, packet::STATUS, packet_body));    
    }

    {
        let my_status = write_handshake(&status, &capabilities, &flow_params);
        proto.handle_packet(&Expect::Nothing, &1, packet::STATUS, &my_status);
    }

    // 1000 requests is far too many for the default flow params.
    let request = encode_request(&Request::Headers(Headers {
        block_num: 1,
        block_hash: provider.client.chain_info().genesis_hash,
        max: 1000,
        skip: 0,
        reverse: false,
    }), 111);

    proto.handle_packet(&Expect::Punish(1), &1, packet::GET_BLOCK_HEADERS, &request);
}

// test the basic request types -- these just make sure that requests are parsed
// and sent to the provider correctly as well as testing response formatting.

#[test]
fn get_block_headers() {
    let flow_params = FlowParams::new(5_000_000.into(), Default::default(), 0.into());
    let capabilities = capabilities();    

    let (provider, proto) = setup(flow_params.clone(), capabilities.clone());    

    let cur_status = status(provider.client.chain_info());
    let my_status = write_handshake(&cur_status, &capabilities, &flow_params);

    provider.client.add_blocks(100, EachBlockWith::Nothing);

    let cur_status = status(provider.client.chain_info());

    {
        let packet_body = write_handshake(&cur_status, &capabilities, &flow_params);    
        proto.on_connect(&1, &Expect::Send(1, packet::STATUS, packet_body));    
        proto.handle_packet(&Expect::Nothing, &1, packet::STATUS, &my_status);
    }

    let request = Headers {
        block_num: 1,
        block_hash: provider.client.block_hash(BlockID::Number(1)).unwrap(),
        max: 10,
        skip: 0,
        reverse: false,
    };
    let req_id = 111;

    let request_body = encode_request(&Request::Headers(request.clone()), req_id);
    let response = {
        let headers: Vec<_> = (0..10).map(|i| provider.client.block_header(BlockID::Number(i + 1)).unwrap()).collect();
        assert_eq!(headers.len(), 10);

        let new_buf = *flow_params.limit() - flow_params.compute_cost(request::Kind::Headers, 10);

        let mut response_stream = RlpStream::new_list(12);
        
        response_stream.append(&req_id).append(&new_buf);
        for header in headers {
            response_stream.append_raw(&header, 1);
        }

        response_stream.out()
    };

    let expected = Expect::Respond(packet::BLOCK_HEADERS, response);
    proto.handle_packet(&expected, &1, packet::GET_BLOCK_HEADERS, &request_body);
}

#[test]
fn get_block_bodies() {
    let flow_params = FlowParams::new(5_000_000.into(), Default::default(), 0.into());
    let capabilities = capabilities();    

    let (provider, proto) = setup(flow_params.clone(), capabilities.clone());    

    let cur_status = status(provider.client.chain_info());
    let my_status = write_handshake(&cur_status, &capabilities, &flow_params);

    provider.client.add_blocks(100, EachBlockWith::Nothing);

    let cur_status = status(provider.client.chain_info());

    {
        let packet_body = write_handshake(&cur_status, &capabilities, &flow_params);    
        proto.on_connect(&1, &Expect::Send(1, packet::STATUS, packet_body));    
        proto.handle_packet(&Expect::Nothing, &1, packet::STATUS, &my_status);
    }

    let request = request::Bodies {
        block_hashes: (0..10).map(|i| provider.client.block_hash(BlockID::Number(i)).unwrap()).collect(),
    };

    let req_id = 111;

    let request_body = encode_request(&Request::Bodies(request.clone()), req_id);
    let response = {
        let bodies: Vec<_> = (0..10).map(|i| provider.client.block_body(BlockID::Number(i + 1)).unwrap()).collect();
        assert_eq!(bodies.len(), 10);

        let new_buf = *flow_params.limit() - flow_params.compute_cost(request::Kind::Bodies, 10);

        let mut response_stream = RlpStream::new_list(12);
        
        response_stream.append(&req_id).append(&new_buf);
        for body in bodies {
            response_stream.append_raw(&body, 1);
        }

        response_stream.out()
    };

    let expected = Expect::Respond(packet::BLOCK_BODIES, response);
    proto.handle_packet(&expected, &1, packet::GET_BLOCK_BODIES, &request_body);
}

#[test]
fn get_block_receipts() {
    let flow_params = FlowParams::new(5_000_000.into(), Default::default(), 0.into());
    let capabilities = capabilities();    

    let (provider, proto) = setup(flow_params.clone(), capabilities.clone());    

    let cur_status = status(provider.client.chain_info());
    let my_status = write_handshake(&cur_status, &capabilities, &flow_params);

    provider.client.add_blocks(1000, EachBlockWith::Nothing);

    let cur_status = status(provider.client.chain_info());

    {
        let packet_body = write_handshake(&cur_status, &capabilities, &flow_params);    
        proto.on_connect(&1, &Expect::Send(1, packet::STATUS, packet_body));    
        proto.handle_packet(&Expect::Nothing, &1, packet::STATUS, &my_status);
    }

    // find the first 10 block hashes starting with `f` because receipts are only provided
    // by the test client in that case.
    let block_hashes: Vec<_> = (0..1000).map(|i| provider.client.block_hash(BlockID::Number(i)).unwrap())
        .filter(|hash| format!("{}", hash).starts_with("f")).take(10).collect();

    let request = request::Receipts {
        block_hashes: block_hashes.clone(),
    };

    let req_id = 111;

    let request_body = encode_request(&Request::Receipts(request.clone()), req_id);
    let response = {
        let receipts: Vec<_> = block_hashes.iter()
            .map(|hash| provider.client.block_receipts(hash).unwrap())
            .collect();

        let new_buf = *flow_params.limit() - flow_params.compute_cost(request::Kind::Receipts, receipts.len());

        let mut response_stream = RlpStream::new_list(2 + receipts.len());
        
        response_stream.append(&req_id).append(&new_buf);
        for block_receipts in receipts {
            response_stream.append_raw(&block_receipts, 1);
        }

        response_stream.out()
    };

    let expected = Expect::Respond(packet::RECEIPTS, response);
    proto.handle_packet(&expected, &1, packet::GET_RECEIPTS, &request_body);
}

#[test]
fn get_block_bodies() {
    let flow_params = FlowParams::new(5_000_000.into(), Default::default(), 0.into());
    let capabilities = capabilities();    

    let (provider, proto) = setup(flow_params.clone(), capabilities.clone());    

    let cur_status = status(provider.client.chain_info());
    let my_status = write_handshake(&cur_status, &capabilities, &flow_params);

    provider.client.add_blocks(100, EachBlockWith::Nothing);

    let cur_status = status(provider.client.chain_info());

    {
        let packet_body = write_handshake(&cur_status, &capabilities, &flow_params);    
        proto.on_connect(&1, &Expect::Send(1, packet::STATUS, packet_body));    
        proto.handle_packet(&Expect::Nothing, &1, packet::STATUS, &my_status);
    }

    let request = request::Bodies {
        block_hashes: (0..10).map(|i| provider.client.block_hash(BlockID::Number(i)).unwrap()).collect(),
    };

    let req_id = 111;

    let request_body = encode_request(&Request::Bodies(request.clone()), req_id);
    let response = {
        let bodies: Vec<_> = (0..10).map(|i| provider.client.block_body(BlockID::Number(i + 1)).unwrap()).collect();
        assert_eq!(bodies.len(), 10);

        let new_buf = *flow_params.limit() - flow_params.compute_cost(request::Kind::Bodies, 10);

        let mut response_stream = RlpStream::new_list(12);
        
        response_stream.append(&req_id).append(&new_buf);
        for body in bodies {
            response_stream.append_raw(&body, 1);
        }

        response_stream.out()
    };

    let expected = Expect::Respond(packet::BLOCK_BODIES, response);
    proto.handle_packet(&expected, &1, packet::GET_BLOCK_BODIES, &request_body);
}