"""Real-time temperature data logger with matplotlib visualization.

Logs TMP117 temperature data to CSV and displays live graphs.
Supports both single and dual sensor configurations.
Supports both USB Serial and WiFi/TCP connections with automatic detection.

Usage (examples):
    # Auto-detect connection (tries WiFi, then USB)
    python pc_realtime_logger.py [--no-plot] [--interval 0.5]

    # Force USB serial
    python pc_realtime_logger.py --port /dev/cu.usbmodem2101

    # Force WiFi with specific IP
    python pc_realtime_logger.py --wifi 192.168.1.100

Options:
    --no-plot           Don't display live plot (just log to file)
    --interval SECONDS  Update interval for plot (default: 0.0s)
    --output FILE       CSV output filename (default: data/tmp117_log_<timestamp>.csv)
    --port PORT         Force USB serial connection on specific port
    --wifi IP[:PORT]    Force WiFi/TCP connection to specific IP (default port: 5025)
    --prefer-usb        Prefer USB over WiFi when both available

(c) 2024-26 Prof. Flavio ABREU ARAUJO. All rights reserved.
"""

from __future__ import annotations

import csv
import sys
import threading
import time
import warnings
from collections import deque
from datetime import datetime
from pathlib import Path

import matplotlib.animation as animation  # type: ignore
import matplotlib.pyplot as plt  # type: ignore
from matplotlib.figure import Figure  # type: ignore

from .tools.myserial.scpi_universal import SCPIUniversal

# Suppress Tcl/Tk cleanup warnings when windows are destroyed
warnings.filterwarnings("ignore", category=DeprecationWarning)

# Configuration
PORT = None  # Auto-detect
MAX_POINTS = 1501  # Keep last 101 data points in memory
DEFAULT_UPDATE_INTERVAL = 0  # milliseconds
TEXT_SCALE = 1.4  # Global text scale factor (1.0 = default, 1.4 = 40% larger)


class TemperatureLogger:
    """Real-time temperature logger with matplotlib visualization."""

    def __init__(
        self,
        port: str | None = None,
        wifi_host: str | None = None,
        wifi_port: int = 5025,
        prefer_wifi: bool = True,
        output_file: str | None = None,
        show_plot: bool = True,
        update_interval_ms: int = DEFAULT_UPDATE_INTERVAL,
        max_points: int = MAX_POINTS,
        avg_level: int = 8,
        verbose: bool = False,
    ):
        self.port = port
        self.wifi_host = wifi_host
        self.wifi_port = wifi_port
        self.prefer_wifi = prefer_wifi
        self.show_plot = show_plot
        self.update_interval_ms = update_interval_ms
        self.max_points = max_points
        self.avg_level = avg_level  # SENS:SYNC averaging level (1, 8, 32, 64)
        self.verbose = verbose  # Enable verbose debugging output
        self.is_dual_sensor = False
        self.sensor_count = 1

        # Calculate timeout based on averaging level
        # AVG 1: ~120ms, AVG 8: ~200ms, AVG 32: ~575ms, AVG 64: ~1065ms
        # Add 500ms buffer for communication overhead
        # self._avg_timeouts = {1: 0.5, 8: 0.55, 32: 0.8, 64: 1.2}
        self._avg_timeouts = {1: 0.2, 8: 0.3, 32: 0.7, 64: 1.2}  #!
        self.timeout = self._avg_timeouts.get(avg_level, 1.0)

        # Data storage
        self.timestamps = deque(maxlen=max_points)
        self.temperatures_1 = deque(maxlen=max_points)
        self.temperatures_2 = deque(maxlen=max_points)
        self.skew_ms_data = deque(maxlen=max_points)
        self.diff_data = deque(maxlen=max_points)
        self.start_time: float | None = None
        self.point_count = 0

        # Output file
        if output_file is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_file = f"data/tmp117_log_{timestamp}.csv"
        self.output_file = Path(output_file)
        self.csv_writer: csv.writer | None = None
        self.csv_file = None

        # Threading control
        self.running = False
        self.data_thread: threading.Thread | None = None
        self.data_lock = threading.Lock()
        self.pause_data_collection = False  # Pause flag for manual SCPI commands

        # Matplotlib figure and animation (must persist)
        self.fig: Figure | None = None
        self.ani: animation.FuncAnimation | None = None

        # Color codes for terminal output
        self.CYAN = "\033[96m"
        self.YELLOW = "\033[93m"
        self.GREEN = "\033[92m"
        self.MAGENTA = "\033[95m"
        self.RED = "\033[91m"
        self.BLUE = "\033[94m"
        self.RESET = "\033[0m"
        self.BOLD = "\033[1m"

        # ANSI symbols
        self.DEG = "°"
        self.THERM = "🌡️"
        self.CHECK = "✓"
        self.WARN = "⚠️"

    def connect(self) -> None:
        """Connect to the TMP117 instrument via WiFi or USB.

        Attempts connection in order based on prefer_wifi setting:
        - If prefer_wifi=True: WiFi → USB
        - If prefer_wifi=False: USB → WiFi

        Auto-discovers WiFi IP address via USB if needed.
        """
        print(f"{self.BOLD}{self.CYAN}Connecting to TMP117...{self.RESET}")

        try:
            # Use SCPIUniversal for automatic WiFi/USB detection
            self.instr = SCPIUniversal(
                port=self.port,
                wifi_host=self.wifi_host,
                wifi_port=self.wifi_port,
                prefer_wifi=self.prefer_wifi,
                timeout=self.timeout,
                auto_connect=True,
                sync_mode=True,
            )

            # Print connection info
            conn_type = self.instr.connection_type.upper()
            if conn_type == "TCP":
                print(f"{self.GREEN}{self.CHECK} Connected via: {self.BOLD}WiFi/TCP{self.RESET}")
            else:
                print(f"{self.GREEN}{self.CHECK} Connected via: {self.BOLD}USB Serial{self.RESET}")

            print(
                f"{self.GREEN}{self.CHECK} Address: {self.BOLD}{self.instr.port_name}{self.RESET}\n"
            )

            # Verify device with basic query and clear any initial garbage
            time.sleep(0.2)  # Brief settling time

            # Clear buffers BEFORE any queries
            try:
                transport = self.instr.serial  # Uses property that returns TCP or serial
                if transport and hasattr(transport, "reset_input_buffer"):
                    transport.reset_input_buffer()
            except Exception:
                pass

            response = self.instr.query(":SENS:TEMP?")
            if not response:
                raise RuntimeError("Device not responding to SCPI queries")

            # Clear errors and reset state
            self.instr.write("*CLS")
            time.sleep(0.3)  # Allow device to process *CLS

            # Clear buffers after *CLS (single attempt is sufficient)
            try:
                transport = self.instr.serial
                if transport and hasattr(transport, "reset_input_buffer"):
                    transport.reset_input_buffer()
            except Exception:
                pass

            # Default to single sensor
            self.sensor_count = 1
            self.is_dual_sensor = False

            # Try to get actual sensor count with retry and validation
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    # Clear buffer before query
                    transport = self.instr.serial
                    if transport and hasattr(transport, "reset_input_buffer"):
                        transport.reset_input_buffer()

                    time.sleep(0.05)  # Brief delay before query

                    count_resp = self.instr.query(":SENS:COUN?").strip()

                    # Aggressive buffer clear immediately after query (CRITICAL for WiFi!)
                    time.sleep(0.05)  # Let response fully arrive
                    if transport and hasattr(transport, "reset_input_buffer"):
                        transport.reset_input_buffer()

                    if count_resp and count_resp.isdigit():
                        self.sensor_count = int(count_resp)
                        self.is_dual_sensor = self.sensor_count >= 2
                        break  # Success
                    else:
                        # Invalid response, try again
                        if self.verbose:
                            print(
                                f"  Attempt {attempt + 1}: Invalid sensor count response: '{count_resp}'"
                            )
                        if attempt < max_retries - 1:
                            time.sleep(0.1)
                except Exception as e:
                    if self.verbose:
                        print(f"  Attempt {attempt + 1} failed: {e}")
                    if attempt < max_retries - 1:
                        time.sleep(0.1)

            # Final buffer clear before continuing
            try:
                transport = self.instr.serial
                if transport and hasattr(transport, "reset_input_buffer"):
                    transport.reset_input_buffer()
            except Exception:
                pass

            # Print sensor mode
            print(f"{self.BOLD}{self.YELLOW}Detected {self.sensor_count} sensor(s){self.RESET}")
            if self.is_dual_sensor:
                print(f"{self.YELLOW}Mode: Synchronized pair measurement (SYNC){self.RESET}\n")
            else:
                print(f"{self.YELLOW}Mode: Single sensor (TEMP){self.RESET}\n")

        except ConnectionError as e:
            print(f"\n{self.RED}{self.BOLD}❌ Connection Error:{self.RESET}")
            print(f"{self.RED}{e}{self.RESET}\n")
            print(f"{self.YELLOW}Possible causes:{self.RESET}")
            print("  • Device is disconnected")
            print("  • Device is powered off")
            print("  • USB cable is loose or faulty")
            print("  • WiFi is not configured or network is unavailable\n")
            print(f"{self.CYAN}Solutions:{self.RESET}")
            print("  1. Check the USB connection to your TMP117 device")
            print("  2. Verify WiFi is enabled on the device (:WIFI:ENAB?)")
            print("  3. Check network connectivity (can you ping the device?)")
            print("  4. Try forcing USB mode: --prefer-usb\n")
            # Re-raise as RuntimeError so GUI callers can catch it; do NOT call
            # SystemExit here — that would kill the entire process (and the GUI).
            raise RuntimeError(str(e)) from e
        except Exception as e:
            print(f"\n{self.RED}{self.BOLD}❌ Error: {e}{self.RESET}\n")
            # Re-raise so GUI callers can catch and display the error instead of
            # the whole process being terminated via SystemExit.
            raise

    def open_csv_file(self) -> None:
        """Open CSV file and write headers."""
        # Ensure output directory exists
        try:
            self.output_file.parent.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            print(
                f"{self.RED}Error creating output directory {self.output_file.parent}: {e}{self.RESET}"
            )
            raise

        self.csv_file = open(self.output_file, "w", newline="")  # noqa: SIM115

        if self.is_dual_sensor:
            headers = ["Elapsed_Time_s", "Timestamp", "T1_C", "T2_C", "Skew_ms", "Delta_T_C"]
        else:
            headers = ["Elapsed_Time_s", "Timestamp", "T1_C"]

        self.csv_writer = csv.writer(self.csv_file)
        self.csv_writer.writerow(headers)
        print(f"{self.GREEN}{self.CHECK} Logging to: {self.output_file}{self.RESET}\n")

    def read_temperature(self) -> dict:
        """Read temperature data from instrument.

        Handles corrupted/malformed responses by validating each field.
        """
        try:
            # Clear any residual buffer data before reading temperature
            try:
                transport = self.instr.serial  # Uses property that returns TCP or serial
                if transport and hasattr(transport, "reset_input_buffer"):
                    transport.reset_input_buffer()
            except Exception:
                pass

            if self.is_dual_sensor:
                response = self.instr.query(f":SENS:SYNC? {self.avg_level}")
                # Strip and parse comma-separated values
                parts = [x.strip() for x in response.split(",")]

                # Validate we have enough parts
                if len(parts) < 5:
                    print(
                        f"{self.RED}Error: Expected 5 values, got {len(parts)}: {response}{self.RESET}"
                    )
                    return None

                try:
                    # Try to convert each part to float
                    ts = float(parts[0])
                    t1 = float(parts[1])
                    t2 = float(parts[2])
                    skew_ms = float(parts[3])
                    diff = float(parts[4])

                    # Sanity check: temperatures should be in TMP117 spec range (-273.15 to +2500°C)
                    # Cap out-of-range values instead of rejecting
                    if t1 < -273.15 or t1 > 2500:
                        t1 = max(-273.15, min(2500, t1))
                    if t2 < -273.15 or t2 > 2500:
                        t2 = max(-273.15, min(2500, t2))

                    return {
                        "ts": ts,
                        "t1": t1,
                        "t2": t2,
                        "skew_ms": skew_ms,
                        "diff": diff,
                    }
                except ValueError as ve:
                    print(
                        f"{self.RED}Error parsing dual sensor values: {ve}\n  Response: {repr(response)}{self.RESET}"
                    )
                    return None
            else:
                response = self.instr.query(":SENS:TEMP?")
                # Strip and parse comma-separated values
                parts = [x.strip() for x in response.split(",")]

                # Validate we have enough parts
                if len(parts) < 2:
                    print(
                        f"{self.RED}Error: Expected 2 values, got {len(parts)}: {response}{self.RESET}"
                    )
                    return None

                try:
                    # Try to convert each part to float
                    ts = float(parts[0])
                    t1 = float(parts[1])

                    # Sanity check: temperature should be in TMP117 spec range (-273.15 to +2500°C)
                    # Cap out-of-range values instead of rejecting
                    if t1 < -273.15 or t1 > 2500:
                        t1 = max(-273.15, min(2500, t1))

                    return {"ts": ts, "t1": t1, "t2": None, "skew_ms": None, "diff": None}
                except ValueError as ve:
                    print(
                        f"{self.RED}Error parsing single sensor value: {ve}\n  Response: {repr(response)}{self.RESET}"
                    )
                    return None
        except Exception as e:
            print(f"{self.RED}Error reading temperature: {e}{self.RESET}")
        return None

    def log_data(self, data: dict) -> None:
        """Log data to CSV and memory."""
        if data is None:
            return

        if self.start_time is None:
            self.start_time = time.time()

        elapsed = time.time() - self.start_time
        timestamp = datetime.now().isoformat()

        # Store in memory
        self.timestamps.append(elapsed)
        self.temperatures_1.append(data["t1"])
        self.temperatures_2.append(data["t2"])
        self.skew_ms_data.append(data["skew_ms"])
        self.diff_data.append(data["diff"])
        self.point_count += 1

        # Write to CSV
        if self.is_dual_sensor:
            row = [
                f"{elapsed:.3f}",
                timestamp,
                f"{data['t1']:.4f}",
                f"{data['t2']:.4f}",
                f"{data['skew_ms']:.3f}" if data["skew_ms"] is not None else "N/A",
                f"{data['diff']:.4f}" if data["diff"] is not None else "N/A",
            ]
        else:
            row = [f"{elapsed:.3f}", timestamp, f"{data['t1']:.4f}"]

        self.csv_writer.writerow(row)
        self.csv_file.flush()

        # Console output every 10 points
        if self.point_count % 10 == 0:
            self._print_status(data, elapsed)

    def _print_status(self, data: dict, elapsed: float) -> None:
        """Print current status to console."""

        # Color code temperatures: blue for cold (<20°C), red for hot (>30°C), yellow for normal
        def get_temp_color(temp):
            if temp < 20:
                return self.BLUE  # Cold
            elif temp > 30:
                return self.RED  # Hot
            else:
                return self.YELLOW  # Normal

        if self.is_dual_sensor:
            t1_color = get_temp_color(data["t1"])
            t2_color = get_temp_color(data["t2"])
            print(
                f"{self.MAGENTA}[{elapsed:8.2f}s]{self.RESET} | "
                f"{self.THERM} {self.BOLD}{self.CYAN}T1{self.RESET} "
                f"{t1_color}{data['t1']:7.2f}{self.DEG}C{self.RESET} | "
                f"{self.THERM} {self.BOLD}{self.GREEN}T2{self.RESET} "
                f"{t2_color}{data['t2']:7.2f}{self.DEG}C{self.RESET} | "
                f"skew: {data['skew_ms']:6.2f}ms | "
                f"ΔT: {self.BOLD}{data['diff']:+7.2f}{self.DEG}C{self.RESET}"
            )
        else:
            t1_color = get_temp_color(data["t1"])
            print(
                f"{self.MAGENTA}[{elapsed:8.2f}s]{self.RESET} | "
                f"{self.THERM} {self.BOLD}{self.CYAN}T1{self.RESET} "
                f"{t1_color}{data['t1']:7.2f}{self.DEG}C{self.RESET}"
            )

    def _on_figure_close(self, event) -> None:
        """Handle matplotlib figure close event - stop logging."""
        print(f"\n{self.YELLOW}Matplotlib window closed. Stopping data collection...{self.RESET}")
        # Signal threads to stop first so main loop breaks immediately
        self.running = False
        self._stop_animation_and_close_fig()

    def _stop_animation_and_close_fig(self) -> None:
        """Stop animation timer and close the matplotlib figure if it exists."""
        try:
            if self.ani and self.ani.event_source:
                self.ani.event_source.stop()
        except Exception:
            pass

        try:
            if self.fig and plt.fignum_exists(self.fig.number):
                plt.close(self.fig)
        except Exception:
            pass
        # Clear references so subsequent checks exit quickly
        self.fig = None
        self.ani = None

    def create_figure(self) -> Figure:
        """Create matplotlib figure with subplots."""
        fig = plt.figure(figsize=(10, 8))
        # Set window title instead of suptitle
        fig.canvas.manager.set_window_title(
            "Real-time Temperature Logger (© 2026 Prof. Flavio ABREU ARAUJO)"
        )

        # Connect figure close event to stop logging
        fig.canvas.mpl_connect("close_event", self._on_figure_close)

        if self.is_dual_sensor:
            # Dual sensor layout - only 2 plots (left column)
            self.ax_temps = fig.add_subplot(2, 1, 1)
            self.ax_diff = fig.add_subplot(2, 1, 2)
            # Set unused axes to None
            self.ax_skew = None
            self.ax_both = None

            self.ax_temps.set_ylabel("Temperature (°C)", fontsize=12 * TEXT_SCALE)
            self.ax_temps.grid(True, alpha=0.3)
            self.ax_temps.tick_params(labelsize=10 * TEXT_SCALE)

            self.ax_diff.set_xlabel("Time (s)", fontsize=12 * TEXT_SCALE)
            self.ax_diff.set_ylabel("ΔT = T1 - T2 (°C)", fontsize=12 * TEXT_SCALE)
            self.ax_diff.grid(True, alpha=0.3)
            self.ax_diff.tick_params(labelsize=10 * TEXT_SCALE)
        else:
            # Single sensor layout
            self.ax_temps = fig.add_subplot(1, 1, 1)
            self.ax_skew = None
            self.ax_diff = None
            self.ax_both = None

            self.ax_temps.set_xlabel("Time (s)", fontsize=12 * TEXT_SCALE)
            self.ax_temps.set_ylabel("Temperature (°C)", fontsize=12 * TEXT_SCALE)
            self.ax_temps.grid(True, alpha=0.3)
            self.ax_temps.tick_params(labelsize=10 * TEXT_SCALE)

        plt.tight_layout(rect=[0, 0, 1, 0.96])
        plt.subplots_adjust(left=0.12, right=0.95, top=0.94, bottom=0.08, wspace=0.3, hspace=0.15)
        return fig

    def _data_collection_thread(self) -> None:
        """Background thread for continuous data acquisition.

        This thread runs independently to collect sensor data without being
        blocked by matplotlib rendering operations.
        """
        consecutive_errors = 0
        max_consecutive_errors = 5

        while self.running:
            try:
                # Check if data collection is paused (for manual SCPI commands)
                # IMPORTANT: Check BEFORE acquiring lock to avoid race condition
                if self.pause_data_collection:
                    time.sleep(0.05)  # Short sleep while paused, check frequently
                    continue

                # Thread-safe serial access - lock protects both read and log
                with self.data_lock:
                    # Double-check pause flag after acquiring lock
                    if self.pause_data_collection:
                        continue

                    data = self.read_temperature()
                    if data:
                        self.log_data(data)
                        consecutive_errors = 0  # Reset error counter on success
                    else:
                        consecutive_errors += 1
                        if consecutive_errors >= max_consecutive_errors:
                            print(
                                f"{self.YELLOW}Warning: {consecutive_errors} consecutive read failures{self.RESET}"
                            )
                            # Try to recover by clearing buffers
                            try:
                                ser = getattr(self.instr, "serial", None)
                                if ser:
                                    ser.reset_input_buffer()
                                    ser.reset_output_buffer()
                                    consecutive_errors = 0
                            except Exception:
                                pass
            except Exception as e:
                print(f"{self.RED}Data collection error: {e}{self.RESET}")
                consecutive_errors += 1

            # Small sleep to prevent busy-waiting and allow other threads to acquire lock
            time.sleep(0.01)

    def update_plot(self, frame: int) -> None:
        """Update plot with latest data from shared deques.

        This runs in the main thread and only updates the matplotlib visualization.
        Data collection happens independently in a background thread.
        """
        # Lock to ensure data consistency while reading
        with self.data_lock:
            if not self.temperatures_1:
                return

            ts_list = list(self.timestamps)
            t1_list = list(self.temperatures_1)

            if self.is_dual_sensor:
                t2_list = list(self.temperatures_2)
                diff_list = list(self.diff_data)

                # Clear both axes
                self.ax_temps.clear()
                self.ax_diff.clear()

                # Plot temperatures
                self.ax_temps.plot(ts_list, t1_list, "b-", label="T1", linewidth=1.5, alpha=0.8)
                self.ax_temps.plot(ts_list, t2_list, "r-", label="T2", linewidth=1.5, alpha=0.8)
                self.ax_temps.set_ylabel("Temperature (°C)", fontsize=12 * TEXT_SCALE)
                self.ax_temps.grid(True, alpha=0.3)
                self.ax_temps.tick_params(labelsize=10 * TEXT_SCALE)
                if len(ts_list) > 1:
                    self.ax_temps.legend(fontsize=11 * TEXT_SCALE, loc="upper left")

                # Plot difference
                self.ax_diff.plot(ts_list, diff_list, "orange", linewidth=1.5, alpha=0.8)
                self.ax_diff.axhline(y=0, color="k", linestyle="--", alpha=0.3)
                self.ax_diff.set_xlabel("Time (s)", fontsize=12 * TEXT_SCALE)
                self.ax_diff.set_ylabel("ΔT = T1 - T2 (°C)", fontsize=12 * TEXT_SCALE)
                self.ax_diff.grid(True, alpha=0.3)
                self.ax_diff.tick_params(labelsize=10 * TEXT_SCALE)

                # Ensure layout is preserved during updates
                plt.subplots_adjust(
                    left=0.12, right=0.95, top=0.94, bottom=0.08, wspace=0.3, hspace=0.15
                )
            else:
                # Single sensor plot
                self.ax_temps.clear()
                self.ax_temps.plot(ts_list, t1_list, "b-", linewidth=1.5, alpha=0.8)
                self.ax_temps.set_xlabel("Time (s)", fontsize=12 * TEXT_SCALE)
                self.ax_temps.set_ylabel("Temperature (°C)", fontsize=12 * TEXT_SCALE)
                self.ax_temps.grid(True, alpha=0.3)
                self.ax_temps.tick_params(labelsize=10 * TEXT_SCALE)

    def run(self) -> None:
        """Main application loop with separate data collection thread."""
        try:
            self.connect()
            self.open_csv_file()

            # Clear any residual data in serial buffers before starting data collection
            try:
                self.instr.write("*CLS")
                time.sleep(0.1)
                if hasattr(self.instr, "_ser") and self.instr._ser:
                    self.instr._ser.reset_input_buffer()
            except Exception:
                pass

            # Start the data collection thread
            self.running = True
            self.data_thread = threading.Thread(
                target=self._data_collection_thread,
                daemon=True,
                name="DataCollectionThread",
            )
            self.data_thread.start()

            if self.show_plot:
                self.fig = self.create_figure()
                self.ani = animation.FuncAnimation(
                    self.fig,
                    self.update_plot,
                    interval=self.update_interval_ms if self.update_interval_ms > 0 else 100,
                    cache_frame_data=False,
                )

                # Non-blocking show loop so close_event can break us out cleanly
                plt.show(block=False)
                try:
                    while self.running and self.fig and plt.fignum_exists(self.fig.number):
                        plt.pause(0.1)
                except KeyboardInterrupt:
                    print(f"\n\n{self.BOLD}{self.MAGENTA}Logging stopped by user{self.RESET}")
                    self.running = False
                    self._stop_animation_and_close_fig()
            else:
                # Non-plotting mode: just let the data thread collect
                print(f"{self.BOLD}{self.YELLOW}Collecting data (Ctrl+C to stop)...{self.RESET}\n")
                try:
                    while True:
                        time.sleep(0.1)
                except KeyboardInterrupt:
                    self.running = False

        except SystemExit:
            # Window closed via _on_figure_close() - cleanup will run in finally block
            pass
        except RuntimeError:
            # Device connection errors - already have helpful messages in connect()
            pass
        except KeyboardInterrupt:
            print(f"\n\n{self.BOLD}{self.MAGENTA}Logging stopped by user{self.RESET}")
            self.running = False
            self._stop_animation_and_close_fig()
        except Exception as e:
            print(f"\n{self.RED}Error: {e}{self.RESET}")
            import traceback

            traceback.print_exc()
        finally:
            self.cleanup()

    def cleanup(self) -> None:
        """Clean up resources and stop data collection thread."""
        # Stop the data collection thread
        self.running = False
        if self.data_thread and self.data_thread.is_alive():
            self.data_thread.join(timeout=2.0)

        if self.csv_file:
            self.csv_file.close()
            print(f"{self.GREEN}{self.CHECK} Data saved to: {self.output_file}{self.RESET}")
            print(f"{self.YELLOW}Total points logged: {self.point_count}{self.RESET}\n")

        try:
            if hasattr(self, "instr") and self.instr:
                self.instr.__exit__(None, None, None)
        except Exception:
            pass


def main() -> int:
    """Main entry point."""
    # Parse command line arguments
    args = sys.argv[1:]
    port = PORT
    wifi_host = None
    wifi_port = 5025
    prefer_wifi = True
    show_plot = True
    update_interval_ms = DEFAULT_UPDATE_INTERVAL
    output_file = None

    RED = "\033[91m"
    RESET = "\033[0m"

    i = 0
    while i < len(args):
        arg = args[i]
        if arg == "--no-plot":
            show_plot = False
        elif arg in ("--wifi", "-w"):
            if i + 1 < len(args):
                # Parse WiFi address (IP or IP:PORT)
                addr = args[i + 1]
                if ":" in addr:
                    parts = addr.rsplit(":", 1)
                    wifi_host = parts[0]
                    try:
                        wifi_port = int(parts[1])
                    except ValueError:
                        print(f"{RED}Invalid WiFi port: {parts[1]}{RESET}")
                        return 1
                else:
                    wifi_host = addr
                i += 1
            else:
                print(f"{RED}--wifi requires an IP address{RESET}")
                return 1
        elif arg in ("--port", "-p"):
            if i + 1 < len(args):
                port = args[i + 1]
                i += 1
            else:
                print(f"{RED}--port requires a serial port path{RESET}")
                return 1
        elif arg == "--prefer-usb":
            prefer_wifi = False
        elif arg == "--interval":
            if i + 1 < len(args):
                try:
                    update_interval_ms = int(float(args[i + 1]) * 1000)
                    i += 1
                except ValueError:
                    print(f"{RED}Invalid interval value{RESET}")
                    return 1
        elif arg == "--output":
            if i + 1 < len(args):
                output_file = args[i + 1]
                i += 1
        elif not arg.startswith("-"):
            # Legacy: bare argument is treated as port
            port = arg
        i += 1

    logger = TemperatureLogger(
        port=port,
        wifi_host=wifi_host,
        wifi_port=wifi_port,
        prefer_wifi=prefer_wifi,
        output_file=output_file,
        show_plot=show_plot,
        update_interval_ms=update_interval_ms,
    )
    logger.run()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
