import rclpy
from pymavlink import mavutil
from time import sleep
import time

import socket
import struct

# used at fdm_switch_example_two
import threading
import queue

import json

def mavlink_analyzer_simple(count = 500):
    # rclpy.init()
    # node = rclpy.create_node('mavlink_analyzer_simple')
 
    print('Inintializing connection')

    connection_string="udp:127.0.0.1:14550"
    connection = mavutil.mavlink_connection(connection_string)

    print('Catch messages')
    msg_count = {}

    start_time = time.time()

    for _ in range(count):
        msg = connection.recv_msg()
        # msg = connection.recv_match(blocking=True)
        if msg:
            msg_type = msg.get_type()
            if msg_type in msg_count:
                msg_count[msg_type] += 1
            else:
                msg_count[msg_type] = 1

            if msg.get_type() == 'SERVO_OUTPUT_RAW':
                print(msg)
                pass

        else:
            print('No message yet, 1 second for data to fill')
            sleep(1)

    print('Messages Type received:')
    print(msg_count)

    end_time = time.time()
    print('Time elapsed: ', end_time - start_time)

    print('Closing connection')
    connection.close()


def fdm_parser_example(data=None):
    if data is None:
        data = bytes.fromhex('1a48e903b1340500e803e803e803e803000000000000000000000000000000000000000000000000')
    # 1a48e903b1340500e803e803e803e803000000000000000000000000000000000000000000000000


    parse_format = 'HHI16H'

    if len(data) != struct.calcsize(parse_format):
        print("got packet of len %u, expected %u" % (len(data), struct.calcsize(parse_format)))
        exit(1)

    decoded = struct.unpack(parse_format,data)

    magic = decoded[0]
    frame_rate_hz = decoded[1]
    frame_count = decoded[2]
    pwm = decoded[3:]


    print("magic: %u, frame_rate_hz: %u, frame_count: %u, pwm: %s" % (magic, frame_rate_hz, frame_count, pwm))

def fdm_switch_example_one():

    '''
    這個範例在 SITL 與 Gazebo 之間 做一個簡單的UDP轉發器 
    沒有轉換數據格式
    '''
    # make a link get fdm from SITL, as a server
    sock_s2g = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock_s2g.bind(('', 9012))
    sock_s2g.settimeout(0.1)

    # set a target link pass JSON to Gazebo, as a client
    sock_g2s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server_address = ('127.0.0.1', 19012)

    packet_count = 0
    start_time = time.time()

    while True:
        try:
            # 接到SITL的數據 並轉發給Gazebo
            data_s2g, addr_s2g = sock_s2g.recvfrom(1024)
            sock_g2s.sendto(data_s2g, server_address)
            # 得到Gazebo的回應 並轉發給SITL
            data_g2s, addr_g2s = sock_g2s.recvfrom(1024)
            sock_s2g.sendto(data_g2s, addr_s2g)

            # 解析並保存接收到的數據到 JSON 檔案 (SITL -> Gazebo)
            try:
                # 解析從 SITL 到 Gazebo 的數據
                decoded = struct.unpack('HHI16H', data_s2g)
                s2g_json = {
                    "magic": decoded[0],
                    "frame_rate_hz": decoded[1],
                    "frame_count": decoded[2],
                    "pwm": list(decoded[3:])
                }
                
                # 解析從 Gazebo 到 SITL 的數據
                # 注意：假設從 Gazebo 返回的數據是字串形式的 JSON
                try:
                    g2s_data_str = data_g2s.decode('utf-8')
                    # 保存原始數據，適合後續分析
                    with open('gazebo_response_raw.txt', 'a') as f:
                        f.write(g2s_data_str + '\n')
                except:
                    # 如果不是 UTF-8 格式，可能是二進制數據
                    g2s_data_str = str(data_g2s)
                
                # 合併數據到一個 JSON 對象
                combined_data = {
                    "sitl_to_gazebo": s2g_json,
                    "gazebo_to_sitl": g2s_data_str,
                    "timestamp": time.time()
                }
                
                # 寫入 JSON 檔案
                with open('fdm_one_data.json', 'w') as f:
                    json.dump(combined_data, f, indent=2)
                    
            except Exception as e:
                print(f"JSON 處理錯誤: {e}")

            packet_count += 1
        except socket.timeout:
            print(f'no value')
            #print(f'address:{addr_s2g}')
            #pass

        current_time = time.time()
        elapsed_time = current_time - start_time

        if elapsed_time >= 1.0:
            #print(f"Packets received in the last second: {packet_count}")
            print(f'JSON Pack:{data_g2s}')
            packet_count = 0
            start_time = current_time
            fdm_parser_example(data_s2g)


def fdm_switch_example_two():
    '''
    這個範例在 SITL 與 Gazebo 之間 做一個簡單的UDP轉發器 
    有轉換數據格式

    time_usec
    servo1_raw 
    servo16_raw
    '''
    # read info from SITL via MAVLink
    connection_string="udp:127.0.0.1:14550" # uart mavlink  # MATLAB
    connection = mavutil.mavlink_connection(connection_string)

    # set a target link pass JSON to Gazebo, as a client
    sock_g2s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server_address = ('127.0.0.1', 19012)

    data_queue = queue.Queue()
    servo_out = [0] * 16
    data_queue.put(servo_out)

    def send_data_from_queue(sock, server_address, q, frame_rate_fdm=600, frame_count_fdm=1):
        interval = 1.0 / frame_rate_fdm
        while True:
            start_time = time.time()
            try:
                data = q.get(timeout=0.1)
                if data is None:
                    break
                last_data = data
            except queue.Empty:
                data = last_data
            data = struct.pack('HHI16H', 0x481a, frame_rate_fdm, frame_count_fdm, *servo_out)
            sock.sendto(data, server_address)

            # 新增：寫入 JSON 檔案
            data_json = {
                "magic": 0x481a,
                "frame_rate_hz": frame_rate_fdm,
                "frame_count": frame_count_fdm,
                "pwm": list(servo_out)
            }
            with open('fdm_data.json', 'w') as f:
                json.dump(data_json, f)

            frame_count_fdm += 1
            fdm_parser_example(data) # debug

            elapsed_time = time.time() - start_time
            sleep_time = interval - elapsed_time
            if sleep_time > 0:
                sleep(sleep_time)

    Running = True
    thread = threading.Thread(target=send_data_from_queue, args=(sock_g2s, server_address, data_queue))
    thread.start()

    while Running:
        msgb1 = None
        msg = connection.recv_msg()
        while msg :
            if msg.get_type() == 'SERVO_OUTPUT_RAW':
                msgb1 = msg
                # break # 這裡不需要break，因為我要最後一個 servo_output_raw 的值
            msg = connection.recv_msg()

        if msgb1:
            for i in range(16):
                servo_out[i] = getattr(msgb1, f'servo{i+1}_raw')
                data_queue.put(servo_out)
            msgb1 = None
            # Running = False # debug
        else:
            # print('No message yet, 10 ms for data to fill')
            sleep(0.01)

        # Example of putting data into the queue
        # data_queue.put(data)
        # To stop the thread, you can put None into the queue
        # data_queue.put(None)


def fdm_switch_example_one_with_remote_forwarding():
    '''
    這個範例在 SITL 與 Gazebo 之間 做一個簡單的UDP轉發器 
    並且將JSON數據轉發到指定的遠端IP
    '''
    # make a link get fdm from SITL, as a server
    sock_s2g = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock_s2g.bind(('', 9012))
    sock_s2g.settimeout(0.1)

    # set a target link pass JSON to Gazebo, as a client
    sock_g2s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server_address = ('127.0.0.1', 19012)

    # 建立與遠端IP的socket連接
    remote_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    remote_address = ('140.120.31.130', 9003)  # 遠端IP與埠號，您可以根據需要更改埠號
    print(f'Setting up forwarding to remote address: {remote_address}')

    packet_count = 0
    start_time = time.time()

    while True:
        try:
            # 接到SITL的數據 並轉發給Gazebo
            data_s2g, addr_s2g = sock_s2g.recvfrom(1024)
            sock_g2s.sendto(data_s2g, server_address)
            
            # 得到Gazebo的回應 並轉發給SITL
            data_g2s, addr_g2s = sock_g2s.recvfrom(1024)
            sock_s2g.sendto(data_g2s, addr_s2g)

            # 解析並保存接收到的數據到 JSON 檔案 (SITL -> Gazebo)
            try:
                # 解析從 SITL 到 Gazebo 的數據
                decoded = struct.unpack('HHI16H', data_s2g)
                s2g_json = {
                    "magic": decoded[0],
                    "frame_rate_hz": decoded[1],
                    "frame_count": decoded[2],
                    "pwm": list(decoded[3:])
                }
                
                # 解析從 Gazebo 到 SITL 的數據
                try:
                    g2s_data_str = data_g2s.decode('utf-8')
                    # 保存原始數據，適合後續分析
                    with open('gazebo_response_raw.txt', 'a') as f:
                        f.write(g2s_data_str + '\n')
                except:
                    # 如果不是 UTF-8 格式，可能是二進制數據
                    g2s_data_str = str(data_g2s)
                
                # 合併數據到一個 JSON 對象
                combined_data = {
                    "sitl_to_gazebo": s2g_json,
                    "gazebo_to_sitl": g2s_data_str,
                    "timestamp": time.time()
                }
                
                # 寫入 JSON 檔案
                with open('fdm_one_data.json', 'w') as f:
                    json.dump(combined_data, f, indent=2)
                
                # 將JSON數據轉發到遠端IP
                remote_socket.sendto(json.dumps(combined_data).encode(), remote_address)
                    
            except Exception as e:
                print(f"JSON 處理錯誤: {e}")

            packet_count += 1
        except socket.timeout:
            print(f'no value')
            #pass

        current_time = time.time()
        elapsed_time = current_time - start_time

        if elapsed_time >= 1.0:
            print(f'JSON Pack:{data_g2s}')
            print(f'Forwarded data to remote IP: 140.120.31.130')
            packet_count = 0
            start_time = current_time
            fdm_parser_example(data_s2g)


def fdm_switch_multiple(num_systems=3):
    '''
    Create multiple UDP forwarders between SITL and Gazebo using threading
    Each forwarder uses its own set of ports following the pattern:
    - SITL ports: 9002, 9012, 9022, ...
    - Gazebo ports: 19002, 19012, 19022, ...
    
    Parameters:
    num_systems (int): Number of forwarding pairs to create
    '''
    # Flag to control thread execution
    running = True
    
    def forwarder_thread(system_id):
        sitl_port = 9002 + (system_id * 10)
        gazebo_port = 19002 + (system_id * 10)
        
        # Create SITL socket (server)
        sitl_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sitl_socket.bind(('', sitl_port))
        sitl_socket.settimeout(0.1)
        
        # Create Gazebo socket (client)
        gazebo_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        gazebo_socket.settimeout(0.1)
        gazebo_address = ('127.0.0.1', gazebo_port)
        
        print(f"System {system_id}: Forwarding {sitl_port} → {gazebo_port}")
        
        packet_count = 0
        start_time = time.time()
        last_data_s2g = None
        
        try:
            while running:
                try:
                    # Receive from SITL and forward to Gazebo
                    data_s2g, addr_s2g = sitl_socket.recvfrom(1024)
                    last_data_s2g = data_s2g
                    gazebo_socket.sendto(data_s2g, gazebo_address)
                    
                    try:
                        # Get response from Gazebo and send back to SITL
                        data_g2s, addr_g2s = gazebo_socket.recvfrom(1024)
                        sitl_socket.sendto(data_g2s, addr_s2g)
                        
                        # Process and store data
                        try:
                            # Parse SITL to Gazebo data
                            decoded = struct.unpack('HHI16H', data_s2g)
                            s2g_json = {
                                "magic": decoded[0],
                                "frame_rate_hz": decoded[1],
                                "frame_count": decoded[2],
                                "pwm": list(decoded[3:])
                            }
                            
                            # Try to parse Gazebo to SITL data
                            try:
                                g2s_data_str = data_g2s.decode('utf-8')
                                if packet_count % 50 == 0:
                                    with open(f'gazebo_response_raw_{system_id}.txt', 'a') as f:
                                        f.write(g2s_data_str + '\n')
                            except:
                                g2s_data_str = str(data_g2s)
                            
                            # Combine data
                            combined_data = {
                                "sitl_to_gazebo": s2g_json,
                                "gazebo_to_sitl": g2s_data_str,
                                "timestamp": time.time(),
                                "system_id": system_id,
                                "sitl_port": sitl_port,
                                "gazebo_port": gazebo_port
                            }
                            
                            # Save to JSON file (once per second)
                            current_time = time.time()
                            elapsed_time = current_time - start_time
                            if elapsed_time >= 1.0:
                                with open(f'fdm_data_system_{system_id}.json', 'w') as f:
                                    json.dump(combined_data, f, indent=2)
                                
                        except Exception as e:
                            print(f"System {system_id} - JSON processing error: {e}")
                        
                        packet_count += 1
                        
                    except socket.timeout:
                        # No response from Gazebo, continue
                        pass
                    
                except socket.timeout:
                    # No data from SITL, continue
                    pass
                
                # Print stats periodically
                current_time = time.time()
                elapsed_time = current_time - start_time
                
                if elapsed_time >= 1.0:
                    # print(f"System {system_id} - Packets: {packet_count} in the last second")
                    if packet_count > 0 and last_data_s2g:
                        print(f'System {system_id} - JSON Pack: {data_g2s}')
                        fdm_parser_example(last_data_s2g)
                    packet_count = 0
                    start_time = current_time
        
        except Exception as e:
            print(f"System {system_id} - Error in forwarder thread: {e}")
        finally:
            # Clean up resources
            sitl_socket.close()
            gazebo_socket.close()
            print(f"System {system_id} - Forwarder thread terminated")
    
    # Create and start threads for each system
    threads = []
    for i in range(num_systems):
        thread = threading.Thread(target=forwarder_thread, args=(i,))
        thread.daemon = True
        threads.append(thread)
        thread.start()
    
    # Keep the main thread running
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\nShutting down all forwarders...")
        running = False
        time.sleep(1)
        print("Done.")
        return


def fdm_switch_multiple_with_remote(num_systems=3, remote_ip='140.120.31.130'):
    '''
    Create multiple UDP forwarders between SITL and Gazebo using threading
    Also forwards the combined data to a remote IP address
    
    Each forwarder uses its own set of ports following the pattern:
    - SITL ports: 9002, 9012, 9022, ...
    - Gazebo ports: 19002, 19012, 19022, ...
    - Remote ports: 9003, 9013, 9023, ... (for the remote forwarding)
    
    Parameters:
    num_systems (int): Number of forwarding pairs to create
    remote_ip (str): IP address for remote forwarding
    '''
    # Flag to control thread execution
    running = True

    # Add shared counters for total packets
    total_packets_sent = 0
    total_packets_lock = threading.Lock()  # For thread safety    
    
    def forwarder_thread(system_id):
        nonlocal total_packets_sent
        sitl_port = 9002 + (system_id * 10)
        gazebo_port = 19002 + (system_id * 10)
        remote_port = 9003 + (system_id * 10)
        
        # Create SITL socket (server)
        sitl_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sitl_socket.bind(('', sitl_port))
        sitl_socket.settimeout(0.1)
        
        # Create Gazebo socket (client)
        gazebo_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        gazebo_socket.settimeout(0.1)
        gazebo_address = ('127.0.0.1', gazebo_port)
        
        # Create remote forwarding socket
        remote_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        remote_address = (remote_ip, remote_port)
        
        print(f"System {system_id}: SITL:{sitl_port} → Gazebo:{gazebo_port}, Remote:{remote_ip}:{remote_port}")
        
        packet_count = 0
        start_time = time.time()
        last_data_s2g = None
        forward_counter = 0  # 添加轉發計數器
        system_packets_sent = 0  # Local counter for this system

        try:
            while running:
                try:
                    # Receive from SITL and forward to Gazebo
                    data_s2g, addr_s2g = sitl_socket.recvfrom(1024)
                    last_data_s2g = data_s2g
                    gazebo_socket.sendto(data_s2g, gazebo_address)
                    
                    try:
                        # Get response from Gazebo and send back to SITL
                        data_g2s, addr_g2s = gazebo_socket.recvfrom(1024)
                        sitl_socket.sendto(data_g2s, addr_s2g)
                        
                        # Process and store data
                        try:
                            # Parse SITL to Gazebo data
                            decoded = struct.unpack('HHI16H', data_s2g)
                            s2g_json = {
                                "magic": decoded[0],
                                "frame_rate_hz": decoded[1],
                                "frame_count": decoded[2],
                                "pwm": list(decoded[3:])
                            }
                            
                            # Try to parse Gazebo to SITL data
                            try:
                                g2s_data_str = data_g2s.decode('utf-8')
                                if packet_count % 50 == 0:
                                    with open(f'gazebo_response_raw_{system_id}.txt', 'a') as f:
                                        f.write(g2s_data_str + '\n')
                            except:
                                g2s_data_str = str(data_g2s)
                            
                            # Combine data
                            combined_data = {
                                "sitl_to_gazebo": s2g_json,
                                "gazebo_to_sitl": g2s_data_str,
                                "timestamp": time.time(),
                                "system_id": system_id,
                                "sitl_port": sitl_port,
                                "gazebo_port": gazebo_port
                            }
                            
                            '''
                            # 每n筆才轉發一次
                            forward_counter += 1
                            if forward_counter >= 18:
                                # Forward data to remote server
                                remote_socket.sendto(json.dumps(combined_data).encode(), remote_address)
                                #print(f"System {system_id} - Send data to remote (第{packet_count}個包)")
                                forward_counter = 0  # 重置計數器
                            '''

                            # 使用模運算代替計數器重置
                            forward_counter += 1
                            if packet_count % 1 == 0:
                                # Forward data to remote server
                                remote_socket.sendto(json.dumps(combined_data).encode(), remote_address) 
                                system_packets_sent += 1
                                with total_packets_lock:
                                    total_packets_sent += 1                                
                                # Debug用，看每秒傳了多少筆出去
                                #print(f"System {system_id} - Send data to remote (第{packet_count}個包)")               
                            
                            # Save to JSON file (once per second)
                            current_time = time.time()
                            elapsed_time = current_time - start_time
                            if elapsed_time >= 1.0:
                                with open(f'fdm_data_system_{system_id}.json', 'w') as f:
                                    json.dump(combined_data, f, indent=2)
                                
                        except Exception as e:
                            print(f"System {system_id} - JSON processing error: {e}")
                        
                        packet_count += 1
                        
                    except socket.timeout:
                        # No response from Gazebo, continue
                        pass
                    
                except socket.timeout:
                    # No data from SITL, continue
                    pass
                
                # Print stats periodically
                current_time = time.time()
                elapsed_time = current_time - start_time
                
                if elapsed_time >= 1.0:
                    print(f"System {system_id} - Packets: {packet_count} in the last second")
                    if packet_count > 0 and last_data_s2g:
                        fdm_parser_example(last_data_s2g)
                    packet_count = 0
                    start_time = current_time
        
        except Exception as e:
            print(f"System {system_id} - Error in forwarder thread: {e}")
        finally:
            # Clean up resources
            sitl_socket.close()
            gazebo_socket.close()
            remote_socket.close()
            print(f"System {system_id} - Forwarder thread terminated")

    # Create statistics display thread
    def stats_display_thread():
        start_time = time.time()
        while running:
            time.sleep(5)  # Display stats every 5 seconds
            elapsed = time.time() - start_time
            with total_packets_lock:
                rate = total_packets_sent / elapsed if elapsed > 0 else 0
                print(f"\n===== STATISTICS =====")
                print(f"Total packets sent: {total_packets_sent}")
                print(f"Average rate: {rate:.2f} packets/sec")
                print(f"Running time: {elapsed:.1f} seconds")
                print(f"======================\n")

    # Create and start threads for each system
    threads = []
    for i in range(num_systems):
        thread = threading.Thread(target=forwarder_thread, args=(i,))
        thread.daemon = True
        threads.append(thread)
        thread.start()

    # Start stats display thread
    stats_thread = threading.Thread(target=stats_display_thread)
    stats_thread.daemon = True
    stats_thread.start()        
    
    # Keep the main thread running
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\nShutting down all forwarders...")
        running = False
        time.sleep(1)
        print(f"Final total packets sent: {total_packets_sent}")
        print("Done.")
        return


print('Start')
#fdm_switch_example_two()
#mavlink_analyzer_simple()
#fdm_switch_example_one()
#fdm_switch_example_one_with_remote_forwarding()
#fdm_switch_multiple(1)  # Create 3 forwarding pairs
fdm_switch_multiple_with_remote(3)  # Create 1 forwarding pair with remote forwarding
#fdm_parser_example()