import re
import pandas as pd
from dataclasses import dataclass
from typing import Optional


@dataclass
class DjiSrt:
    """解析 DJI SRT 文件中的元数据"""
    FrameCnt: int
    Datetime: int  # 时间戳（毫秒）
    rel_alt: float
    abs_alt: float
    gb_yaw: float
    gb_pitch: float
    gb_roll: float
    dehaze_level: int
    dehaze_mode: int
    iso: Optional[int] = None
    shutter: Optional[str] = None
    fnum: Optional[float] = None
    ev: Optional[int] = None
    ae_meter_md: Optional[int] = None
    focal_len: Optional[float] = None
    dzoom_ratio: Optional[float] = None
    latitude: Optional[float] = None
    longitude: Optional[float] = None


def parse_srt_file(srt_file: str) -> list[DjiSrt]:
    """解析 SRT 文件，返回 DjiSrt 对象列表"""
    data = []
    frame_pattern = re.compile(r"FrameCnt: (\d+) (\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3})")
    rel_abs_alt_pattern = re.compile(r"\[rel_alt: ([\d\.\-]+) abs_alt: ([\d\.\-]+)\]")
    gb_pattern = re.compile(r"\[gb_yaw: ([\d\.\-]+) gb_pitch: ([\d\.\-]+) gb_roll: ([\d\.\-]+)\]")
    dehaze_pattern = re.compile(r"\[dehaze_level: (\d+)\] \[dehaze_mode: (\d+)\]")
    fields_pattern = re.compile(r"\[([a-zA-Z_]+): ([\d\.\-\/]+)\]")  # 匹配其他字段

    with open(srt_file, 'r', encoding='utf-8') as file:
        content = file.read()

    entries = content.split('\n\n')  # 按 SRT 块分割

    for entry in entries:
        if not entry.strip():
            continue

        # 初始化 DjiSrt 对象（先设为 None，后续填充）
        frame_info = DjiSrt(
            FrameCnt=0, Datetime=0, rel_alt=0.0, abs_alt=0.0,
            gb_yaw=0.0, gb_pitch=0.0, gb_roll=0.0,
            dehaze_level=0, dehaze_mode=0
        )

        # 匹配 FrameCnt 和 Datetime
        frame_match = frame_pattern.search(entry)
        if frame_match:
            frame_info.FrameCnt = int(frame_match.group(1))
            datetime_str = frame_match.group(2)
            frame_info.Datetime = int(pd.to_datetime(datetime_str).timestamp() * 1000)

        # 匹配 rel_alt 和 abs_alt
        rel_abs_alt_match = rel_abs_alt_pattern.search(entry)
        if rel_abs_alt_match:
            frame_info.rel_alt = float(rel_abs_alt_match.group(1))
            frame_info.abs_alt = float(rel_abs_alt_match.group(2))

        # 匹配 gb_yaw, gb_pitch, gb_roll
        gb_match = gb_pattern.search(entry)
        if gb_match:
            frame_info.gb_yaw = float(gb_match.group(1))
            frame_info.gb_pitch = float(gb_match.group(2))
            frame_info.gb_roll = float(gb_match.group(3))

        # 匹配 dehaze_level 和 dehaze_mode
        dehaze_match = dehaze_pattern.search(entry)
        if dehaze_match:
            frame_info.dehaze_level = int(dehaze_match.group(1))
            frame_info.dehaze_mode = int(dehaze_match.group(2))

        # 匹配其他字段（iso, shutter, fnum, ev, ae_meter_md, focal_len, dzoom_ratio, latitude, longitude）
        fields_matches = fields_pattern.findall(entry)
        for field, value in fields_matches:
            if field == "iso":
                frame_info.iso = int(value)
            elif field == "shutter":
                frame_info.shutter = value
            elif field == "fnum":
                frame_info.fnum = float(value)
            elif field == "ev":
                frame_info.ev = int(value)
            elif field == "ae_meter_md":
                frame_info.ae_meter_md = int(value)
            elif field == "focal_len":
                frame_info.focal_len = float(value)
            elif field == "dzoom_ratio":
                frame_info.dzoom_ratio = float(value)
            elif field == "latitude":
                frame_info.latitude = float(value)
            elif field == "longitude":
                frame_info.longitude = float(value)

        data.append(frame_info)

    return data


if __name__ == '__main__':
    srt_data = parse_srt_file(r"D:\project\verification\DJI_SRT_Tool-main\example\hongxian.srt")
    for entry in srt_data[:2]:  # 打印前两条数据
        print(entry)