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(('', 9002))
    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', 19002)

    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', 19002)

    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(('', 9002))
    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', 19002)

    # 建立與遠端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)

print('Start')
#fdm_switch_example_two()
#mavlink_analyzer_simple()
#fdm_switch_example_one()
fdm_switch_example_one_with_remote_forwarding()
#fdm_parser_example()