controller/src/ws/mod.rs

use actix::Addr;
use emgauwa_common::constants::WEBSOCKET_RETRY_TIMEOUT;
use emgauwa_common::db::{DbController, DbJunctionRelaySchedule, DbRelay, DbSchedule};
use emgauwa_common::errors::{DatabaseError, EmgauwaError};
use emgauwa_common::models::{Controller, Relay};
use emgauwa_common::types::{ControllerWsAction, ScheduleUid};
use futures::{future, pin_mut, SinkExt, StreamExt};
use sqlx::pool::PoolConnection;
use sqlx::{Pool, Sqlite};
use tokio::time;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{connect_async, tungstenite};

use crate::app_state::AppState;
use crate::utils;
use crate::utils::{app_state_get_relay_notifier, app_state_get_this};

pub async fn run_ws_loop(pool: Pool<Sqlite>, app_state: Addr<AppState>, url: String) {
	log::debug!("Spawned ws loop");
	loop {
		let run_result = run_websocket(pool.clone(), &app_state, &url).await;
		if let Err(err) = run_result {
			log::error!("Error running websocket: {}", err);
		}

		log::info!(
			"Retrying to connect in {} seconds...",
			WEBSOCKET_RETRY_TIMEOUT.as_secs()
		);
		time::sleep(WEBSOCKET_RETRY_TIMEOUT).await;
	}
}

async fn run_websocket(
	pool: Pool<Sqlite>,
	app_state: &Addr<AppState>,
	url: &str,
) -> Result<(), EmgauwaError> {
	log::debug!("Trying to connect to {}", url);
	match connect_async(url).await {
		Ok(connection) => {
			log::info!("Websocket connected");
			let (ws_stream, _) = connection;

			let (mut write, read) = ws_stream.split();

			let ws_action = ControllerWsAction::Register(app_state_get_this(app_state).await?);

			let ws_action_json = serde_json::to_string(&ws_action)?;
			if let Err(err) = write.send(Message::text(ws_action_json)).await {
				log::error!("Failed to register at websocket: {}", err);
				return Ok(());
			}

			let (app_state_tx, app_state_rx) = futures_channel::mpsc::unbounded::<Message>();
			tokio::spawn(read_app_state(app_state.clone(), app_state_tx));
			let app_state_to_ws = app_state_rx.map(Ok).forward(write);

			let read_handler = read.for_each(|msg| handle_message(pool.clone(), app_state, msg));

			pin_mut!(app_state_to_ws, read_handler);
			future::select(app_state_to_ws, read_handler).await;

			log::warn!("Lost connection to websocket");
		}
		Err(err) => {
			log::warn!("Failed to connect to websocket: {}", err,);
		}
	}
	Ok(())
}

async fn read_app_state(
	app_state: Addr<AppState>,
	tx: futures_channel::mpsc::UnboundedSender<Message>,
) -> Result<(), EmgauwaError> {
	let notifier = &*app_state_get_relay_notifier(&app_state).await?;
	loop {
		notifier.notified().await;
		log::debug!("Relay change detected");
		let this = app_state_get_this(&app_state).await?;
		let relay_states = this.get_relay_states();
		let ws_action = ControllerWsAction::RelayStates((this.c.uid, relay_states));

		let ws_action_json = serde_json::to_string(&ws_action)?;
		tx.unbounded_send(Message::text(ws_action_json))
			.map_err(|_| {
				EmgauwaError::Other(String::from(
					"Failed to forward message from app state to websocket",
				))
			})?;
	}
}

async fn handle_message(
	pool: Pool<Sqlite>,
	app_state: &Addr<AppState>,
	message_result: Result<Message, tungstenite::Error>,
) {
	let msg = match message_result {
		Ok(msg) => msg,
		Err(err) => {
			log::error!("Error reading message: {}", err);
			return;
		}
	};
	if let Message::Text(text) = msg {
		match serde_json::from_str(&text) {
			Ok(action) => {
				log::debug!("Received action: {:?}", action);
				let mut pool_conn = match pool.acquire().await {
					Ok(conn) => conn,
					Err(err) => {
						log::error!("Failed to acquire database connection: {:?}", err);
						return;
					}
				};
				let action_res = handle_action(&mut pool_conn, app_state, action).await;
				if let Err(e) = action_res {
					log::error!("Error handling action: {:?}", e);
				}
			}
			Err(e) => {
				log::error!("Error deserializing action: {:?}", e);
			}
		}
	}
}

pub async fn handle_action(
	conn: &mut PoolConnection<Sqlite>,
	app_state: &Addr<AppState>,
	action: ControllerWsAction,
) -> Result<(), EmgauwaError> {
	let this = app_state_get_this(app_state).await?;

	match action {
		ControllerWsAction::Controller(controller) => {
			handle_controller(conn, &this, controller).await?
		}
		ControllerWsAction::Relays(relays) => handle_relays(conn, &this, relays).await?,
		ControllerWsAction::Schedules(schedules) => handle_schedules(conn, schedules).await?,
		ControllerWsAction::RelayPulse((relay_num, duration)) => {
			handle_relay_pulse(app_state, relay_num, duration).await?
		}
		_ => return Ok(()),
	};

	utils::app_state_reload(app_state).await
}

async fn handle_controller(
	conn: &mut PoolConnection<Sqlite>,
	this: &Controller,
	controller: Controller,
) -> Result<(), EmgauwaError> {
	if controller.c.uid != this.c.uid {
		return Err(EmgauwaError::Other(String::from(
			"Controller UID mismatch during update",
		)));
	}
	DbController::get_by_uid(conn, &controller.c.uid)
		.await?
		.ok_or(DatabaseError::NotFound)?
		.update(conn, controller.c.name.as_str(), this.c.relay_count)
		.await?;

	Ok(())
}

async fn handle_schedules(
	conn: &mut PoolConnection<Sqlite>,
	schedules: Vec<DbSchedule>,
) -> Result<(), EmgauwaError> {
	let mut handled_uids = vec![
		// on and off schedules are always present and should not be updated
		ScheduleUid::On,
		ScheduleUid::Off,
	];
	for schedule in schedules {
		if handled_uids.contains(&schedule.uid) {
			continue;
		}
		handled_uids.push(schedule.uid.clone());

		log::debug!("Handling schedule: {:?}", schedule);
		let schedule_db = DbSchedule::get_by_uid(conn, &schedule.uid).await?;

		if let Some(schedule_db) = schedule_db {
			schedule_db
				.update(conn, schedule.name.as_str(), &schedule.periods)
				.await?;
		} else {
			DbSchedule::create(
				conn,
				schedule.uid.clone(),
				schedule.name.as_str(),
				&schedule.periods,
			)
			.await?;
		}
	}

	Ok(())
}

async fn handle_relays(
	conn: &mut PoolConnection<Sqlite>,
	this: &Controller,
	relays: Vec<Relay>,
) -> Result<(), EmgauwaError> {
	for relay in relays {
		if relay.controller.uid != this.c.uid {
			return Err(EmgauwaError::Other(String::from(
				"Controller UID mismatch during relay update",
			)));
		}
		let db_relay = DbRelay::get_by_controller_and_num(conn, &this.c, relay.r.number)
			.await?
			.ok_or(DatabaseError::NotFound)?;

		db_relay.update(conn, relay.r.name.as_str()).await?;

		handle_schedules(conn, relay.schedules.clone()).await?;

		let mut schedules = Vec::new(); // We need to get the schedules from the database to have the right IDs
		for schedule in relay.schedules {
			schedules.push(
				DbSchedule::get_by_uid(conn, &schedule.uid)
					.await?
					.ok_or(DatabaseError::NotFound)?,
			);
		}

		DbJunctionRelaySchedule::set_schedules(conn, &db_relay, schedules.iter().collect()).await?;
	}

	Ok(())
}

async fn handle_relay_pulse(
	app_state: &Addr<AppState>,
	relay_num: i64,
	duration: Option<u32>,
) -> Result<(), EmgauwaError> {
	utils::app_state_relay_pulse(app_state, relay_num, duration).await
}