// This file is part of Substrate.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

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

// This program 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 this program. If not, see <https://www.gnu.org/licenses/>.

//! Warp sync support.

pub use sp_consensus_grandpa::{AuthorityList, SetId};

use crate::{
	schema::v1::{StateRequest, StateResponse},
	state::{ImportResult, StateSync},
};
use codec::{Decode, Encode};
use futures::channel::oneshot;
use log::error;
use sc_client_api::ProofProvider;
use sc_network_common::sync::message::{
	BlockAttributes, BlockData, BlockRequest, Direction, FromBlock,
};
use sp_blockchain::HeaderBackend;
use sp_runtime::traits::{Block as BlockT, Header, NumberFor, Zero};
use std::{fmt, sync::Arc};

/// Log target for this file.
const LOG_TARGET: &'static str = "sync";

/// Scale-encoded warp sync proof response.
pub struct EncodedProof(pub Vec<u8>);

/// Warp sync request
#[derive(Encode, Decode, Debug, Clone)]
pub struct WarpProofRequest<B: BlockT> {
	/// Start collecting proofs from this block.
	pub begin: B::Hash,
}

/// Proof verification result.
pub enum VerificationResult<Block: BlockT> {
	/// Proof is valid, but the target was not reached.
	Partial(SetId, AuthorityList, Block::Hash),
	/// Target finality is proved.
	Complete(SetId, AuthorityList, Block::Header),
}

/// Warp sync backend. Handles retrieving and verifying warp sync proofs.
pub trait WarpSyncProvider<Block: BlockT>: Send + Sync {
	/// Generate proof starting at given block hash. The proof is accumulated until maximum proof
	/// size is reached.
	fn generate(
		&self,
		start: Block::Hash,
	) -> Result<EncodedProof, Box<dyn std::error::Error + Send + Sync>>;
	/// Verify warp proof against current set of authorities.
	fn verify(
		&self,
		proof: &EncodedProof,
		set_id: SetId,
		authorities: AuthorityList,
	) -> Result<VerificationResult<Block>, Box<dyn std::error::Error + Send + Sync>>;
	/// Get current list of authorities. This is supposed to be genesis authorities when starting
	/// sync.
	fn current_authorities(&self) -> AuthorityList;
}

/// Reported warp sync phase.
#[derive(Clone, Eq, PartialEq, Debug)]
pub enum WarpSyncPhase<Block: BlockT> {
	/// Waiting for peers to connect.
	AwaitingPeers { required_peers: usize },
	/// Waiting for target block to be received.
	AwaitingTargetBlock,
	/// Downloading and verifying grandpa warp proofs.
	DownloadingWarpProofs,
	/// Downloading target block.
	DownloadingTargetBlock,
	/// Downloading state data.
	DownloadingState,
	/// Importing state.
	ImportingState,
	/// Downloading block history.
	DownloadingBlocks(NumberFor<Block>),
}

impl<Block: BlockT> fmt::Display for WarpSyncPhase<Block> {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		match self {
			Self::AwaitingPeers { required_peers } =>
				write!(f, "Waiting for {required_peers} peers to be connected"),
			Self::AwaitingTargetBlock => write!(f, "Waiting for target block to be received"),
			Self::DownloadingWarpProofs => write!(f, "Downloading finality proofs"),
			Self::DownloadingTargetBlock => write!(f, "Downloading target block"),
			Self::DownloadingState => write!(f, "Downloading state"),
			Self::ImportingState => write!(f, "Importing state"),
			Self::DownloadingBlocks(n) => write!(f, "Downloading block history (#{})", n),
		}
	}
}

/// Reported warp sync progress.
#[derive(Clone, Eq, PartialEq, Debug)]
pub struct WarpSyncProgress<Block: BlockT> {
	/// Estimated download percentage.
	pub phase: WarpSyncPhase<Block>,
	/// Total bytes downloaded so far.
	pub total_bytes: u64,
}

/// The different types of warp syncing, passed to `build_network`.
pub enum WarpSyncParams<Block: BlockT> {
	/// Standard warp sync for the chain.
	WithProvider(Arc<dyn WarpSyncProvider<Block>>),
	/// Skip downloading proofs and wait for a header of the state that should be downloaded.
	///
	/// It is expected that the header provider ensures that the header is trusted.
	WaitForTarget(oneshot::Receiver<<Block as BlockT>::Header>),
}

/// Warp sync configuration as accepted by [`WarpSync`].
pub enum WarpSyncConfig<Block: BlockT> {
	/// Standard warp sync for the chain.
	WithProvider(Arc<dyn WarpSyncProvider<Block>>),
	/// Skip downloading proofs and wait for a header of the state that should be downloaded.
	///
	/// It is expected that the header provider ensures that the header is trusted.
	WaitForTarget,
}

impl<Block: BlockT> WarpSyncParams<Block> {
	/// Split `WarpSyncParams` into `WarpSyncConfig` and warp sync target block header receiver.
	pub fn split(
		self,
	) -> (WarpSyncConfig<Block>, Option<oneshot::Receiver<<Block as BlockT>::Header>>) {
		match self {
			WarpSyncParams::WithProvider(provider) =>
				(WarpSyncConfig::WithProvider(provider), None),
			WarpSyncParams::WaitForTarget(rx) => (WarpSyncConfig::WaitForTarget, Some(rx)),
		}
	}
}

/// Warp sync phase.
enum Phase<B: BlockT, Client> {
	/// Downloading warp proofs.
	WarpProof {
		set_id: SetId,
		authorities: AuthorityList,
		last_hash: B::Hash,
		warp_sync_provider: Arc<dyn WarpSyncProvider<B>>,
	},
	/// Waiting for target block to be set externally if we skip warp proofs downloading,
	/// and start straight from the target block (used by parachains warp sync).
	PendingTargetBlock,
	/// Downloading target block.
	TargetBlock(B::Header),
	/// Downloading state.
	State(StateSync<B, Client>),
}

/// Import warp proof result.
pub enum WarpProofImportResult {
	/// Import was successful.
	Success,
	/// Bad proof.
	BadResponse,
}

/// Import target block result.
pub enum TargetBlockImportResult {
	/// Import was successful.
	Success,
	/// Invalid block.
	BadResponse,
}

/// Warp sync state machine. Accumulates warp proofs and state.
pub struct WarpSync<B: BlockT, Client> {
	phase: Phase<B, Client>,
	client: Arc<Client>,
	total_proof_bytes: u64,
}

impl<B, Client> WarpSync<B, Client>
where
	B: BlockT,
	Client: HeaderBackend<B> + ProofProvider<B> + 'static,
{
	/// Create a new instance. When passing a warp sync provider we will be checking for proof and
	/// authorities. Alternatively we can pass a target block when we want to skip downloading
	/// proofs, in this case we will continue polling until the target block is known.
	pub fn new(client: Arc<Client>, warp_sync_config: WarpSyncConfig<B>) -> Self {
		let last_hash = client.hash(Zero::zero()).unwrap().expect("Genesis header always exists");
		match warp_sync_config {
			WarpSyncConfig::WithProvider(warp_sync_provider) => {
				let phase = Phase::WarpProof {
					set_id: 0,
					authorities: warp_sync_provider.current_authorities(),
					last_hash,
					warp_sync_provider: warp_sync_provider.clone(),
				};
				Self { client, phase, total_proof_bytes: 0 }
			},
			WarpSyncConfig::WaitForTarget =>
				Self { client, phase: Phase::PendingTargetBlock, total_proof_bytes: 0 },
		}
	}

	/// Set target block externally in case we skip warp proof downloading.
	pub fn set_target_block(&mut self, header: B::Header) {
		let Phase::PendingTargetBlock = self.phase else {
			error!(
				target: LOG_TARGET,
				"Attempt to set warp sync target block in invalid phase.",
			);
			debug_assert!(false);
			return
		};

		self.phase = Phase::TargetBlock(header);
	}

	///  Validate and import a state response.
	pub fn import_state(&mut self, response: StateResponse) -> ImportResult<B> {
		match &mut self.phase {
			Phase::WarpProof { .. } | Phase::TargetBlock(_) | Phase::PendingTargetBlock { .. } => {
				log::debug!(target: "sync", "Unexpected state response");
				ImportResult::BadResponse
			},
			Phase::State(sync) => sync.import(response),
		}
	}

	///  Validate and import a warp proof response.
	pub fn import_warp_proof(&mut self, response: EncodedProof) -> WarpProofImportResult {
		match &mut self.phase {
			Phase::State(_) | Phase::TargetBlock(_) | Phase::PendingTargetBlock { .. } => {
				log::debug!(target: "sync", "Unexpected warp proof response");
				WarpProofImportResult::BadResponse
			},
			Phase::WarpProof { set_id, authorities, last_hash, warp_sync_provider } =>
				match warp_sync_provider.verify(&response, *set_id, authorities.clone()) {
					Err(e) => {
						log::debug!(target: "sync", "Bad warp proof response: {}", e);
						WarpProofImportResult::BadResponse
					},
					Ok(VerificationResult::Partial(new_set_id, new_authorities, new_last_hash)) => {
						log::debug!(target: "sync", "Verified partial proof, set_id={:?}", new_set_id);
						*set_id = new_set_id;
						*authorities = new_authorities;
						*last_hash = new_last_hash;
						self.total_proof_bytes += response.0.len() as u64;
						WarpProofImportResult::Success
					},
					Ok(VerificationResult::Complete(new_set_id, _, header)) => {
						log::debug!(target: "sync", "Verified complete proof, set_id={:?}", new_set_id);
						self.total_proof_bytes += response.0.len() as u64;
						self.phase = Phase::TargetBlock(header);
						WarpProofImportResult::Success
					},
				},
		}
	}

	/// Import the target block body.
	pub fn import_target_block(&mut self, block: BlockData<B>) -> TargetBlockImportResult {
		match &mut self.phase {
			Phase::WarpProof { .. } | Phase::State(_) | Phase::PendingTargetBlock { .. } => {
				log::debug!(target: "sync", "Unexpected target block response");
				TargetBlockImportResult::BadResponse
			},
			Phase::TargetBlock(header) =>
				if let Some(block_header) = &block.header {
					if block_header == header {
						if block.body.is_some() {
							let state_sync = StateSync::new(
								self.client.clone(),
								header.clone(),
								block.body,
								block.justifications,
								false,
							);
							self.phase = Phase::State(state_sync);
							TargetBlockImportResult::Success
						} else {
							log::debug!(
								target: "sync",
								"Importing target block failed: missing body.",
							);
							TargetBlockImportResult::BadResponse
						}
					} else {
						log::debug!(
							target: "sync",
							"Importing target block failed: different header.",
						);
						TargetBlockImportResult::BadResponse
					}
				} else {
					log::debug!(target: "sync", "Importing target block failed: missing header.");
					TargetBlockImportResult::BadResponse
				},
		}
	}

	/// Produce next state request.
	pub fn next_state_request(&self) -> Option<StateRequest> {
		match &self.phase {
			Phase::WarpProof { .. } | Phase::TargetBlock(_) | Phase::PendingTargetBlock { .. } =>
				None,
			Phase::State(sync) => Some(sync.next_request()),
		}
	}

	/// Produce next warp proof request.
	pub fn next_warp_proof_request(&self) -> Option<WarpProofRequest<B>> {
		match &self.phase {
			Phase::WarpProof { last_hash, .. } => Some(WarpProofRequest { begin: *last_hash }),
			Phase::TargetBlock(_) | Phase::State(_) | Phase::PendingTargetBlock { .. } => None,
		}
	}

	/// Produce next target block request.
	pub fn next_target_block_request(&self) -> Option<(NumberFor<B>, BlockRequest<B>)> {
		match &self.phase {
			Phase::WarpProof { .. } | Phase::State(_) | Phase::PendingTargetBlock { .. } => None,
			Phase::TargetBlock(header) => {
				let request = BlockRequest::<B> {
					id: 0,
					fields: BlockAttributes::HEADER |
						BlockAttributes::BODY | BlockAttributes::JUSTIFICATION,
					from: FromBlock::Hash(header.hash()),
					direction: Direction::Ascending,
					max: Some(1),
				};
				Some((*header.number(), request))
			},
		}
	}

	/// Return target block hash if it is known.
	pub fn target_block_hash(&self) -> Option<B::Hash> {
		match &self.phase {
			Phase::WarpProof { .. } | Phase::TargetBlock(_) | Phase::PendingTargetBlock { .. } =>
				None,
			Phase::State(s) => Some(s.target()),
		}
	}

	/// Return target block number if it is known.
	pub fn target_block_number(&self) -> Option<NumberFor<B>> {
		match &self.phase {
			Phase::WarpProof { .. } | Phase::PendingTargetBlock { .. } => None,
			Phase::TargetBlock(header) => Some(*header.number()),
			Phase::State(s) => Some(s.target_block_num()),
		}
	}

	/// Check if the state is complete.
	pub fn is_complete(&self) -> bool {
		match &self.phase {
			Phase::WarpProof { .. } | Phase::TargetBlock(_) | Phase::PendingTargetBlock { .. } =>
				false,
			Phase::State(sync) => sync.is_complete(),
		}
	}

	/// Returns state sync estimated progress (percentage, bytes)
	pub fn progress(&self) -> WarpSyncProgress<B> {
		match &self.phase {
			Phase::WarpProof { .. } => WarpSyncProgress {
				phase: WarpSyncPhase::DownloadingWarpProofs,
				total_bytes: self.total_proof_bytes,
			},
			Phase::TargetBlock(_) => WarpSyncProgress {
				phase: WarpSyncPhase::DownloadingTargetBlock,
				total_bytes: self.total_proof_bytes,
			},
			Phase::PendingTargetBlock { .. } => WarpSyncProgress {
				phase: WarpSyncPhase::AwaitingTargetBlock,
				total_bytes: self.total_proof_bytes,
			},
			Phase::State(sync) => WarpSyncProgress {
				phase: if self.is_complete() {
					WarpSyncPhase::ImportingState
				} else {
					WarpSyncPhase::DownloadingState
				},
				total_bytes: self.total_proof_bytes + sync.progress().size,
			},
		}
	}
}