From ccaad6eda23ea159b3fe667aed6b2fc72ae85158 Mon Sep 17 00:00:00 2001
From: Erwin Boskma <erwin@datarift.nl>
Date: Sat, 14 Oct 2023 16:44:23 +0200
Subject: [PATCH] Fixed reconnecting for realsies now. I hope.

---
 src/homeassistant.rs | 87 ++++++++++++++------------------------------
 1 file changed, 28 insertions(+), 59 deletions(-)

diff --git a/src/homeassistant.rs b/src/homeassistant.rs
index aa178c9..00fcb48 100644
--- a/src/homeassistant.rs
+++ b/src/homeassistant.rs
@@ -6,6 +6,7 @@ use futures::{SinkExt, StreamExt};
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use thiserror::Error;
+use tokio::time::timeout;
 use tracing::{debug, error, trace};
 
 use crate::output::{i3blocks::I3Blocks, waybar::Waybar, OutputFormat};
@@ -215,20 +216,6 @@ impl Message {
         }
     }
 
-    fn ping() -> Self {
-        Self {
-            message_type: MessageType::Ping,
-            ..Default::default()
-        }
-    }
-
-    fn pong() -> Self {
-        Self {
-            message_type: MessageType::Pong,
-            ..Default::default()
-        }
-    }
-
     fn to_json(&self) -> String {
         serde_json::to_string(&self).unwrap()
     }
@@ -260,35 +247,33 @@ impl HomeAssistant {
 
         let (mut ws_stream, _) = connect_async(&api_url).await?;
 
-        let mut ping_interval = tokio::time::interval(Duration::from_millis(5000));
-
         loop {
-            tokio::select! {
-                msg = ws_stream.next() => match msg {
+            match timeout(Duration::from_secs(5), ws_stream.next()).await {
+                Ok(msg) => match msg {
                     Some(msg) => match msg {
-                        Ok(msg) => {
-                            match self.handle_message(msg).await {
-                                Ok(response_messages) => {
-                                    for response in response_messages {
-                                        ws_stream.send(response).await?
-                                    }
+                        Ok(msg) => match self.handle_message(msg).await {
+                            Ok(response_messages) => {
+                                for response in response_messages {
+                                    ws_stream.send(response).await?
                                 }
-                                Err(e) => {
-                                    if let Some(err) = e.downcast_ref::<HomeAssistantError>() {
-                                        match err {
-                                            HomeAssistantError::EmptyMessage => {
-                                                debug!("Received empty message, ignoring...")
-                                            }
-                                            HomeAssistantError::UnhandledMessage(message_type) => {
-                                                debug!(
+                            }
+                            Err(e) => {
+                                if let Some(err) = e.downcast_ref::<HomeAssistantError>() {
+                                    match err {
+                                        HomeAssistantError::EmptyMessage => {
+                                            debug!("Received empty message, ignoring...")
+                                        }
+                                        HomeAssistantError::UnhandledMessage(message_type) => {
+                                            debug!(
                                                     "Received '{message_type}', we're currently not handling that."
                                                 )
-                                            }
-                                            HomeAssistantError::JsonError(e) => error!("Error deserializing JSON: {e}"),
                                         }
-                                    } else {
-                                        error!("{e}");
+                                        HomeAssistantError::JsonError(e) => {
+                                            error!("Error deserializing JSON: {e}")
+                                        }
                                     }
+                                } else {
+                                    error!("{e}");
                                 }
                             }
                         },
@@ -296,23 +281,19 @@ impl HomeAssistant {
                             error!("Error in receiving message: {e}");
                             // break;
                         }
-                    }
+                    },
                     None => {
                         // I recon we should reconnect here. Probably.
                         error!("Received 'None' from handle_message");
                         break;
-                    },
-                },
-                _ = ping_interval.tick() => {
-                    if self.auth_complete {
-                        let mut ping_message = Message::ping();
-                        ping_message.id = Some(self.incrementing_id());
-                        debug!("[S] Ping!");
-
-                        ws_stream.send(ping_message.to_message()).await?;
                     }
                 },
-            }
+                Err(e) => {
+                    error!("Timeout reading message: {e}");
+                    // return Err(HomeAssistantError::EmptyMessage);
+                    bail!(e);
+                }
+            };
         }
 
         Ok(())
@@ -384,18 +365,6 @@ impl HomeAssistant {
 
                 vec![]
             }
-            MessageType::Ping => {
-                debug!("[R] Ping!");
-                let mut pong = Message::pong();
-                pong.id = Some(self.incrementing_id());
-
-                vec![pong.to_message()]
-            }
-            MessageType::Pong => {
-                debug!("[R] Pong!");
-
-                vec![]
-            }
             _ => bail!(HomeAssistantError::UnhandledMessage(
                 message.message_type.to_string()
             )),