/*! Provides Websocket transport running in Tokio runtime
*/

use std::{collections::HashMap, convert::TryInto, sync::{Arc, Mutex, RwLock}, time::{Duration, Instant}};

use async_tungstenite::{tokio::connect_async, tungstenite};
use futures::{stream::{SplitSink, SplitStream}, SinkExt, StreamExt};
use tokio::{net::TcpListener, time::{sleep, timeout}};
use tracing::{error, info, trace, Level};
use url::Url;

use crate::{events::EventStream, message::Message, node::remote::{Descriptor, RemoteNode}, LocalNode};


// tiny state machine engine
struct State<T: Clone> {
	state: RwLock<T>,
	notify: tokio::sync::Notify,
}

#[derive(Debug)]
enum StateError<T> {
	UnexpectedState(T),
}


impl<T: std::fmt::Debug> std::fmt::Display for StateError<T> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self {
			StateError::UnexpectedState(_) => write!(f, "State error {:?}", self),
		}
	}
}

impl<T: std::fmt::Debug> std::error::Error for StateError<T> {}


impl<T: Clone + Eq> State<T> {
	fn new(initial: T) -> State<T> {
		State { state: RwLock::new(initial), notify: tokio::sync::Notify::new() }
	}

	fn change<F: FnOnce(T) -> T>(&self, closure: F) -> (T, T) {
		let mut state = self.state.write().unwrap();
		let old_state = state.clone();
		*state = (closure)(old_state.clone());
		self.notify.notify_waiters();
		(old_state, state.clone())
	}

	fn change_from_to(&self, from: T, to: T) -> Result<(), StateError<T>> {
		let mut state = self.state.write().unwrap();
		let old_state = state.clone();
		if old_state == from {
			*state = to;
			self.notify.notify_waiters();
			Ok(())
		} else {
			Err(StateError::UnexpectedState(old_state))
		}
	}

	fn notify(&self) -> &tokio::sync::Notify {
		&self.notify
	}

	fn current(&self) -> T {
		self.state.read().unwrap().clone()
	}
}

#[derive(Clone)]
struct ConnectionParameters {
	timeout: Duration,
	diff_request: Option<bool>,
	diff_produce: Option<bool>,
}

impl ConnectionParameters {
	fn from_url(url: &Url) -> Result<ConnectionParameters, String> {
		let mut timeout = 20.0;
		let mut diff_request = None;
		let mut diff_produce = None;
		if let Some(fragment) = url.fragment() {
			//TODO: split_once
			for key_val in fragment.split(',').map(|key_value| key_value.splitn(2, '=').collect::<Vec<&str>>()).filter(|key_value| key_value.len() == 2) {
				match key_val[0] {
					"timeout" => timeout = key_val[1].parse::<f32>().expect("can't parse 'timeout' parameter's value"),
					"diff-request" =>
						diff_request = Some(match key_val[1] {
							"true" => true,
							"false" => false,
							_ => return Err("'diff-request' parameter can only be set to 'true' or 'false'".to_string()),
						}),
					"diff-produce" =>
						diff_produce = Some(match key_val[1] {
							"true" => true,
							"false" => false,
							_ => return Err("'diff-request' parameter can only be set to 'true' or 'false'".to_string()),
						}),
					other => return Err(format!("unknown parameter ('{}') passed in url fragment", other)),
				}
			}
		};
		Ok(ConnectionParameters { timeout: Duration::from_secs_f32(timeout), diff_request, diff_produce })
	}
}


pub struct WebsocketListener {
	shared: Arc<WebsocketListenerShared>,
}

struct WebsocketListenerShared {
	state: State<WebsocketListenerState>,
	url: Url,
	parameters: ConnectionParameters,
}

#[derive(Eq, PartialEq, Clone, Debug)]
enum WebsocketListenerState {
	Created,
	Listening,
	Stop,
	Stopped,
}


impl WebsocketListener {
	#[must_use = "Creating WebsocketListener does nothing useful by itself. You'll have to start() the listener to make it work."]
	pub fn new<T: TryInto<Url> + std::fmt::Debug>(address: T) -> Result<WebsocketListener, String>
	where
		T::Error: std::fmt::Debug,
	{
		let url: Url = address.try_into().map_err(|error| format!("{:?}", error))?;
		Ok(WebsocketListener {
			shared: Arc::new(WebsocketListenerShared {
				state: State::new(WebsocketListenerState::Created),
				parameters: ConnectionParameters::from_url(&url)?,
				url,
			}),
		})
	}

	pub fn url(&self) -> Url {
		self.shared.url.clone()
	}

	pub async fn start(&self, local_node: LocalNode) -> Result<(), String> {
		self.shared.state.change_from_to(WebsocketListenerState::Created, WebsocketListenerState::Listening).expect("Unexpected state");
		let bind_host = self.shared.url.host_str().ok_or("Missing host")?.to_string();
		let bind_port = self.shared.url.port().ok_or("Missing port")?;
		match TcpListener::bind((bind_host.clone(), bind_port)).await.map_err(|err| format!("{:?}", err)) {
			Ok(tcp_listener) => {
				let address = tcp_listener.local_addr().unwrap();
				let url = Url::parse(format!("ws://{}:{}", address.ip(), address.port()).as_str()).unwrap();
				info!("Listening on: {}", url.as_str());
				tokio::task::spawn(self.shared.clone().process(local_node, tcp_listener));
				Ok(())
			}
			Err(error) => Err(error),
		}
	}

	pub fn stop(&self) {
		self.shared.state.change(|current| match current {
			WebsocketListenerState::Listening => WebsocketListenerState::Stop,
			WebsocketListenerState::Stopped => WebsocketListenerState::Stopped,
			_ => panic!(),
		});
	}
}

impl WebsocketListenerShared {
	async fn process(self: Arc<WebsocketListenerShared>, local_node: LocalNode, tcp_listener: TcpListener) {
		let connection_states = Arc::new(Mutex::new(HashMap::new()));

		loop {
			let notified = self.state.notify().notified();
			if self.state.change_from_to(WebsocketListenerState::Stop, WebsocketListenerState::Stopped).is_ok() {
				break;
			};
			tokio::select! {
				//TODO: is this safe? or do we have to keep this future across loop iterations?
				new_connection = tcp_listener.accept() => match new_connection {
					Ok((stream, address)) => {
						let address = Url::parse(format!("ws://{:?}",address).as_str()).unwrap();
						info!("Incomming connection from {:?}", address.as_str());
						let local_node_for_closure = local_node.clone();
						let shared_for_closure = self.clone();
						let connection_states_for_closure = connection_states.clone();
						let connection_parameters = self.parameters.clone();
						tokio::task::spawn(async move {
							let mut node = RemoteNode::new(&local_node_for_closure, false, shared_for_closure.parameters.diff_produce, shared_for_closure.parameters.diff_request);
							//TODO: timeout
							match async_tungstenite::tokio::accept_async(stream).await {
								Ok(ws_stream) => {
									let mut connection = WebsocketConnection::new(address.clone(),ws_stream,connection_parameters);
									let (_old_state, new_state) = shared_for_closure.state.change(|current| {
										if current == WebsocketListenerState::Listening {
											connection_states_for_closure.lock().unwrap().insert(connection.url.clone(),connection.state.clone());
										};
										current
									});
									if new_state == WebsocketListenerState::Listening {
										if let Err(error) = connection.remote_node_process(&mut node).await {
											match error {
												Termination::Proper => { info!("Connection ended: {:?}", address.as_str()) }
												Termination::BenignError(error) => { info!("RemoteNode loop error: {:?} {}", address.as_str(), error) }
												Termination::SeriousError(error) => { error!("RemoteNode loop error: {:?} {}", address.as_str(), error) }
											};
										}
										connection_states_for_closure.lock().unwrap().remove(&connection.url.clone());
										connection.close().await;
									}
								}
								Err(error) => {
									info!("Socket accept error: {:?}", error);
								}
							};
						});
					}
					Err(_) => todo!(), //TODO what errors can happen here?
				},
				_ = notified => {}
			}
		}
		for (_url, connection_state) in connection_states.lock().unwrap().drain() {
			connection_state.change_from_to(WebsocketConnectionState::Processing, WebsocketConnectionState::Stop).expect("Unexpected state");
		}
	}
}



#[derive(Clone)]
pub struct WebsocketConnector {
	shared: Arc<WebsocketConnectorShared>,
}

struct WebsocketConnectorShared {
	url: Url,
	state: State<WebsocketConnectorState>,
	connection_state: Mutex<Option<Arc<State<WebsocketConnectionState>>>>,
	parameters: ConnectionParameters,
}

#[derive(Eq, PartialEq, Clone, Debug)]
enum WebsocketConnectorState {
	Created,
	Connecting,
	Connected,
	Stop,
	Stopped,
}

#[derive(Debug)]
pub enum AddressError<E: std::fmt::Debug> {
	UrlConversionError(E),
	ConnectionParametersParsingError(String),
}

impl<E: std::fmt::Debug> std::fmt::Display for AddressError<E> {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		write!(f, "{:?}", self)
	}
}

impl<E: std::fmt::Debug> std::error::Error for AddressError<E> {}


impl WebsocketConnector {
	#[must_use = "Creating WebsocketConnector does nothing useful by itself. You'll have to start() the connector to make it work."]
	pub fn new<T: TryInto<Url>>(address: T) -> Result<Self, AddressError<T::Error>>
	where
		T::Error: std::fmt::Debug,
	{
		let url = address.try_into().map_err(AddressError::UrlConversionError)?;
		let shared = Arc::new(WebsocketConnectorShared {
			parameters: ConnectionParameters::from_url(&url).map_err(AddressError::ConnectionParametersParsingError)?,
			url,
			state: State::new(WebsocketConnectorState::Created),
			connection_state: Mutex::new(None),
		});

		Ok(WebsocketConnector { shared })
	}

	pub fn start(&self, local_node: LocalNode) {
		self.shared.state.change_from_to(WebsocketConnectorState::Created, WebsocketConnectorState::Connecting).expect("Unexpected state");
		let shared = self.shared.clone();
		tokio::task::spawn(async move {
			shared.event_loop(local_node).await;
		});
	}

	pub fn stop(&self) {
		self.shared.state.change(|current| match current {
			WebsocketConnectorState::Connecting => WebsocketConnectorState::Stop,
			WebsocketConnectorState::Connected => {
				let _ = self
					.shared
					.connection_state
					.lock()
					.unwrap()
					.take()
					.unwrap()
					.change_from_to(WebsocketConnectionState::Processing, WebsocketConnectionState::Stop);
				WebsocketConnectorState::Stop
			}
			WebsocketConnectorState::Stopped => WebsocketConnectorState::Stopped,
			_ => panic!(),
		});
	}

	pub fn url(&self) -> Url {
		self.shared.url.clone()
	}
}


//TODO: stop, drop
impl WebsocketConnectorShared {
	async fn event_loop(&self, local_node: LocalNode) {
		let mut node = RemoteNode::new(&local_node, true, self.parameters.diff_produce, self.parameters.diff_request);
		let mut connection_future = None;
		loop {
			//TODO: timeout here
			let notified = self.state.notify().notified();
			// some futures, like connect_async, have side effects, and select will happily drop them on the floor mid-execution (at await point ofc)
			// specifically, connect_async may get terminated after accepting tcp connection, but before handshake is over. leading to connection being dropped.
			if connection_future.is_none() {
				trace!("Connecting to {:?}", self.url.as_str());
				connection_future = Some(Box::pin(connect_async(self.url.to_string())));
			}
			tokio::select! {
				_ = notified => {},
				maybe_ws_stream = connection_future.as_mut().unwrap() => {
					connection_future.take();
					match maybe_ws_stream {
					Ok((ws_stream, _)) => {
						info!("Connected to {:?}", self.url.as_str());
						let mut connection = WebsocketConnection::new(self.url.clone(), ws_stream, self.parameters.clone());
						*self.connection_state.lock().unwrap() = Some(connection.state.clone());
						if self.state.change_from_to(WebsocketConnectorState::Connecting, WebsocketConnectorState::Connected).is_ok() {
							if let Err(error) = connection.remote_node_process(&mut node).await {
								match error {
									Termination::Proper => { info!("Connection ended: {:?}", self.url.as_str()) }
									Termination::BenignError(error) => {info!("RemoteNode loop error: {:?} {:?}", self.url.as_str(), error)}
									Termination::SeriousError(error) => {error!("RemoteNode loop error: {:?} {:?}", self.url.as_str(), error)}
								};
							}
							let _ = self.state.change_from_to(WebsocketConnectorState::Connected, WebsocketConnectorState::Connecting);
							connection.close().await;
						}
						self.connection_state.lock().unwrap().take();
					}
					Err(error) => {
						info!("Connector error: {:?}", error);
					}
				}
			}
			}

			if self.state.change_from_to(WebsocketConnectorState::Stop, WebsocketConnectorState::Stopped).is_ok() {
				return;
			};

			sleep(Duration::from_millis(100)).await;
		}
	}
}

#[derive(Debug)]
enum Termination {
	Proper,
	BenignError(String),
	SeriousError(String),
}

// should be From trait with WebsocketConnection::T::Error, but:
// https://stackoverflow.com/questions/37347311/how-is-there-a-conflicting-implementation-of-from-when-using-a-generic-type
impl Termination {
	fn from_tungstenite_error(error: tungstenite::Error) -> Self {
		match error {
			tungstenite::Error::ConnectionClosed
			| tungstenite::Error::Io(_)
			| tungstenite::Error::Tls(_)
			| tungstenite::Error::Protocol(tungstenite::error::ProtocolError::ResetWithoutClosingHandshake) => Termination::Proper,
			error => Termination::SeriousError(format!("{:?}", error)),
		}
	}

	fn from_other_error(error: impl std::fmt::Debug) -> Self {
		Self::SeriousError(format!("{:?}", error))
	}
}


impl std::fmt::Display for Termination {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		match self {
			Termination::Proper => write!(f, "Polite termination"),
			Termination::BenignError(text) => write!(f, "Benign error termination: {:?}", text),
			Termination::SeriousError(text) => write!(f, "Serious error termination: {:?}", text),
		}
	}
}

impl std::error::Error for Termination {}


#[derive(Eq, PartialEq, Clone, Debug)]
enum WebsocketConnectionState {
	Created,
	Processing,
	Stop,
	Stopped,
}

struct WebsocketConnection<T>
where
	T: Send + Sync + Unpin + futures::Sink<tungstenite::Message> + futures::Stream<Item = Result<tungstenite::Message, tungstenite::Error>>,
	<T as futures::Sink<async_tungstenite::tungstenite::Message>>::Error: std::fmt::Debug,
{
	url: Url,
	parameters: ConnectionParameters,
	ws_stream_read: tokio::sync::Mutex<Option<SplitStream<T>>>,
	ws_stream_write: tokio::sync::Mutex<Option<SplitSink<T, async_tungstenite::tungstenite::Message>>>,
	last_sent_at: std::sync::Mutex<Instant>,
	last_received_at: std::sync::Mutex<Instant>,
	state: Arc<State<WebsocketConnectionState>>,
}


impl<T> WebsocketConnection<T>
where
	T: Send
		+ Sync
		+ Unpin
		+ 'static
		+ futures::Sink<tungstenite::Message, Error = async_tungstenite::tungstenite::Error>
		+ futures::Stream<Item = Result<tungstenite::Message, tungstenite::Error>>,
{
	fn new(address: Url, ws_stream: T, parameters: ConnectionParameters) -> WebsocketConnection<T> {
		let (write, read) = ws_stream.split();
		WebsocketConnection {
			url: address,
			parameters,
			ws_stream_read: tokio::sync::Mutex::new(Some(read)),
			ws_stream_write: tokio::sync::Mutex::new(Some(write)),
			last_sent_at: std::sync::Mutex::new(Instant::now()),
			last_received_at: std::sync::Mutex::new(Instant::now()),
			state: Arc::new(State::new(WebsocketConnectionState::Created)),
		}
	}

	async fn remote_node_process(&mut self, node: &mut RemoteNode) -> Result<(), Termination> {
		self.state.change_from_to(WebsocketConnectionState::Created, WebsocketConnectionState::Processing).map_err(Termination::from_other_error)?;
		if let Some(message) = node.connected().map_err(Termination::from_other_error)? {
			//println!("-> {}  {:?}",node.name(),message);
			self.send(vec![message]).await?;
		}
		let ret = self.remote_node_event_loop(node).await;
		node.disconnected();
		self.state.change(|_| WebsocketConnectionState::Stopped);
		ret
	}

	async fn remote_node_event_loop(&mut self, node: &mut RemoteNode) -> Result<(), Termination> {
		let mut local_node_events: EventStream<Descriptor> = node.local_node_events.clone().into();
		let mut receive_future = None;
		loop {
			let span = tracing::span!(Level::TRACE, "RemoteNode", name = node.name().as_str());
			//{
			//	  let _enter = span.enter();
			//	  trace!("{}", node.local_node.inspect());
			//}
			let cloned_state = self.state.clone();
			let notified = cloned_state.notify().notified();

			if receive_future.is_none() {
				receive_future = Some(Box::pin(self.receive()));
			};

			let mut messages_to_send = Vec::new();
			tokio::select! {
				_close = notified => {},
				message_result = receive_future.as_mut().unwrap() => {
					receive_future.take();
					let _enter = span.enter();
					//println!("<- {}  {:?}",node.name(),&message_result);
					if let Some(message) = node.received_message(message_result?).map_err(Termination::from_other_error)? {
						messages_to_send.push(message);
					}
				},
				maybe_event = local_node_events.next() => {
					let _enter = span.enter();
					if let Some(message) = node.received_local_node_event(maybe_event.unwrap()).map_err(Termination::from_other_error)? {
						messages_to_send.push(message);
					};
				},
			};
			if self.state.current() == WebsocketConnectionState::Stop {
				return Ok(());
			}
			self.send(messages_to_send).await?;
		}
	}

	async fn close(&self) {
		if let Some(read) = self.ws_stream_read.lock().await.take() {
			let write = self.ws_stream_write.lock().await.take().unwrap();
			let _ = read.reunite(write).unwrap().close().await;
		}
	}

	async fn send(&self, messages: Vec<Message>) -> Result<(), Termination> {
		let mut sink_mutex = self.ws_stream_write.lock().await;
		let sink = sink_mutex.as_mut().unwrap();
		if ! messages.is_empty() {
			for message in messages {
				//trace!("Sending to {:?} <- {:?}", self.url, message);
				let tungstenite_message = tungstenite::Message::Binary(rmp_serde::to_vec(&message).unwrap());
				sink.send(tungstenite_message).await.map_err(Termination::from_tungstenite_error)?;
			}
			*self.last_sent_at.lock().unwrap() = Instant::now();  //FORMER-BUG, REGRESSION-TEST-NEEDED: this was called even if messages vector was empty
		}
		Ok(())
	}

	async fn receive(&self) -> Result<Message, Termination> {
		let mut ws_stream_read = self.ws_stream_read.lock().await;
		loop {
			// sending keep-alives if needed
			if Instant::now() > { *self.last_sent_at.lock().unwrap() } + self.parameters.timeout / 2 {
				*self.last_sent_at.lock().unwrap() = Instant::now();
				self.ws_stream_write
					.lock()
					.await
					.as_mut()
					.unwrap()
					.send(tungstenite::Message::Ping(vec![]))
					.await
					.map_err(Termination::from_tungstenite_error)?;
			};
			// trying to receive something
			if let Ok(message) = timeout(self.parameters.timeout/4, ws_stream_read.as_mut().unwrap().next()).await {
				match message {
					Some(Ok(tungstenite::Message::Binary(message_msgpack))) => {
						*self.last_received_at.lock().unwrap() = Instant::now();
						return rmp_serde::from_read_ref::<Vec<u8>, Message>(&message_msgpack).map_err(Termination::from_other_error);
					}
					Some(Ok(tungstenite::Message::Ping(_payload))) => {
						*self.last_received_at.lock().unwrap() = Instant::now();
					}
					Some(Ok(tungstenite::Message::Pong(_payload))) => {
						*self.last_received_at.lock().unwrap() = Instant::now();
					}
					Some(Ok(tungstenite::Message::Close(_payload))) => {
						return Err(Termination::Proper);
					}
					Some(Ok(_other_message)) => {
						todo!();
					}
					Some(Err(error)) => return Err(Termination::from_tungstenite_error(error)),
					None => todo!(),
				}
			};
			// checking timeout
			if Instant::now() > { *self.last_received_at.lock().unwrap() } + self.parameters.timeout {
				return Err(Termination::BenignError("Read timeout".to_string()));
			}
		}
	}
}


pub mod ffi {

	use super::*;
	use crate::ffi::FFIError;


	pub struct TokioRuntime {
		runtime: tokio::runtime::Runtime,
	}



	#[no_mangle]
	pub extern "C" fn hakuban_tokio_init_multi_thread(worker_threads: usize) -> *mut TokioRuntime {
		let mut builder = tokio::runtime::Builder::new_multi_thread(); //new_current_thread()
		if worker_threads > 0 {
			builder.worker_threads(worker_threads);
		};
		builder.enable_all();
		builder.thread_name("hakuban");
		let runtime = builder.build().unwrap();
		Box::into_raw(Box::new(TokioRuntime { runtime }))
	}

	//TODO: deinit

	#[repr(C)]
	pub struct FFITokioWebsocketConnectorNewResult {
		error: FFIError,
		websocket_connector: *mut WebsocketConnector,
	}


	#[no_mangle]
	extern "C" fn hakuban_tokio_websocket_connector_new(runtime: *mut TokioRuntime, address: *const i8) -> FFITokioWebsocketConnectorNewResult {
		let runtime: Box<TokioRuntime> = unsafe { Box::from_raw(runtime) };
		let _runtime_context = runtime.runtime.enter();
		std::mem::forget(runtime);
		let address = unsafe { std::ffi::CStr::from_ptr(address).to_string_lossy().into_owned() };
		let connector = match WebsocketConnector::new(address.as_str()) {
			Ok(connector) => FFITokioWebsocketConnectorNewResult { error: FFIError::None, websocket_connector: Box::into_raw(Box::new(connector)) },
			Err(_error) => FFITokioWebsocketConnectorNewResult { error: FFIError::InvalidURL, websocket_connector: std::ptr::null_mut() },
		};
		connector
	}

	#[no_mangle]
	extern "C" fn hakuban_tokio_websocket_connector_start(runtime: *mut TokioRuntime, connector: *mut WebsocketConnector, local_node: *mut LocalNode) {
		let runtime: Box<TokioRuntime> = unsafe { Box::from_raw(runtime) };
		let _runtime_context = runtime.runtime.enter();
		let local_node: Box<LocalNode> = unsafe { Box::from_raw(local_node) };
		let connector: Box<WebsocketConnector> = unsafe { Box::from_raw(connector) };
		connector.start(*local_node.clone());
		std::mem::forget(local_node);
		std::mem::forget(connector);
		std::mem::forget(runtime);
	}

	//TODO: connector stop, drop
	//TODO: listener new, start, stop, drop
}