#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from std_msgs.msg import String
import subprocess
import time
import re
from datetime import datetime

class WiFiMonitorNode(Node):
    def __init__(self):
        super().__init__('wifi_monitor_node')
        self.declare_parameter('interface', 'wlan0')
        self.declare_parameter('roam_threshold_dbm', -70)
        self.declare_parameter('gateway_check_interval', 5.0)
        self.declare_parameter('offline_warning_timeout', 30.0)
        self.declare_parameter('target_ssid', 'SaltyLab')
        self.declare_parameter('fallback_tether', True)
        self.declare_parameter('coverage_log_file', '/tmp/wifi_coverage.log')

        self.interface = self.get_parameter('interface').value
        self.roam_threshold = self.get_parameter('roam_threshold_dbm').value
        self.gateway_interval = self.get_parameter('gateway_check_interval').value
        self.offline_timeout = self.get_parameter('offline_warning_timeout').value
        self.target_ssid = self.get_parameter('target_ssid').value
        self.fallback_tether = self.get_parameter('fallback_tether').value
        self.log_file = self.get_parameter('coverage_log_file').value

        self.current_ssid = ""
        self.current_signal_dbm = -100
        self.connected = False
        self.roaming = False
        self.last_gateway_ping = time.time()
        self.offline_start = None
        self.tether_active = False

        self.wifi_state_pub = self.create_publisher(String, '/saltybot/wifi_state', 10)
        self.warning_pub = self.create_publisher(String, '/saltybot/speech_text', 10)
        self.create_subscription(String, '/saltybot/wifi_cmd', self.cmd_callback, 10)

        self.monitor_timer = self.create_timer(1.0, self.monitor_loop)
        self.enable_roaming()
        self.get_logger().info("WiFi monitor initialized")

    def enable_roaming(self):
        try:
            subprocess.run(["wpa_cli", "-i", self.interface, "set_network", "0", "bgscan", '"simple:30:-70:600"'], capture_output=True, timeout=5)
            self.get_logger().info("Roaming enabled")
        except Exception as e:
            self.get_logger().warn(f"Roaming setup: {e}")

    def scan_aps(self):
        try:
            result = subprocess.run(["nmcli", "dev", "wifi", "list", "--rescan", "auto"], capture_output=True, text=True, timeout=10)
            aps = {}
            for line in result.stdout.split('\n')[1:]:
                if not line.strip():
                    continue
                parts = line.split()
                if len(parts) >= 7:
                    ssid = parts[1]
                    signal = int(parts[6])
                    aps[ssid] = signal
            return aps
        except Exception as e:
            self.get_logger().debug(f"Scan error: {e}")
            return {}

    def get_wifi_status(self):
        try:
            result = subprocess.run(["iwconfig", self.interface], capture_output=True, text=True, timeout=5)
            status = {'connected': False, 'ssid': '', 'signal_dbm': -100, 'ap_addr': '', 'link_quality': 0, 'tx_rate': 0, 'rx_rate': 0}
            output = result.stdout
            ssid_match = re.search(r'ESSID:"([^"]*)"', output)
            if ssid_match:
                status['ssid'] = ssid_match.group(1)
                status['connected'] = len(status['ssid']) > 0
            signal_match = re.search(r'Signal level[=:\s]*(-?\d+)', output)
            if signal_match:
                status['signal_dbm'] = int(signal_match.group(1))
            ap_match = re.search(r'Access Point:\s*([0-9A-Fa-f:]+)', output)
            if ap_match:
                status['ap_addr'] = ap_match.group(1)
            quality_match = re.search(r'Link Quality[=:\s]*(\d+)', output)
            if quality_match:
                status['link_quality'] = int(quality_match.group(1))
            return status
        except Exception as e:
            self.get_logger().error(f"iwconfig error: {e}")
            return None

    def gateway_ping(self):
        try:
            result = subprocess.run(["ping", "-c", "1", "-W", "2", "8.8.8.8"], capture_output=True, timeout=5)
            return result.returncode == 0
        except:
            return False

    def activate_tether(self):
        if not self.fallback_tether or self.tether_active:
            return
        try:
            self.get_logger().warn("Activating USB tethering fallback")
            subprocess.run(["nmcli", "device", "connect", "usb0"], capture_output=True, timeout=10)
            self.tether_active = True
        except Exception as e:
            self.get_logger().error(f"Tether activation failed: {e}")

    def deactivate_tether(self):
        if not self.tether_active:
            return
        try:
            subprocess.run(["nmcli", "device", "disconnect", "usb0"], capture_output=True, timeout=10)
            self.tether_active = False
        except Exception as e:
            self.get_logger().debug(f"Tether deactivation: {e}")

    def check_roaming_needed(self):
        if self.current_signal_dbm < self.roam_threshold:
            aps = self.scan_aps()
            for ssid, signal in aps.items():
                if ssid == self.target_ssid and signal > self.current_signal_dbm + 5:
                    self.get_logger().info(f"Roaming: {ssid} ({signal}dBm) > current ({self.current_signal_dbm}dBm)")
                    self.roaming = True
                    return True
        return False

    def log_coverage(self):
        try:
            with open(self.log_file, 'a') as f:
                timestamp = datetime.now().isoformat()
                f.write(f"{timestamp} | SSID:{self.current_ssid} | Signal:{self.current_signal_dbm}dBm | Connected:{self.connected} | Roaming:{self.roaming}\n")
        except Exception as e:
            self.get_logger().debug(f"Log error: {e}")

    def publish_state(self):
        msg = String()
        msg.data = f"ssid:{self.current_ssid}|signal:{self.current_signal_dbm}dBm|connected:{self.connected}|roaming:{self.roaming}|tether:{self.tether_active}"
        self.wifi_state_pub.publish(msg)

    def cmd_callback(self, msg):
        if msg.data == "rescan":
            self.get_logger().info("WiFi rescan requested")
            self.scan_aps()
        elif msg.data == "enable_tether":
            self.activate_tether()
        elif msg.data == "disable_tether":
            self.deactivate_tether()

    def monitor_loop(self):
        status = self.get_wifi_status()
        if status is None:
            return

        prev_ssid = self.current_ssid
        prev_signal = self.current_signal_dbm
        prev_connected = self.connected

        self.current_ssid = status['ssid']
        self.current_signal_dbm = status['signal_dbm']
        self.connected = status['connected']

        if self.current_ssid != prev_ssid:
            self.get_logger().info(f"AP change: {prev_ssid} → {self.current_ssid}")
            self.log_coverage()

        if abs(self.current_signal_dbm - prev_signal) >= 5:
            self.log_coverage()

        if self.gateway_ping():
            self.last_gateway_ping = time.time()
            self.offline_start = None
            if self.tether_active:
                self.deactivate_tether()
        else:
            if self.offline_start is None:
                self.offline_start = time.time()
            offline_time = time.time() - self.offline_start
            if offline_time > self.offline_timeout:
                warning = String()
                warning.data = "Warning: Lost WiFi connectivity. Activating USB tethering."
                self.warning_pub.publish(warning)
                self.activate_tether()

        self.check_roaming_needed()
        self.publish_state()

def main(args=None):
    rclpy.init(args=args)
    node = WiFiMonitorNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
