// 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::{self, Edge, EdgeType};
use pow::error::{Error, ErrorKind};
use pow::num::ToPrimitive;
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,
{
	edge_bits: u8,
	siphash_keys: [u64; 4],
	easiness: T,
	proof_size: usize,
	edge_mask: 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,
		easiness_pct: u32,
		max_sols: u32,
	) -> Result<Box<Self>, Error> {
		Ok(Box::new(CuckooContext::<T>::new_impl(
			edge_bits,
			proof_size,
			easiness_pct,
			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,
		easiness_pct: u32,
		max_sols: u32,
	) -> Result<CuckooContext<T>, Error> {
		let num_edges = 1 << edge_bits;
		let easiness = to_u64!(easiness_pct) * num_edges / 100;
		Ok(CuckooContext {
			edge_bits: edge_bits,
			siphash_keys: [0; 4],
			easiness: to_edge!(easiness),
			proof_size: proof_size,
			edge_mask: to_edge!(num_edges / 2 - 1),
			graph: vec![T::zero(); num_edges as usize + 1],
			_max_sols: max_sols,
		})
	}

	fn reset(&mut self) -> Result<(), Error> {
		let num_edges = 1 << self.edge_bits;
		self.graph = vec![T::zero(); num_edges + 1];
		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.siphash_keys = common::set_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> {
		common::sipnode::<T>(&self.siphash_keys, edge, &self.edge_mask, 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.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.proof_size];
		for nonce in 0..to_usize!(self.easiness) {
			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.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..to_usize!(self.easiness) {
			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.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 easiness, 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 easiness = to_u64!(self.easiness);
		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] >= easiness || (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] = [
		0x3bbd, 0x4e96, 0x1013b, 0x1172b, 0x1371b, 0x13e6a, 0x1aaa6, 0x1b575, 0x1e237, 0x1ee88,
		0x22f94, 0x24223, 0x25b4f, 0x2e9f3, 0x33b49, 0x34063, 0x3454a, 0x3c081, 0x3d08e, 0x3d863,
		0x4285a, 0x42f22, 0x43122, 0x4b853, 0x4cd0c, 0x4f280, 0x557d5, 0x562cf, 0x58e59, 0x59a62,
		0x5b568, 0x644b9, 0x657e9, 0x66337, 0x6821c, 0x7866f, 0x7e14b, 0x7ec7c, 0x7eed7, 0x80643,
		0x8628c, 0x8949e,
	];
	static V2: [u64; 42] = [
		0x5e3a, 0x8a8b, 0x103d8, 0x1374b, 0x14780, 0x16110, 0x1b571, 0x1c351, 0x1c826, 0x28228,
		0x2909f, 0x29516, 0x2c1c4, 0x334eb, 0x34cdd, 0x38a2c, 0x3ad23, 0x45ac5, 0x46afe, 0x50f43,
		0x51ed6, 0x52ddd, 0x54a82, 0x5a46b, 0x5dbdb, 0x60f6f, 0x60fcd, 0x61c78, 0x63899, 0x64dab,
		0x6affc, 0x6b569, 0x72639, 0x73987, 0x78806, 0x7b98e, 0x7c7d7, 0x7ddd4, 0x7fa88, 0x8277c,
		0x832d9, 0x8ba6f,
	];
	static V3: [u64; 42] = [
		0x308b, 0x9004, 0x91fc, 0x983e, 0x9d67, 0xa293, 0xb4cb, 0xb6c8, 0xccc8, 0xdddc, 0xf04d,
		0x1372f, 0x16ec9, 0x17b61, 0x17d03, 0x1e3bc, 0x1fb0f, 0x29e6e, 0x2a2ca, 0x2a719, 0x3a078,
		0x3b7cc, 0x3c71d, 0x40daa, 0x43e17, 0x46adc, 0x4b359, 0x4c3aa, 0x4ce92, 0x4d06e, 0x51140,
		0x565ac, 0x56b1f, 0x58a8b, 0x5e410, 0x5e607, 0x5ebb5, 0x5f8ae, 0x7aeac, 0x7b902, 0x7d6af,
		0x7f400,
	];
	// 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 mut cuckoo_ctx = CuckooContext::<T>::new(20, 42, 75, 10)?;
		cuckoo_ctx.set_header_nonce([49].to_vec(), None, 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, 70, 10)?;
		cuckoo_ctx.set_header_nonce([50].to_vec(), None, 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([51].to_vec(), None, 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 mut cuckoo_ctx = CuckooContext::<T>::new(20, 42, 75, 10)?;
		cuckoo_ctx.set_header_nonce([49].to_vec(), None, false)?;
		assert!(cuckoo_ctx.verify(&Proof::new(V1.to_vec().clone())).is_ok());
		let mut cuckoo_ctx = CuckooContext::<T>::new(20, 42, 70, 10)?;
		cuckoo_ctx.set_header_nonce([50].to_vec(), None, false)?;
		assert!(cuckoo_ctx.verify(&Proof::new(V2.to_vec().clone())).is_ok());
		cuckoo_ctx.set_header_nonce([51].to_vec(), None, 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, 75, 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, 50, 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,
	{
		for n in 1..5 {
			let h = [n; 32];
			let mut cuckoo_ctx = CuckooContext::<T>::new(16, 42, 75, 10)?;
			cuckoo_ctx.set_header_nonce(h.to_vec(), None, false)?;
			let res = cuckoo_ctx.find_cycles()?;
			assert!(cuckoo_ctx.verify(&res[0]).is_ok())
		}
		Ok(())
	}
}