#!/usr/bin/env python3
import struct
from pathlib import Path


def progress(current, total):
    percent = current / total
    bar = '#' * int(percent * 50) + '-' * int((1 - percent) * 50)
    print(f'\r进度: [{bar}] {percent * 100:.2f}%', end=' ')


def remove_param_dataflash(input_file, output_file):
    start_marker = b'\xa3\x95\x20'
    next_marker = b'\xa3\x95'

    with open(input_file, 'rb') as f_in, open(output_file, 'wb', buffering=0) as f_out:
        data = f_in.read()
        index = 0
        data_len = len(data)
        written_count = 0
        removed_msgs = 0

        print(f'开始处理文件: {input_file}')
        print(f'文件大小: {data_len} 字节\n')

        while index < data_len:
            progress(index, data_len)

            if data[index : index + len(start_marker)] == start_marker:
                removed_msgs += 1
                index += len(start_marker)
                found = False
                while index < data_len:
                    if data[index : index + len(next_marker)] == next_marker:
                        found = True
                        break
                    index += 1
                if not found:
                    break
            else:
                f_out.write(data[index : index + 1])
                written_count += 1
                index += 1
        f_out.flush()

    print()
    print('-' * 40)
    print(f'处理完成！输出文件: {output_file}')
    print('统计信息:')
    print(f'  删除消息数: {removed_msgs}')
    print(f'原始大小: {data_len} 字节，输出大小: {written_count} 字节')


def parse_mavlink_message(data, start_pos):
    data_len = len(data)
    if start_pos >= data_len:
        return None, None, None

    # 检查起始标记
    if data[start_pos] == 0xFE:  # MAVLink v1
        # v1 消息结构: STX(1) + LEN(1) + SEQ(1) + SYSID(1) + COMPID(1) + MSGID(3) + PAYLOAD + CRC(2)
        if start_pos + 6 > data_len:
            return None, None, None

        payload_len = data[start_pos + 1]
        msg_id = struct.unpack('<I', data[start_pos + 2 : start_pos + 6])[0] & 0xFFFFFF
        total_len = 1 + 1 + 1 + 1 + 1 + 3 + payload_len + 2

        return msg_id, total_len, 1

    elif data[start_pos] == 0xFD:  # MAVLink v2
        # v2 消息结构: STX(1) + LEN(1) + INCOMPAT_FLAGS(1) + COMPAT_FLAGS(1) + SEQ(1) + SYSID(1) + COMPID(1) + MSGID(3) + PAYLOAD + CRC(2)
        if start_pos + 9 > data_len:
            return None, None, None

        payload_len = data[start_pos + 1]
        msg_id = struct.unpack('<I', data[start_pos + 6 : start_pos + 10])[0] & 0xFFFFFF
        total_len = 1 + 1 + 1 + 1 + 1 + 1 + 1 + 3 + payload_len + 2

        return msg_id, total_len, 2

    else:
        return None, None, None


def remove_param_telemetry(input_file, output_file):
    MAVLINK_V1_STX = 0xFE
    MAVLINK_V2_STX = 0xFD
    PARAM_VALUE_MSG_ID = 22  # PARAM_VALUE 消息 ID (v1/v2 相同)

    with open(input_file, 'rb') as infile, open(output_file, 'wb') as outfile:
        data = infile.read()
        index = 0
        data_len = len(data)
        written_count = 0
        total_msgs = 0
        removed_msgs = 0
        v1_count = 0
        v2_count = 0

        print(f'开始处理文件: {input_file}')
        print(f'文件大小: {data_len} 字节\n')

        while index < data_len:
            progress(index, data_len)

            if data[index] not in (MAVLINK_V1_STX, MAVLINK_V2_STX):
                index += 1
                continue

            msg_id, msg_len, version = parse_mavlink_message(data, index)

            if msg_len is None:
                index += 1
                continue

            if index + msg_len > data_len:
                break

            total_msgs += 1
            if version == 1:
                v1_count += 1
            else:
                v2_count += 1

            if msg_id == PARAM_VALUE_MSG_ID:
                removed_msgs += 1
                index += msg_len
                continue

            outfile.write(data[index : index + msg_len])
            written_count += msg_len
            index += msg_len

        print()
        print('-' * 40)
        print(f'处理完成！输出文件: {output_file}')
        print('统计信息:')
        print(f'  总消息数: {total_msgs}')
        print(f'  MAVLink v1 消息: {v1_count}')
        print(f'  MAVLink v2 消息: {v2_count}')
        print(f'  删除消息数: {removed_msgs}')
        print(f'  保留消息数: {total_msgs - removed_msgs}')
        print(f'原始大小: {data_len} 字节，输出大小: {written_count} 字节')


def main(args):
    if args.input.suffix.lower() in ['.bin']:
        remove_param_dataflash(args.input, args.output)
    elif args.input.suffix.lower() in ['.tlog']:
        remove_param_telemetry(args.input, args.output)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument('-i', '--input', type=Path, required=True, help='input file')
    parser.add_argument('-o', '--output', type=Path, required=True, help='output file')

    parser.set_defaults(func=main)
    args = parser.parse_args()
    args.func(args)