// Copyright 2018 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Implementation of Cuckoo Cycle designed by John Tromp. Ported to Rust from
//! the C and Java code at https://github.com/tromp/cuckoo. Note that only the
//! simple miner is included, mostly for testing purposes. John Tromp's Tomato
//! miner will be much faster in almost every environment.

use pow::common::{CuckooParams, Edge, EdgeType};
use pow::error::{Error, ErrorKind};
use pow::{PoWContext, Proof};

use std::cmp;
use std::collections::HashSet;

const MAXPATHLEN: usize = 8192;

/// Cuckoo cycle context
pub struct CuckooContext<T>
where
	T: EdgeType,
{
	params: CuckooParams<T>,
	graph: Vec<T>,
	_max_sols: u32,
}

impl<T> PoWContext<T> for CuckooContext<T>
where
	T: EdgeType,
{
	fn new(
		edge_bits: u8,
		proof_size: usize,
		max_sols: u32,
	) -> Result<Box<Self>, Error> {
		Ok(Box::new(CuckooContext::<T>::new_impl(
			edge_bits,
			proof_size,
			max_sols,
		)?))
	}

	fn set_header_nonce(
		&mut self,
		header: Vec<u8>,
		nonce: Option<u32>,
		solve: bool,
	) -> Result<(), Error> {
		self.set_header_nonce_impl(header, nonce, solve)
	}

	fn find_cycles(&mut self) -> Result<Vec<Proof>, Error> {
		self.find_cycles_impl()
	}

	fn verify(&self, proof: &Proof) -> Result<(), Error> {
		self.verify_impl(proof)
	}
}

impl<T> CuckooContext<T>
where
	T: EdgeType,
{
	/// Create a new cuckoo context with given parameters
	pub fn new_impl(
		edge_bits: u8,
		proof_size: usize,
		max_sols: u32,
	) -> Result<CuckooContext<T>, Error> {
		let params = CuckooParams::new(edge_bits, proof_size)?;
		let num_nodes = 2 * params.num_edges as usize;
		Ok(CuckooContext {
			params: params,
			graph: vec![T::zero(); num_nodes],
			_max_sols: max_sols,
		})
	}

	fn reset(&mut self) -> Result<(), Error> {
		let num_nodes = 2 * self.params.num_edges as usize;
		self.graph = vec![T::zero(); num_nodes];
		Ok(())
	}

	/// Set the header and optional nonce in the last part of the header
	/// and create siphash keys
	pub fn set_header_nonce_impl(
		&mut self,
		header: Vec<u8>,
		nonce: Option<u32>,
		solve: bool,
	) -> Result<(), Error> {
		self.params.reset_header_nonce(header, nonce)?;
		if solve {
			self.reset()?;
		}
		Ok(())
	}

	/// Generates a node in the cuckoo graph generated from our seed. A node is
	/// simply materialized as a u64 from a nonce and an offset (generally 0 or
	/// 1).
	fn new_node(&self, edge: T, uorv: u64) -> Result<T, Error> {
		self.params.sipnode(edge, uorv, true)
	}

	/// Creates a new edge in the cuckoo graph generated by our seed from a
	/// nonce. Generates two node coordinates from the nonce and links them
	/// together.
	fn new_edge(&self, nonce: T) -> Result<Edge<T>, Error> {
		Ok(Edge {
			u: self.new_node(nonce, 0)?,
			v: self.new_node(nonce, 1)?,
		})
	}

	fn path(&self, mut u: T, us: &mut [T]) -> Result<T, Error> {
		let mut nu = 0;
		while u != T::zero() {
			nu += 1;
			if nu >= MAXPATHLEN {
				while nu != 0 && us[(nu - 1) as usize] != u {
					nu -= 1;
				}
				return Err(ErrorKind::Path)?;
			}
			us[nu as usize] = u;
			u = self.graph[to_usize!(u)];
		}
		Ok(to_edge!(nu))
	}

	fn update_graph(
		&mut self,
		mut nu: usize,
		us: &[T],
		mut nv: usize,
		vs: &[T],
	) -> Result<(), Error> {
		if nu < nv {
			while nu != 0 {
				nu -= 1;
				self.graph[to_usize!(us[nu + 1])] = us[nu];
			}
			self.graph[to_usize!(us[0])] = vs[0];
		} else {
			while nv != 0 {
				nv -= 1;
				self.graph[to_usize!(vs[nv + 1])] = vs[nv];
			}
			self.graph[to_usize!(vs[0])] = us[0];
		}
		Ok(())
	}

	fn find_sol(
		&mut self,
		mut nu: usize,
		us: &[T],
		mut nv: usize,
		vs: &[T],
	) -> Result<Vec<T>, Error> {
		if us[nu] == vs[nv] {
			let min = cmp::min(nu, nv);
			nu -= min;
			nv -= min;
			while us[nu] != vs[nv] {
				nu += 1;
				nv += 1;
			}
			if nu + nv + 1 == self.params.proof_size {
				self.solution(&us, nu as u64, &vs, nv as u64)
			} else {
				Err(ErrorKind::InvalidCycle(nu + nv + 1))?
			}
		} else {
			Err(ErrorKind::NoCycle)?
		}
	}

	fn solution(&mut self, us: &[T], mut nu: u64, vs: &[T], mut nv: u64) -> Result<Vec<T>, Error> {
		let mut cycle = HashSet::new();
		cycle.insert(Edge { u: us[0], v: vs[0] });
		while nu != 0 {
			// u's in even position; v's in odd
			nu = nu - 1;
			cycle.insert(Edge {
				u: us[((nu + 1) & !1) as usize],
				v: us[(nu | 1) as usize],
			});
		}
		while nv != 0 {
			// u's in odd position; v's in even
			nv -= 1;
			cycle.insert(Edge {
				u: vs[(nv | 1) as usize],
				v: vs[((nv + 1) & !1) as usize],
			});
		}
		let mut n = 0;
		let mut sol = vec![T::zero(); self.params.proof_size];
		for nonce in 0..self.params.num_edges {
			let edge = self.new_edge(to_edge!(nonce))?;
			if cycle.contains(&edge) {
				sol[n] = to_edge!(nonce);
				n += 1;
				cycle.remove(&edge);
			}
		}
		return if n == self.params.proof_size {
			Ok(sol)
		} else {
			Err(ErrorKind::NoCycle)?
		};
	}

	/// Searches for a solution (simple implementation)
	pub fn find_cycles_impl(&mut self) -> Result<Vec<Proof>, Error> {
		let mut us = [T::zero(); MAXPATHLEN];
		let mut vs = [T::zero(); MAXPATHLEN];
		for nonce in 0..self.params.num_edges {
			us[0] = self.new_node(to_edge!(nonce), 0)?;
			vs[0] = self.new_node(to_edge!(nonce), 1)?;
			let u = self.graph[to_usize!(us[0])];
			let v = self.graph[to_usize!(vs[0])];
			if us[0] == T::zero() {
				continue; // ignore duplicate edges
			}
			let nu = to_usize!(self.path(u, &mut us)?);
			let nv = to_usize!(self.path(v, &mut vs)?);

			let sol = self.find_sol(nu, &us, nv, &vs);
			match sol {
				Ok(s) => {
					let mut proof = Proof::new(map_vec!(s.to_vec(), |&n| n.to_u64().unwrap_or(0)));
					proof.cuckoo_sizeshift = self.params.edge_bits;
					return Ok(vec![proof]);
				}
				Err(e) => match e.kind() {
					ErrorKind::InvalidCycle(_) => continue,
					ErrorKind::NoCycle => self.update_graph(nu, &us, nv, &vs)?,
					_ => return Err(e),
				},
			}
		}
		Err(ErrorKind::NoSolution)?
	}

	/// Assuming increasing nonces all smaller than #edges, verifies the
	/// nonces form a cycle in a Cuckoo graph. Each nonce generates an edge, we
	/// build the nodes on both side of that edge and count the connections.
	pub fn verify_impl(&self, proof: &Proof) -> Result<(), Error> {
		let num_nonces = self.params.num_edges;
		let nonces = &proof.nonces;
		let mut us = vec![T::zero(); proof.proof_size()];
		let mut vs = vec![T::zero(); proof.proof_size()];
		for n in 0..proof.proof_size() {
			if nonces[n] >= num_nonces || (n != 0 && nonces[n] <= nonces[n - 1]) {
				return Err(ErrorKind::Verification("edge wrong size".to_owned()))?;
			}
			us[n] = self.new_node(to_edge!(nonces[n]), 0)?;
			vs[n] = self.new_node(to_edge!(nonces[n]), 1)?;
		}
		let mut i = 0;
		let mut count = proof.proof_size();
		loop {
			let mut j = i;
			for k in 0..proof.proof_size() {
				// find unique other j with same vs[j]
				if k != i && vs[k] == vs[i] {
					if j != i {
						return Err(ErrorKind::Verification("".to_owned()))?;
					}
					j = k;
				}
			}
			if j == i {
				return Err(ErrorKind::Verification("".to_owned()))?;
			}
			i = j;
			for k in 0..proof.proof_size() {
				// find unique other i with same us[i]
				if k != j && us[k] == us[j] {
					if i != j {
						return Err(ErrorKind::Verification("".to_owned()))?;
					}
					i = k;
				}
			}
			if i == j {
				return Err(ErrorKind::Verification("".to_owned()))?;
			}
			count -= 2;
			if i == 0 {
				break;
			}
		}
		match count == 0 {
			true => Ok(()),
			false => Err(ErrorKind::Verification("Invalid solution".to_owned()))?,
		}
	}
}

#[cfg(test)]
mod test {
	use super::*;

	static V1: [u64; 42] = [
		0x8702, 0x12003, 0x2043f, 0x24cf8, 0x27631, 0x2beda, 0x325e5, 0x345b4, 0x36f5c, 0x3b3bc,
		0x4cef6, 0x4dfdf, 0x5036b, 0x5d528, 0x7d76b, 0x80958, 0x81649, 0x8a064, 0x935fe, 0x93c28,
		0x93fc9, 0x9aec5, 0x9c5c8, 0xa00a7, 0xa7256, 0xaa35e, 0xb9e04, 0xc8835, 0xcda49, 0xd72ea,
		0xd7f80, 0xdaa3a, 0xdafce, 0xe03fe, 0xe55a2, 0xe6e60, 0xebb9d, 0xf5248, 0xf6a4b, 0xf6d32,
		0xf7c61, 0xfd9e9
	];
	static V2: [u64; 42] = [
		0xab0, 0x403c, 0x509c, 0x127c0, 0x1a0b3, 0x1ffe4, 0x26180, 0x2a20a, 0x35559, 0x36dd3,
		0x3cb20, 0x4992f, 0x55b20, 0x5b507, 0x66e58, 0x6784d, 0x6fda8, 0x7363d, 0x76dd6, 0x7f13b,
		0x84672, 0x85724, 0x991cf, 0x9a6fe, 0x9b0c5, 0xa5019, 0xa7207, 0xaf32f, 0xc29f3, 0xc39d3,
		0xc78ed, 0xc9e75, 0xcd0db, 0xcd81e, 0xd02e0, 0xd05c4, 0xd8f99, 0xd9359, 0xdff3b, 0xea623,
		0xf9100, 0xfc966
	];
	static V3: [u64; 42] = [
		0x14ca, 0x1e80, 0x587c, 0xa2d4, 0x14f6b, 0x1b100, 0x1b74c, 0x2477d, 0x29ba4, 0x33f25,
		0x4c55f, 0x4d280, 0x50ffa, 0x53900, 0x5cf62, 0x63f66, 0x65623, 0x6fb19, 0x7a19e, 0x82eef,
		0x83d2d, 0x88015, 0x8e6c5, 0x91086, 0x97429, 0x9aa27, 0xa01b7, 0xa304b, 0xafa06, 0xb1cb3,
		0xbb9fc, 0xbf345, 0xc0761, 0xc0e78, 0xc5b99, 0xc9f09, 0xcc62c, 0xceb6e, 0xd98ad, 0xeecb3,
		0xef966, 0xfef9b
	];
	// cuckoo28 at 50% edges of letter 'u'
	static V4: [u64; 42] = [
		0xf7243, 0x11f130, 0x193812, 0x23b565, 0x279ac3, 0x69b270, 0xe0778f, 0xef51fc, 0x10bf6e8,
		0x13ccf7d, 0x1551177, 0x1b6cfd2, 0x1f872c3, 0x2075681, 0x2e23ccc, 0x2e4c0aa, 0x2f607f1,
		0x3007eeb, 0x3407e9a, 0x35423f9, 0x39e48bf, 0x45e3bf6, 0x46aa484, 0x47c0fe1, 0x4b1d5a6,
		0x4bae0ba, 0x4dfdbaf, 0x5048eda, 0x537da6b, 0x5402887, 0x56b8897, 0x5bd8e8b, 0x622de20,
		0x62be5ce, 0x62d538e, 0x6464518, 0x650a6d5, 0x66ec4fa, 0x66f9476, 0x6b1e5f6, 0x6fd5d88,
		0x701f37b,
	];

	#[test]
	fn cuckoo_context() {
		let ret = mine20_vectors::<u32>();
		if let Err(r) = ret {
			panic!("mine20_vectors u32: Error: {}", r);
		}
		let ret = mine20_vectors::<u64>();
		if let Err(r) = ret {
			panic!("mine20_vectors u64: Error: {}", r);
		}
		let ret = validate20_vectors::<u32>();
		if let Err(r) = ret {
			panic!("validate20_vectors u32: Error: {}", r);
		}
		let ret = validate20_vectors::<u64>();
		if let Err(r) = ret {
			panic!("validate20_vectors u64: Error: {}", r);
		}
		let ret = validate_fail::<u32>();
		if let Err(r) = ret {
			panic!("validate_fail u32: Error: {}", r);
		}
		let ret = validate_fail::<u64>();
		if let Err(r) = ret {
			panic!("validate_fail u64: Error: {}", r);
		}
		let ret = mine16_validate::<u32>();
		if let Err(r) = ret {
			panic!("mine16_validate u32: Error: {}", r);
		}
		let ret = mine16_validate::<u64>();
		if let Err(r) = ret {
			panic!("mine16_validate u64: Error: {}", r);
		}
	}

	/// Find a 42-cycle on Cuckoo20 at 75% easiness and verify against a few
	/// known cycle proofs
	/// generated by other implementations.
	fn mine20_vectors<T>() -> Result<(), Error>
	where
		T: EdgeType,
	{
		let header = [0; 4].to_vec();
		let mut cuckoo_ctx = CuckooContext::<T>::new(20, 42, 10)?;
		cuckoo_ctx.set_header_nonce(header.clone(), Some(39), true)?;
		let res = cuckoo_ctx.find_cycles()?;
		let mut proof = Proof::new(V1.to_vec());
		proof.cuckoo_sizeshift = 20;
		assert_eq!(proof, res[0]);

		let mut cuckoo_ctx = CuckooContext::<T>::new(20, 42, 10)?;
		cuckoo_ctx.set_header_nonce(header.clone(), Some(56), true)?;
		let res = cuckoo_ctx.find_cycles()?;
		let mut proof = Proof::new(V2.to_vec());
		proof.cuckoo_sizeshift = 20;
		assert_eq!(proof, res[0]);

		//re-use context
		cuckoo_ctx.set_header_nonce(header, Some(66), true)?;
		let res = cuckoo_ctx.find_cycles()?;
		let mut proof = Proof::new(V3.to_vec());
		proof.cuckoo_sizeshift = 20;
		assert_eq!(proof, res[0]);
		Ok(())
	}

	fn validate20_vectors<T>() -> Result<(), Error>
	where
		T: EdgeType,
	{
		let header = [0; 4].to_vec();
		let mut cuckoo_ctx = CuckooContext::<T>::new(20, 42, 10)?;
		cuckoo_ctx.set_header_nonce(header.clone(), Some(39), false)?;
		assert!(cuckoo_ctx.verify(&Proof::new(V1.to_vec().clone())).is_ok());
		let mut cuckoo_ctx = CuckooContext::<T>::new(20, 42, 10)?;
		cuckoo_ctx.set_header_nonce(header.clone(), Some(56), false)?;
		assert!(cuckoo_ctx.verify(&Proof::new(V2.to_vec().clone())).is_ok());
		cuckoo_ctx.set_header_nonce(header.clone(), Some(66), false)?;
		assert!(cuckoo_ctx.verify(&Proof::new(V3.to_vec().clone())).is_ok());
		Ok(())
	}

	fn validate_fail<T>() -> Result<(), Error>
	where
		T: EdgeType,
	{
		// edge checks
		let mut cuckoo_ctx = CuckooContext::<T>::new(20, 42, 10)?;
		cuckoo_ctx.set_header_nonce([49].to_vec(), None, false)?;
		// edge checks
		assert!(!cuckoo_ctx.verify(&Proof::new(vec![0; 42])).is_ok());
		assert!(!cuckoo_ctx.verify(&Proof::new(vec![0xffff; 42])).is_ok());
		// wrong data for proof
		cuckoo_ctx.set_header_nonce([50].to_vec(), None, false)?;
		assert!(!cuckoo_ctx.verify(&Proof::new(V1.to_vec().clone())).is_ok());
		let mut test_header = [0; 32];
		test_header[0] = 24;
		let mut cuckoo_ctx = CuckooContext::<T>::new(20, 42, 10)?;
		cuckoo_ctx.set_header_nonce(test_header.to_vec(), None, false)?;
		assert!(!cuckoo_ctx.verify(&Proof::new(V4.to_vec().clone())).is_ok());
		Ok(())
	}

	fn mine16_validate<T>() -> Result<(), Error>
	where
		T: EdgeType,
	{
		let h = [0 as u8; 32];
		for n in [45 as u32, 49,131,143,151].iter() {
			let mut cuckoo_ctx = CuckooContext::<T>::new(16, 42, 10)?;
			cuckoo_ctx.set_header_nonce(h.to_vec(), Some(*n), false)?;
			let res = cuckoo_ctx.find_cycles()?;
			assert!(cuckoo_ctx.verify(&res[0]).is_ok())
		}
		Ok(())
	}
}