From eed69173e4189fe17cb69180e9e84587a14d9355 Mon Sep 17 00:00:00 2001
From: Tq <tqmeng123@gmail.com>
Date: Thu, 16 Jan 2025 20:05:13 -0500
Subject: [PATCH] Update recieve_statustext.py

---
 modules/recieve_statustext.py | 78 ++++++++++++++++++-----------------
 1 file changed, 41 insertions(+), 37 deletions(-)

diff --git a/modules/recieve_statustext.py b/modules/recieve_statustext.py
index d7e718a..dc93b45 100644
--- a/modules/recieve_statustext.py
+++ b/modules/recieve_statustext.py
@@ -6,9 +6,10 @@
 import argparse
 from pathlib import Path
 
+
 from pymavlink import mavutil
-import modules.common.modules.data_encoding.message_encoding_decoding
-import modules.common.modules.data_encoding.metadata_encoding_decoding
+from modules.common.modules.data_encoding import message_encoding_decoding
+from modules.common.modules.data_encoding import metadata_encoding_decoding
 from modules.common.modules.data_encoding import worker_enum
 import modules.common.modules.kml.kml_conversion
 
@@ -31,7 +32,12 @@ def main(save_directory: str, document_name_prefix: str) -> int:
     Returns:
         int: 0 on success, -1 on error (connection failure, invalid data, or failure to save KML).
     """
+    save_directory = Path(args.save_directory)
+    document_name_prefix = args.document_name_prefix
+    
     vehicle = mavutil.mavlink_connection(CONNECTION_ADDRESS, source_system=255, source_component=0)
+    # Catch library and other unexpected errors
+    # pylint: disable=broad-exception-caught
     try:
         vehicle.wait_heartbeat()
         print("connected")
@@ -52,42 +58,43 @@ def main(save_directory: str, document_name_prefix: str) -> int:
             success, worker_id, expected_positions_count = (
                 metadata_encoding_decoding.decode_metadata(msg.data)
             )
-            if success and worker_id == worker_enum.WorkerEnum.COMMUNICATIONS_WORKER:
-                received_positions_count = 0
-                while received_positions_count < expected_positions_count:
-                    gps_msg = vehicle.recv_match(type="STATUSTEXT", blocking=True)
-                    if not gps_msg:
-                        print("no message")
-                    elif gps_msg.get_type() == "BAD_DATA":
-                        if mavutil.all_printable(gps_msg.data):
-                            print(gps_msg.data)
-                        return -1
-                    else:
-                        success, worker_id, global_position = (
-                            message_encoding_decoding.decode_bytes_to_position_global(gps_msg.data)
-                        )
-                        if success and worker_id == worker_enum.WorkerEnum.COMMUNICATIONS_WORKER:
-                            print(
-                                f"Decoded GPS Data: {global_position.latitude}, {global_position.longitude}, {global_position.altitude}"
-                            )
-                            positions.append(global_position)
-                            received_positions_count += 1
-                        else:
-                            print(
-                                f"Unsuccessful or Non-communications worker message (worker_id: {worker_id})"
-                            )
-                            return -1
-                success, kml_path = kml_conversion.positions_to_kml(
-                    positions, document_name_prefix, save_directory
+            if not success or worker_id != worker_enum.WorkerEnum.COMMUNICATIONS_WORKER:
+                return
+            
+            received_positions_count = 0
+            while received_positions_count < expected_positions_count:
+                gps_msg = vehicle.recv_match(type="STATUSTEXT", blocking=True)
+                if not gps_msg:
+                    print("no message")
+                    continue
+            
+                if gps_msg.get_type() == "BAD_DATA":
+                    if mavutil.all_printable(gps_msg.data):
+                        print(gps_msg.data)
+                    return -1
+            
+                success, worker_id, global_position = (
+                    message_encoding_decoding.decode_bytes_to_position_global(gps_msg.data)
                 )
-                if success:
-                    print(f"KML file saved to {kml_path}")
-                else:
-                    print("Failed to save KML file")
+                if not success or worker_id != worker_enum.WorkerEnum.COMMUNICATIONS_WORKER:
+                    print(f"Unsuccessful or Non-communications worker message (worker_id: {worker_id})")
                     return -1
+            
+                print(
+                    f"Decoded GPS Data: {global_position.latitude}, {global_position.longitude}, {global_position.altitude}"
+                )
+                positions.append(global_position)
+                received_positions_count += 1
+            
+            success, kml_path = kml_conversion.positions_to_kml(
+                positions, document_name_prefix, save_directory
+            )
+            if not success:
+                print("Failed to save KML file")
+                return -1
+            print(f"KML file saved to {kml_path}")
     return 0
 
-
 DEFAULT_SAVE_DIRECTORY = "logs"
 DEFAULT_DOCUMENT_NAME_PREFIX = "IR hotspot locations"
 
@@ -107,7 +114,4 @@ def main(save_directory: str, document_name_prefix: str) -> int:
     )
 
     args = parser.parse_args()
-    save_directory = Path(args.save_directory)
-    document_name_prefix = args.document_name_prefix
-
     main(save_directory, document_name_prefix)