use crate::{NodeWithCache, PreimageCache};
use jam_std_common::{Node, VersionedParameters};
use jam_types::ProtocolParameters;
use jsonrpsee::ws_client::{WsClient, WsClientBuilder};
use parking_lot::Mutex;
use std::{ops::ControlFlow, sync::Arc, time::Duration};
use tokio::time::sleep;
use url::Url;

// TODO @ivan Values need to be updated.
const MAX_REQUEST_SIZE: u32 = 100 * 1024 * 1024;
const MAX_RESPONSE_SIZE: u32 = 100 * 1024 * 1024;

const REQUEST_TIMEOUT: Duration = Duration::from_secs(6);
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);

/// Manages connections to JAM nodes.
pub struct ConnectionMan {
	urls: Vec<Url>,
	node: Option<Arc<NodeWithCache<WsClient>>>,
	cached_preimages: Arc<Mutex<PreimageCache>>,
}

impl ConnectionMan {
	/// Connects to the first available JAM node and applies its parameters _globally_.
	///
	/// NOTE Parameters should be applied as early as is possible in the program to ensure
	/// correctness.
	///
	/// If you've already applied parameters, simply use [`Self::new`].
	pub async fn apply_parameters(
		urls: Vec<Url>,
		cached_preimages: Arc<Mutex<PreimageCache>>,
	) -> Self {
		let apply_params = |url: &Url, params: ProtocolParameters| {
			if let Err(e) = params.apply() {
				log::warn!("Failed to apply chain parameters from JAM node {url}: {e}");
				return ControlFlow::Continue(());
			}
			ControlFlow::Break(())
		};
		let mut man = Self::new(urls, cached_preimages);
		man.connect(apply_params).await;
		man
	}

	/// Creates new connection manager.
	///
	/// - Doesn't automatically connect to any of the nodes.
	/// - Doesn't apply protocol parameters.
	pub fn new(urls: Vec<Url>, cached_preimages: Arc<Mutex<PreimageCache>>) -> Self {
		Self { urls, node: None, cached_preimages }
	}

	/// Returns a connection to one of the JAM nodes.
	///
	/// Initiates a new connection if there are no active connections.
	pub async fn get(&mut self) -> Arc<dyn Node> {
		match self.node {
			Some(ref node) if node.get().is_connected() => node.clone(),
			_ => {
				let check_params = |url: &Url, params: ProtocolParameters| {
					if params != ProtocolParameters::get() {
						log::warn!("Skipping JAM node {url}: incompatible protocol parameters");
						return ControlFlow::Continue(());
					}
					ControlFlow::Break(())
				};
				self.connect(check_params).await
			},
		}
	}

	async fn connect(
		&mut self,
		mut check_on_connect: impl FnMut(&Url, ProtocolParameters) -> ControlFlow<()>,
	) -> Arc<NodeWithCache<WsClient>> {
		loop {
			let mut num_incompatible_nodes: usize = 0;
			for url in self.urls.iter() {
				match WsClientBuilder::new()
					.max_request_size(MAX_REQUEST_SIZE)
					.max_response_size(MAX_RESPONSE_SIZE)
					.request_timeout(REQUEST_TIMEOUT)
					.connection_timeout(CONNECT_TIMEOUT)
					.build(url)
					.await
				{
					Ok(client) => {
						let node =
							Arc::new(NodeWithCache::new(client, self.cached_preimages.clone()));
						let Ok(VersionedParameters::V1(params)) = node.parameters().await else {
							continue;
						};
						match check_on_connect(url, params) {
							ControlFlow::Continue(()) => {
								num_incompatible_nodes += 1;
								continue;
							},
							ControlFlow::Break(()) => return self.node.insert(node).clone(),
						}
					},
					Err(e) => log::debug!("Failed to connect to JAM node {url}: {e}; retrying..."),
				}
			}
			let sleep_time = if num_incompatible_nodes == self.urls.len() {
				log::warn!("All JAM nodes have incompatible protocol parameters");
				Duration::from_secs(30)
			} else {
				Duration::from_secs(1)
			};
			sleep(sleep_time).await;
		}
	}
}
