abpl 0.1.3

use std::{
	collections::HashMap,
	error::Error,
	fmt::Display,
	io::{Error as IoError, Result as IoResult},
	mem,
	sync::Arc,
	thread::{self, JoinHandle},
};

use axum::Router;
use futures::{FutureExt as _, channel::oneshot};

use crate::{
	future::block_on,
	providers::ProvidesExitCode,
	types::{
		hotswap_state::HotswapState,
		http::{SocketAddr, SocketListener},
	},
};

#[derive(Debug)]
pub struct HotReloadingAxumService<S: Send + Sync + 'static> {
	state: HotswapState<S>,
	router: Router<()>,
	// using the futures version of oneshot because it has a simpler error type
	active_routers: HashMap<SocketAddr, (oneshot::Sender<()>, JoinHandle<IoResult<()>>)>,
}
impl<S: Send + Sync + 'static> HotReloadingAxumService<S> {
	/// Makes a new hot-reloading axum service.
	///
	/// Note that the `router` should not have been called with `.with_state` yet. That's the job of this struct.
	///
	/// It's important to know that the axum instance(s) run in their dedicated tokio runtime, if you make external
	/// requests, be mindful of <https://github.com/tokio-rs/tokio/issues/7337>
	pub fn new(state: S, router: impl Fn(HotswapState<S>) -> Router<HotswapState<S>>) -> Self {
		let parent_state = HotswapState::new(state);
		let router = router(parent_state.clone_as_parent()).with_state(parent_state.clone_as_parent());
		Self {
			state: parent_state,
			router,
			active_routers: HashMap::new(),
		}
	}

	/// All future requests will have the `state` specified (wrapped in a child `HotswapState`). Any requests
	/// currently running will still have the old state and this function doesn't wait for those requests to finish.
	///
	/// This function returns the old state. Since requests may be in-flight, the result is not guaranteed to be
	/// unique, hence the `Arc`.
	///
	/// If you're making a [super::ReloadableService], this is where you'd put in the new state generated by the new
	/// config.
	pub fn replace_state(&self, state: S) -> Arc<S> {
		self.state.replace_parent_state(state).expect("must be a state parent")
	}

	/// Returns `true` if any axum instances panicked or otherwise ended unexpectedly.
	pub fn has_dead_threads(&self) -> bool {
		self.active_routers
			.iter()
			.any(|(_, (signal, thread))| thread.is_finished() || signal.is_canceled())
	}

	/// If any of the axum instances panicked or otherwise finished unexpectedly, a restart will be attempted.
	///
	/// This function only returns an error if there are no active axum instances.
	///
	/// If you're making a [super::ReloadableService], you should call this function during the interval tick
	pub fn restart_dead_threads(&mut self) -> Result<(), HotReloadAxumError> {
		if !self.has_dead_threads() {
			return Ok(());
		}
		let sockets: Vec<SocketAddr> = self.active_routers.keys().cloned().collect();
		self.bind_sockets(sockets)
	}

	/// Shuts down all axum instances. This is a blocking function which returns when all pending requests are complete.
	pub fn stop(&mut self) {
		for (socket, (signal, thread)) in mem::take(&mut self.active_routers) {
			let _ = signal.send(());
			Self::join_thread_and_log_error(&socket, thread, true);
		}
	}

	/// Starts axum with the sockets specified.
	///
	/// If the specified socket was already bound, it will be untouched.
	///
	/// Any previously bound sockets not in the list provided will be unbound and the axum instances terminated.
	///
	/// This function only returns an error if there are no active axum instances.
	///
	/// If you're making a [super::ReloadableService], you should call this during initial startup, and during every
	/// reload.
	pub fn bind_sockets(&mut self, sockets: impl IntoIterator<Item = SocketAddr>) -> Result<(), HotReloadAxumError> {
		//let mut threads_to_join: Vec<(SocketAddr, JoinHandle<IoResult<()>>)> = Vec::new();
		let mut previous_routers = mem::take(&mut self.active_routers);
		for socket in sockets.into_iter() {
			if previous_routers
				.get(&socket)
				.is_none_or(|(signal, thread)| thread.is_finished() || signal.is_canceled())
			{
				if let Some((_, dead_thread)) = previous_routers.remove(&socket) {
					// We're only here because because of the "or" part in the "is_none_or" clause, so the
					// thread is dead.
					Self::join_thread_and_log_error(&socket, dead_thread, false);
				}
				// previous_routers doesn't have socket or socket is dead, spawn new
				let (bind_error_send, bind_error_recv) = oneshot::channel::<IoError>();
				let (shutdown_send, shutdown_recv) = oneshot::channel::<()>();
				let router = self.router.clone();
				let l_socket = socket.clone();
				match thread::Builder::new().name("hot-reload axum".into()).spawn(move || {
					// Axum instance gets a dedicated thread
					match tokio::runtime::Builder::new_current_thread()
						.enable_all()
						.build()
						.and_then(|runtime| {
							let listener = runtime.block_on(SocketListener::bind(l_socket))?;
							Ok((runtime, listener))
						}) {
						Ok((runtime, listener)) => {
							tracing::debug!("axum thread successfully bound to {}", listener.original_addr());
							// Droping the sender as a means to send a signal saying the bind was successful
							drop(bind_error_send);
							// runtime and binding successfully complete, now to dedicate the thread to running axum
							runtime.block_on(
								axum::serve(listener, router)
									.with_graceful_shutdown(shutdown_recv.map(|_| {}))
									.into_future(),
							)
						},
						Err(runtime_err) => {
							let _ = bind_error_send.send(runtime_err);
							Ok(())
						},
					}
				}) {
					Ok(new_thread) => {
						match block_on(bind_error_recv) {
							Ok(bind_err) => {
								tracing::error!("Unable to bind to {socket}: {bind_err}");
								Self::join_thread_and_log_error(&socket, new_thread, true);
							},
							Err(oneshot::Canceled) => {
								// if the sender was dropped, that means binding was successful, yay!
								self.active_routers.insert(socket, (shutdown_send, new_thread));
							},
						}
					},
					Err(err) => {
						tracing::error!("failed to spawn axum thread serving {socket}: {err}");
					},
				}
			} else if let Some(active_router) = previous_routers.remove(&socket) {
				// previous_routers already has socket, reuse
				self.active_routers.insert(socket, active_router);
			}
		}
		// At this point, everything in previous_routers should be bindings that we don't want anymore
		for (socket, (signal, thread)) in previous_routers {
			let _ = signal.send(());
			Self::join_thread_and_log_error(&socket, thread, true);
		}
		if self.active_routers.is_empty() {
			Err(HotReloadAxumError {})
		} else {
			Ok(())
		}
	}

	fn join_thread_and_log_error(socket: &SocketAddr, thread: JoinHandle<IoResult<()>>, expected: bool) {
		match thread.join() {
			Ok(Ok(())) => {
				if expected {
					tracing::info!("axum thread serving {socket} has finished");
				} else {
					tracing::error!("axum thread serving {socket} finished unexpectedly");
				}
			},
			Ok(Err(axum_err)) => {
				tracing::error!("axum instance serving {socket} encountered an irrecoverable error: {axum_err}");
			},
			Err(panic) => {
				if let Some(panic_str) = panic.downcast_ref::<&str>() {
					tracing::error!("axum thread serving {socket} panicked with message: {panic_str}");
				} else if let Some(panic_string) = panic.downcast_ref::<String>() {
					tracing::error!("axum thread serving {socket} panicked with message: {panic_string}");
				} else {
					tracing::error!("axum thread serving {socket} panicked with unknown error type");
				}
			},
		}
	}
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct HotReloadAxumError {}
impl Display for HotReloadAxumError {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.write_str("no sockets could be listened to")
	}
}
impl Error for HotReloadAxumError {}
impl ProvidesExitCode for HotReloadAxumError {
	fn exit_code(&self) -> std::process::ExitCode {
		// ChatGPT historically told me that no successful binds still counts as a misconfiguration
		// ...but I wonder what Claude will think
		crate::app::consts::EX_CONFIG.into()
	}
}

#[cfg(test)]
#[path = "../tests/app/axum.rs"]
mod tests;