from rclpy.node import Node
from PyQt6.QtCore import QObject, pyqtSignal
import math
import re
import threading
from threading import Lock
import asyncio
import websockets
import json
import socket
from pymavlink import mavutil
from geometry_msgs.msg import Point, Vector3
from sensor_msgs.msg import BatteryState, NavSatFix, Imu
from std_msgs.msg import Float64
from mavros_msgs.msg import State, VfrHud
from mavros_msgs.srv import CommandBool, CommandTOL

class DroneSignals(QObject):
    update_signal = pyqtSignal(str, str, object)  # (msg_type, drone_id, data)

class UDPMavlinkReceiver(threading.Thread):
    """UDP MAVLink 接收器"""
    def __init__(self, ip, port, signals, connection_name):
        super().__init__(daemon=True)
        self.ip = ip
        self.port = port
        self.signals = signals
        self.connection_name = connection_name
        self.running = False
        self.sock = None
        
    def run(self):
        """執行 UDP 接收循環"""
        self.running = True
        try:
            print(f"UDP MAVLink receiver started on {self.ip}:{self.port}")
            
            # 創建 MAVLink 連接
            mav = mavutil.mavlink_connection(f'udpin:{self.ip}:{self.port}')
            while self.running:
                try:
                    msg = mav.recv_match(blocking=True, timeout=1.0)
                    if msg is None:
                        continue
                    
                    self.process_mavlink_message(msg)
                    
                except socket.timeout:
                    continue
                except Exception as e:
                    print(f"Error receiving MAVLink message: {e}")
                    
        except Exception as e:
            print(f"UDP receiver error: {e}")
        finally:
            if self.sock:
                self.sock.close()
                
    def process_mavlink_message(self, msg):
        """處理 MAVLink 訊息"""
        try:
            msg_type = msg.get_type()
            system_id = msg.get_srcSystem()
            drone_id = f"s8_{system_id}"  # 使用 s8_ 前綴表示 UDP 來源
            
            if msg_type == "HEARTBEAT":
                mode = mavutil.mode_string_v10(msg)
                armed = bool(msg.base_mode & 128)
                self.signals.update_signal.emit('state', drone_id, {
                    'mode': mode,
                    'armed': armed
                })

            elif msg_type == "BATTERY_STATUS":
                voltage = msg.voltages[0] / 1000
                self.signals.update_signal.emit('battery', drone_id, {
                    'voltage': voltage
                })

            elif msg_type == "GLOBAL_POSITION_INT":
                latitude = msg.lat / 1e7
                longitude = msg.lon / 1e7
                self.signals.update_signal.emit('gps', drone_id, {
                    'lat': latitude,
                    'lon': longitude
                })

            elif msg_type == "GPS_RAW_INT":
                fix_type = msg.fix_type

            elif msg_type == "LOCAL_POSITION_NED":
                x = msg.y
                y = msg.x
                z = -msg.z
                self.signals.update_signal.emit('local_pose', drone_id, {
                    'x': x,
                    'y': y,
                    'z': z
                })
                self.signals.update_signal.emit('altitude', drone_id, {
                    'altitude': z
                })
                self.signals.update_signal.emit('velocity', drone_id, {
                    'vx': msg.vx,
                    'vy': msg.vy,
                    'vz': msg.vz
                })

            elif msg_type == "ATTITUDE":
                pitch = math.degrees(msg.pitch)
                self.signals.update_signal.emit('attitude', drone_id, {
                    'pitch': pitch,
                    'roll': 0,
                    'yaw': 0,
                    'rates': (0, 0, 0)
                })

            elif msg_type == "VFR_HUD":
                groundspeed = msg.groundspeed
                heading = msg.heading
                self.signals.update_signal.emit('hud', drone_id, {
                    'heading': heading,
                    'groundspeed': groundspeed
                })
                
        except Exception as e:
            print(f"Error processing MAVLink message: {e}")
            
    def stop(self):
        """停止接收器"""
        self.running = False

class SerialMavlinkReceiver(threading.Thread):
    """串口 MAVLink 接收器"""
    def __init__(self, port, baudrate, signals, connection_name):
        super().__init__(daemon=True)
        self.port = port
        self.baudrate = baudrate
        self.signals = signals
        self.connection_name = connection_name
        self.running = False
        self.mav = None
        
    def run(self):
        """執行串口接收循環"""
        self.running = True
        try:
            print(f"Serial MAVLink receiver started on {self.port} at {self.baudrate} baud")
            
            # 創建 MAVLink 串口連接
            self.mav = mavutil.mavlink_connection(
                self.port,
                baud=self.baudrate,
                source_system=255
            )
            
            print(f"Waiting for heartbeat from {self.port}...")
            self.mav.wait_heartbeat()
            print(f"Heartbeat received from system {self.mav.target_system}, component {self.mav.target_component}")
            
            while self.running:
                try:
                    msg = self.mav.recv_match(blocking=True, timeout=1.0)
                    if msg is None:
                        continue
                    
                    self.process_mavlink_message(msg)
                    
                except Exception as e:
                    if self.running:
                        print(f"Error receiving MAVLink message from serial: {e}")
                    
        except Exception as e:
            print(f"Serial receiver error: {e}")
        finally:
            if self.mav:
                try:
                    self.mav.close()
                except:
                    pass
                
    def process_mavlink_message(self, msg):
        """處理 MAVLink 訊息"""
        try:
            msg_type = msg.get_type()
            system_id = msg.get_srcSystem()
            drone_id = f"s5_{system_id}"  # 使用 serial_ 前綴表示串口來源
            
            if msg_type == "HEARTBEAT":
                mode = mavutil.mode_string_v10(msg)
                armed = bool(msg.base_mode & 128)
                self.signals.update_signal.emit('state', drone_id, {
                    'mode': mode,
                    'armed': armed
                })

            elif msg_type == "BATTERY_STATUS":
                voltage = msg.voltages[0] / 1000
                self.signals.update_signal.emit('battery', drone_id, {
                    'voltage': voltage
                })

            elif msg_type == "GLOBAL_POSITION_INT":
                latitude = msg.lat / 1e7
                longitude = msg.lon / 1e7
                self.signals.update_signal.emit('gps', drone_id, {
                    'lat': latitude,
                    'lon': longitude
                })

            elif msg_type == "GPS_RAW_INT":
                fix_type = msg.fix_type

            elif msg_type == "LOCAL_POSITION_NED":
                x = msg.y
                y = msg.x
                z = -msg.z
                self.signals.update_signal.emit('local_pose', drone_id, {
                    'x': x,
                    'y': y,
                    'z': z
                })
                self.signals.update_signal.emit('altitude', drone_id, {
                    'altitude': z
                })
                self.signals.update_signal.emit('velocity', drone_id, {
                    'vx': msg.vx,
                    'vy': msg.vy,
                    'vz': msg.vz
                })

            elif msg_type == "ATTITUDE":
                pitch = math.degrees(msg.pitch)
                self.signals.update_signal.emit('attitude', drone_id, {
                    'pitch': pitch,
                    'roll': 0,
                    'yaw': 0,
                    'rates': (0, 0, 0)
                })

            elif msg_type == "VFR_HUD":
                groundspeed = msg.groundspeed
                heading = msg.heading
                self.signals.update_signal.emit('hud', drone_id, {
                    'heading': heading,
                    'groundspeed': groundspeed
                })
                
        except Exception as e:
            print(f"Error processing MAVLink message from serial: {e}")
            
    def stop(self):
        """停止接收器"""
        self.running = False

class WebSocketMavlinkReceiver(threading.Thread):
    """WebSocket MAVLink 接收器"""
    def __init__(self, url, signals, connection_name):
        super().__init__(daemon=True)
        self.url = url
        self.signals = signals
        self.connection_name = connection_name
        self.running = False
        self.max_retries = 5
        self.base_delay = 1.0
        
    def run(self):
        """執行 WebSocket 接收循環"""
        self.running = True
        asyncio.set_event_loop(asyncio.new_event_loop())
        asyncio.get_event_loop().run_until_complete(self.ws_client_loop())
    
    async def ws_client_loop(self):
        """WebSocket 連接的主循環"""
        retry_count = 0
        
        print(f"Starting WebSocket client for {self.connection_name} at {self.url}")
        
        while self.running and retry_count < self.max_retries:
            try:
                async with websockets.connect(self.url) as websocket:
                    print(f"WebSocket {self.connection_name} connected to {self.url}")
                    retry_count = 0  # 重置重試計數
                    
                    async for message in websocket:
                        if not self.running:
                            break
                            
                        try:
                            data = json.loads(message)
                            data['_connection_source'] = self.connection_name
                            self.process_websocket_message(data)
                        except json.JSONDecodeError as e:
                            print(f"WebSocket {self.connection_name} JSON decode error: {e}")
                        except Exception as e:
                            print(f"WebSocket {self.connection_name} message processing error: {e}")
                            
            except websockets.exceptions.ConnectionClosedError:
                print(f"WebSocket {self.connection_name} connection closed")
                if self.running:
                    retry_count += 1
                    if retry_count < self.max_retries:
                        delay = self.base_delay * (2 ** min(retry_count, 4))
                        print(f"Reconnecting in {delay}s...")
                        await asyncio.sleep(delay)
                else:
                    break
            except Exception as e:
                retry_count += 1
                if retry_count < self.max_retries and self.running:
                    delay = self.base_delay * (2 ** min(retry_count, 4))
                    print(f"WebSocket {self.connection_name} connection error: {e}, retrying in {delay}s (attempt {retry_count}/{self.max_retries})")
                    await asyncio.sleep(delay)
                else:
                    break
        
        print(f"WebSocket client {self.connection_name} stopped")
    
    def process_websocket_message(self, data):
        """處理 WebSocket 訊息"""
        try:
            drone_id = data.get('system_id')
            drone_id = f"s9_{drone_id}" if drone_id else None
            if not drone_id:
                return

            # 模式
            if 'mode' in data:
                self.signals.update_signal.emit('state', drone_id, {
                    'mode': data['mode'],
                })

            # 電池
            if 'battery' in data:
                self.signals.update_signal.emit('battery', drone_id, {
                    'percentage': data['battery']  
                })

            # 位置
            if 'position' in data:
                pos = data['position']
                self.signals.update_signal.emit('gps', drone_id, {
                    'lat': pos.get('lat', 0),
                    'lon': pos.get('lon', 0)
                })

                # Local position - 設定 x, y 為 0.0
                self.signals.update_signal.emit('local_pose', drone_id, {
                    'x': 0.0,
                    'y': 0.0,
                    'z': 0.0
                })
                
                # Altitude - 設定為 0.0
                self.signals.update_signal.emit('altitude', drone_id, {
                    'altitude': 0.0
                })

            # 航向
            if 'heading' in data:
                self.signals.update_signal.emit('hud', drone_id, {
                    'heading': data['heading'],
                    'groundspeed': 0.0
                })
            
        except Exception as e:
            print(f"WebSocket message processing error: {e}")
            
    def stop(self):
        """停止接收器"""
        self.running = False

class DroneMonitor(Node):
    def __init__(self):
        super().__init__('drone_monitor')
        self.signals = DroneSignals()
        self.drone_topics = {}
        self.lock = Lock()
        
        self.arm_clients = {}
        self.takeoff_clients = {}
        self.setpoint_pubs = {}
        self.selected_drones = set()
        self.latest_data = {}

        # 定義需要過濾的模式
        self.filtered_modes = ['Mode(0x000000c0)']

        # WebSocket 接收器列表
        self.ws_receivers = []
        
        # 串口接收器列表
        
        # ================================================================================
        # 【新增】儲存 GPS 資料的字典
        # ================================================================================
        self.drone_gps = {}  # {drone_id: {'lat': ..., 'lon': ..., 'alt': ...}}
        # ================================================================================
        self.serial_receivers = []
        
        # 主题检测定时器
        self.create_timer(1.0, self.scan_topics)
       
    def scan_topics(self):
        topics = self.get_topic_names_and_types()
        drone_pattern = re.compile(r'/MavLinkBus/(s\d+_\d+)/(\w+)')
        
        found_drones = set()
        for topic_name, _ in topics:
            if match := drone_pattern.match(topic_name):
                drone_id, topic_type = match.groups()
                found_drones.add(drone_id)
                with self.lock:
                    self.drone_topics.setdefault(drone_id, set()).add(topic_type)
        
        for drone_id in found_drones:
            if not hasattr(self, f'drone_{drone_id}_subs'):
                self.setup_drone(drone_id)
    
    def setup_drone(self, drone_id):
        base_topic = f'/MavLinkBus/{drone_id}'
        
        # Add service clients
        self.arm_clients[drone_id] = self.create_client(
            CommandBool,
            f'{base_topic}/cmd/arming'
        )
        self.takeoff_clients[drone_id] = self.create_client(
            CommandTOL,
            f'{base_topic}/cmd/takeoff'
        )

        # Add setpoint publisher
        self.setpoint_pubs[drone_id] = self.create_publisher(
            Point,
            f'{base_topic}/setpoint_position/local',
            10
        )

        subs = {
            'attitude': self.create_subscription(
                Imu,
                f'{base_topic}/attitude',
                lambda msg, did=drone_id: self.attitude_callback(did, msg),
                10
            ),
            'battery': self.create_subscription(
                BatteryState,
                f'{base_topic}/battery',
                lambda msg, did=drone_id: self.battery_callback(did, msg),
                10
            ),
            'global': self.create_subscription(
                NavSatFix,
                f'{base_topic}/global_position/global',
                lambda msg, did=drone_id: self.gps_callback(did, msg),
                10
            ),
            'rel_alt': self.create_subscription(
                Float64,
                f'{base_topic}/global_position/rel_alt',
                lambda msg, did=drone_id: self.altitude_callback(did, msg),
                10
            ),
            'local_pose': self.create_subscription(
                Point,
                f'{base_topic}/local_position/pose',
                lambda msg, did=drone_id: self.local_pose_callback(did, msg),
                10
            ),
            'local_vel': self.create_subscription(
                Vector3,
                f'{base_topic}/local_position/velocity',
                lambda msg, did=drone_id: self.local_vel_callback(did, msg),
                10
            ),
            'loss_rate': self.create_subscription(
                Float64,
                f'{base_topic}/packet_loss_rate',
                lambda msg, did=drone_id: self.loss_rate_callback(did, msg),
                10
            ),
            'state': self.create_subscription(
                State,
                f'{base_topic}/state',
                lambda msg, did=drone_id: self.state_callback(did, msg),
                10
            ),
            'ping': self.create_subscription(
                Float64,
                f'{base_topic}/ping',
                lambda msg, did=drone_id: self.ping_callback(did, msg),
                10
            ),
            'vfr_hud': self.create_subscription(
                VfrHud,
                f'{base_topic}/vfr_hud',
                lambda msg, did=drone_id: self.hud_callback(did, msg),
                10
            )
        }
        
        setattr(self, f'drone_{drone_id}_subs', subs)

    async def arm_drone(self, drone_id, arm):
        if drone_id not in self.arm_clients:
            return False
        
        client = self.arm_clients[drone_id]
        if not client.wait_for_service(timeout_sec=1.0):
            return False
            
        request = CommandBool.Request()
        request.value = arm
        
        future = client.call_async(request)
        try:
            response = await future
            return response.success
        except Exception as e:
            self.get_logger().error(f'Arm service call failed: {e}')
            return False

    async def takeoff_drone(self, drone_id, altitude=10.0):
        if drone_id not in self.takeoff_clients:
            return False
            
        client = self.takeoff_clients[drone_id]
        if not client.wait_for_service(timeout_sec=1.0):
            return False
            
        request = CommandTOL.Request()
        request.altitude = altitude
        request.min_pitch = 0.0
        request.yaw = 0.0
        
        future = client.call_async(request)
        try:
            response = await future
            return response.success
        except Exception as e:
            self.get_logger().error(f'Takeoff service call failed: {e}')
            return False

    def send_setpoint(self, drone_id, x, y, z):
        """Send setpoint position command"""
        if drone_id not in self.setpoint_pubs:
            return False
            
        msg = Point()
        msg.x = float(x)
        msg.y = float(y)
        msg.z = float(z)
        
        self.setpoint_pubs[drone_id].publish(msg)
        return True
    
    def quaternion_to_euler(self, q):
        sinr_cosp = 2 * (q.w * q.x + q.y * q.z)
        cosr_cosp = 1 - 2 * (q.x**2 + q.y**2)
        roll = math.atan2(sinr_cosp, cosr_cosp)
        
        sinp = 2 * (q.w * q.y - q.z * q.x)
        pitch = math.asin(sinp) if abs(sinp) < 1 else math.copysign(math.pi/2, sinp)
        
        siny_cosp = 2 * (q.w * q.z + q.x * q.y)
        cosy_cosp = 1 - 2 * (q.y**2 + q.z**2)
        yaw = math.atan2(siny_cosp, cosy_cosp)
        
        return math.degrees(roll), math.degrees(pitch), math.degrees(yaw)

    # callbacks
    def attitude_callback(self, drone_id, msg):
        if hasattr(msg, 'orientation'):
            roll, pitch, yaw = self.quaternion_to_euler(msg.orientation)
            self.latest_data[(drone_id, 'attitude')] = {
                'roll': roll,
                'pitch': pitch,
                'yaw': yaw,
                'rates': (msg.angular_velocity.x, 
                         msg.angular_velocity.y,
                         msg.angular_velocity.z)
            }

    def battery_callback(self, drone_id, msg):
        self.latest_data[(drone_id, 'battery')] = {
            'voltage': msg.voltage
        }

    def state_callback(self, drone_id, msg):
        mode = msg.mode
        if mode in self.filtered_modes:
            return
        self.latest_data[(drone_id, 'state')] = {
            'mode': msg.mode,
            'armed': msg.armed
        }

    def gps_callback(self, drone_id, msg):
        self.latest_data[(drone_id, 'gps')] = {
            'lat': msg.latitude,
            'lon': msg.longitude,
            'alt': msg.altitude
        }
        
        # ================================================================================
        # 【新增】儲存 GPS 資料到 drone_gps 字典
        # ================================================================================
        self.drone_gps[drone_id] = {
            'lat': msg.latitude,
            'lon': msg.longitude,
            'alt': msg.altitude
        }
        # ================================================================================

    def local_vel_callback(self, drone_id, msg):
        self.latest_data[(drone_id, 'velocity')] = {
            'vx': msg.x,
            'vy': msg.y,
            'vz': msg.z
        }

    def altitude_callback(self, drone_id, msg):
        self.latest_data[(drone_id, 'altitude')] = {
            'altitude': msg.data
        }

    def local_pose_callback(self, drone_id, msg):
        self.latest_data[(drone_id, 'local_pose')] = {
            'x': msg.x,
            'y': msg.y,
            'z': msg.z
        }

    def hud_callback(self, drone_id, msg):
        self.latest_data[(drone_id, 'hud')] = {
            'airspeed': msg.airspeed,
            'groundspeed': msg.groundspeed,
            'heading': msg.heading,
            'throttle': msg.throttle,
            'alt': msg.altitude,
            'climb': msg.climb
        }

    def loss_rate_callback(self, drone_id, msg):
        self.latest_data[(drone_id, 'loss_rate')] = {
            'loss_rate': msg.data
        }

    def ping_callback(self, drone_id, msg):
        self.latest_data[(drone_id, 'ping')] = {
            'ping': msg.data
        }

    def start_serial_connection(self, port='/dev/ttyUSB0', baudrate=57600):
        """啟動串口 MAVLink 連接"""
        connection_name = f"Serial_{port.replace('/', '_')}"
        receiver = SerialMavlinkReceiver(port, baudrate, self.signals, connection_name)
        receiver.start()
        self.serial_receivers.append(receiver)
        print(f"Started serial connection on {port} at {baudrate} baud")
        return receiver