import sys
import os
from dataclasses import dataclass
from typing import Optional, Tuple
import numpy as np

from PyQt5 import QtCore, QtGui, QtWidgets


# -----------------------------
# 数据模型与枚举
# -----------------------------

class PixelFormat:
    # Packed YUV422
    YUYV = "YUYV"  # Y0 U0 Y1 V0
    YVYU = "YVYU"  # Y0 V0 Y1 U0
    UYVY = "UYVY"  # U0 Y0 V0 Y1
    VYUY = "VYUY"  # V0 Y0 U0 Y1

    # Semi-planar / planar variants
    NV12 = "NV12"  # Y + interleaved UV (4:2:0), UV order = U,V
    NV21 = "NV21"  # Y + interleaved VU (4:2:0), UV order = V,U
    NV16 = "NV16"  # Y + interleaved UV (4:2:2), UV order = U,V
    NV61 = "NV61"  # Y + interleaved VU (4:2:2), UV order = V,U

    GRAY = "GRAY"  # Only Y

    # YUV444
    YUV444_PLANAR = "YUV444_PLANAR"  # YYYY... UUUU... VVVV...
    YUV444_INTERLEAVED_UV = "YUV444_INTERLEAVED_UV"  # Y U V Y U V ... (UV顺序=U,V)
    YUV444_INTERLEAVED_VU = "YUV444_INTERLEAVED_VU"  # Y V U Y V U ... (UV顺序=V,U)
    
    # RGB/BGR formats
    RGB565 = "RGB565"  # 16-bit RGB: RRRRR GGGGGG BBBBB
    BGR565 = "BGR565"  # 16-bit BGR: BBBBB GGGGGG RRRRR
    RGB888 = "RGB888"  # 24-bit RGB: R G B
    BGR888 = "BGR888"  # 24-bit BGR: B G R
    RGB8888 = "RGB8888"  # 32-bit RGB: R G B (alpha ignored)
    BGR8888 = "BGR8888"  # 32-bit BGR: B G R (alpha ignored)
    NRGB8888 = "NRGB8888"  # 32-bit native RGB (no alpha)
    NBGR8888 = "NBGR8888"  # 32-bit native BGR (no alpha)


SUPPORTED_FORMATS = [
    PixelFormat.YUYV, PixelFormat.YVYU, PixelFormat.UYVY, PixelFormat.VYUY,
    PixelFormat.NV12, PixelFormat.NV21, PixelFormat.NV16, PixelFormat.NV61,
    PixelFormat.GRAY,
    PixelFormat.YUV444_PLANAR, PixelFormat.YUV444_INTERLEAVED_UV, PixelFormat.YUV444_INTERLEAVED_VU,
    PixelFormat.RGB565, PixelFormat.BGR565,
    PixelFormat.RGB888, PixelFormat.BGR888,
    PixelFormat.RGB8888, PixelFormat.BGR8888,
    PixelFormat.NRGB8888, PixelFormat.NBGR8888
]

# 预定义的尺寸选项
PREDEFINED_SIZES = [
    "640x480 (VGA)",
    "800x600 (SVGA)",
    "1024x768 (XGA)",
    "1280x720 (HD)",
    "1920x1080 (Full HD)",
    "2560x1440 (2K)",
    "3840x2160 (4K)",
    "自定义..."
]


@dataclass
class ImageSpec:
    width: int
    height: int
    fmt: str


# -----------------------------
# YUV -> RGB 转换工具
# -----------------------------

def yuv_to_rgb_numpy(Y, U, V):
    """
    Y, U, V: float32 ndarray in [0..255], same shape.
    Returns RGB uint8 ndarray with shape (..., 3).
    BT.601 full-range近似：
    R = Y + 1.402*(V-128)
    G = Y - 0.344136*(U-128) - 0.714136*(V-128)
    B = Y + 1.772*(U-128)
    """
    Yf = Y.astype(np.float32)
    Uf = U.astype(np.float32) - 128.0
    Vf = V.astype(np.float32) - 128.0

    R = Yf + 1.402 * Vf
    G = Yf - 0.344136 * Uf - 0.714136 * Vf
    B = Yf + 1.772 * Uf

    RGB = np.stack([R, G, B], axis=-1)
    np.clip(RGB, 0, 255, out=RGB)
    return RGB.astype(np.uint8)


def unpack_packed_422_to_planes(data: np.ndarray, width: int, height: int, order: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    将打包的YUV422（每2像素4字节）解成等大Y平面和下采样的U/V再上采样到全分辨率。
    order: "YUYV", "YVYU", "UYVY", "VYUY"
    返回 Y, U_full, V_full
    """
    # 每行应有 width*2 字节
    expected = width * height * 2
    if data.size != expected:
        raise ValueError(f"数据长度不匹配，期望 {expected} 字节，实际 {data.size} 字节")

    # 先 reshape 成 (H, W*2)
    row = data.reshape((height, width * 2))
    # 两像素为一组的4字节块
    blocks = row.reshape((height, width // 2, 4))  # shape (H, W/2, 4)

    if order == PixelFormat.YUYV:
        Y0 = blocks[:, :, 0]
        U0 = blocks[:, :, 1]
        Y1 = blocks[:, :, 2]
        V0 = blocks[:, :, 3]
    elif order == PixelFormat.YVYU:
        Y0 = blocks[:, :, 0]
        V0 = blocks[:, :, 1]
        Y1 = blocks[:, :, 2]
        U0 = blocks[:, :, 3]
    elif order == PixelFormat.UYVY:
        U0 = blocks[:, :, 0]
        Y0 = blocks[:, :, 1]
        V0 = blocks[:, :, 2]
        Y1 = blocks[:, :, 3]
    elif order == PixelFormat.VYUY:
        V0 = blocks[:, :, 0]
        Y0 = blocks[:, :, 1]
        U0 = blocks[:, :, 2]
        Y1 = blocks[:, :, 3]
    else:
        raise ValueError("不支持的打包422顺序")

    # 还原Y为 (H, W)
    Y = np.empty((height, width), dtype=np.uint8)
    Y[:, 0::2] = Y0
    Y[:, 1::2] = Y1

    # U/V 为每2像素共享，横向采样因子2，纵向不变 -> 需要水平上采样到全分辨率
    U = np.repeat(U0, 2, axis=1)
    V = np.repeat(V0, 2, axis=1)
    return Y, U, V


def parse_nvxx(data: np.ndarray, width: int, height: int, chroma_order: str, is_420: bool) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    NV12/NV21 (420) 或 NV16/NV61 (422) 半平面解析。
    chroma_order: "UV" or "VU"
    is_420: True for 4:2:0, False for 4:2:2
    返回 Y, U_full, V_full
    """
    if is_420:
        y_size = width * height
        c_h = height // 2
        c_size = width * c_h  # interleaved 2 channels across width
        expected = y_size + c_size
        if data.size != expected:
            raise ValueError(f"数据长度不匹配(420)，期望 {expected} 字节，实际 {data.size} 字节")
        Y = data[:y_size].reshape((height, width))
        UV = data[y_size:].reshape((c_h, width))
        # UV 行内是交错: U,V,U,V,...
        U_sub = UV[:, 0::2]
        V_sub = UV[:, 1::2]
        if chroma_order == "VU":
            U_sub, V_sub = V_sub, U_sub
        # 上采样到全分辨率（双线性可选，这里用最近邻）
        U = np.repeat(np.repeat(U_sub, 2, axis=0), 2, axis=1)
        V = np.repeat(np.repeat(V_sub, 2, axis=0), 2, axis=1)
    else:
        # 422: 色度高方向与亮度同高，水平方向为1/2
        y_size = width * height
        c_h = height
        c_size = width * c_h // 1 // 1 // 2  # 每行交错UV对，宽度一样但每2像素一组 → 列数等于宽
        expected = y_size + (width * c_h)  # 因为UV交错，每行宽度与Y一致
        if data.size != expected:
            raise ValueError(f"数据长度不匹配(422)，期望 {expected} 字节，实际 {data.size} 字节")
        Y = data[:y_size].reshape((height, width))
        UV = data[y_size:].reshape((c_h, width))
        U_sub = UV[:, 0::2]
        V_sub = UV[:, 1::2]
        if chroma_order == "VU":
            U_sub, V_sub = V_sub, U_sub
        U = np.repeat(U_sub, 2, axis=1)
        V = np.repeat(V_sub, 2, axis=1)
    return Y, U, V


def parse_gray(data: np.ndarray, width: int, height: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    expected = width * height
    if data.size != expected:
        raise ValueError(f"数据长度不匹配(GRAY)，期望 {expected} 字节，实际 {data.size} 字节")
    Y = data.reshape((height, width))
    U = np.full_like(Y, 128, dtype=np.uint8)
    V = np.full_like(Y, 128, dtype=np.uint8)
    return Y, U, V


def parse_yuv444_planar(data: np.ndarray, width: int, height: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    plane_size = width * height
    expected = plane_size * 3
    if data.size != expected:
        raise ValueError(f"数据长度不匹配(YUV444 planar)，期望 {expected} 字节，实际 {data.size} 字节")
    Y = data[0:plane_size].reshape((height, width))
    U = data[plane_size:2*plane_size].reshape((height, width))
    V = data[2*plane_size:3*plane_size].reshape((height, width))
    return Y, U, V


def parse_yuv444_interleaved(data: np.ndarray, width: int, height: int, uv_order: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    交错式：按像素 Y U V 或 Y V U 排列（总步长3字节）。
    uv_order: "UV" 或 "VU"
    """
    expected = width * height * 3
    if data.size != expected:
        raise ValueError(f"数据长度不匹配(YUV444 interleaved)，期望 {expected} 字节，实际 {data.size} 字节")
    pix = data.reshape((height, width, 3))
    Y = pix[:, :, 0]
    if uv_order == "UV":
        U = pix[:, :, 1]
        V = pix[:, :, 2]
    else:
        V = pix[:, :, 1]
        U = pix[:, :, 2]
    return Y, U, V


# -----------------------------
# RGB/BGR 格式解析函数
# -----------------------------

def parse_rgb565(data: np.ndarray, width: int, height: int, is_bgr: bool = False) -> np.ndarray:
    """
    解析RGB565或BGR565格式
    每个像素2字节：RRRRR GGGGGG BBBBB (16位)
    返回RGB uint8数组 (H, W, 3)
    """
    expected = width * height * 2
    if data.size != expected:
        raise ValueError(f"数据长度不匹配(RGB565/BGR565)，期望 {expected} 字节，实际 {data.size} 字节")
    
    # 将数据重新解释为16位整数（小端序）
    pixels = data.view(dtype=np.uint16).reshape((height, width))
    
    # 提取RGB分量
    if is_bgr:
        # BGR565: BBBBB GGGGGG RRRRR
        b = ((pixels >> 11) & 0x1F) << 3  # 5-bit to 8-bit
        g = ((pixels >> 5) & 0x3F) << 2   # 6-bit to 8-bit
        r = (pixels & 0x1F) << 3          # 5-bit to 8-bit
    else:
        # RGB565: RRRRR GGGGGG BBBBB
        r = ((pixels >> 11) & 0x1F) << 3
        g = ((pixels >> 5) & 0x3F) << 2
        b = (pixels & 0x1F) << 3
    
    # 组合成RGB图像
    rgb = np.stack([r, g, b], axis=-1).astype(np.uint8)
    return rgb


def parse_rgb888(data: np.ndarray, width: int, height: int, is_bgr: bool = False) -> np.ndarray:
    """
    解析RGB888或BGR888格式
    每个像素3字节：R G B 或 B G R
    返回RGB uint8数组 (H, W, 3)
    """
    expected = width * height * 3
    if data.size != expected:
        raise ValueError(f"数据长度不匹配(RGB888/BGR888)，期望 {expected} 字节，实际 {data.size} 字节")
    
    # 重塑为(H, W, 3)
    pixels = data.reshape((height, width, 3))
    
    if is_bgr:
        # BGR -> RGB
        rgb = np.zeros_like(pixels)
        rgb[:, :, 0] = pixels[:, :, 2]  # R
        rgb[:, :, 1] = pixels[:, :, 1]  # G
        rgb[:, :, 2] = pixels[:, :, 0]  # B
        return rgb
    else:
        # 已经是RGB格式
        return pixels


def parse_rgb8888(data: np.ndarray, width: int, height: int, is_bgr: bool = False, is_native: bool = False) -> np.ndarray:
    """
    解析32位RGB/BGR格式
    每个像素4字节：R G B (alpha ignored) 或 B G R (alpha ignored)
    或NRGB8888/NBGR8888（native格式，alpha通道可能位于不同位置）
    返回RGB uint8数组 (H, W, 3)
    """
    expected = width * height * 4
    if data.size != expected:
        raise ValueError(f"数据长度不匹配(32-bit RGB/BGR)，期望 {expected} 字节，实际 {data.size} 字节")
    
    # 重塑为(H, W, 4)
    pixels = data.reshape((height, width, 4))
    
    if is_native:
        # 对于NRGB8888/NBGR8888，格式可能为 A R G B 或 R G B A，取决于系统字节序
        # 简单起见，我们假设前三个字节是RGB/BGR分量（忽略alpha）
        if is_bgr:
            # NBGR8888: B G R (alpha)
            rgb = np.zeros((height, width, 3), dtype=np.uint8)
            rgb[:, :, 0] = pixels[:, :, 2]  # R
            rgb[:, :, 1] = pixels[:, :, 1]  # G
            rgb[:, :, 2] = pixels[:, :, 0]  # B
        else:
            # NRGB8888: R G B (alpha)
            rgb = pixels[:, :, :3].copy()
    else:
        # RGB8888或BGR8888
        if is_bgr:
            # BGR8888: B G R (alpha)
            rgb = np.zeros((height, width, 3), dtype=np.uint8)
            rgb[:, :, 0] = pixels[:, :, 2]  # R
            rgb[:, :, 1] = pixels[:, :, 1]  # G
            rgb[:, :, 2] = pixels[:, :, 0]  # B
        else:
            # RGB8888: R G B (alpha)
            rgb = pixels[:, :, :3].copy()
    
    return rgb


def decode_raw_frame(raw: bytes, spec: ImageSpec, yuv444_uv_order: str = "UV") -> np.ndarray:
    """
    将原始帧数据按 spec 解码为 RGB np.ndarray (H, W, 3), uint8
    yuv444_uv_order: 针对 YUV444_INTERLEAVED 时 UV顺序选择 "UV" 或 "VU"
    """
    width, height, fmt = spec.width, spec.height, spec.fmt
    arr = np.frombuffer(raw, dtype=np.uint8)

    # YUV格式解析
    if fmt in (PixelFormat.YUYV, PixelFormat.YVYU, PixelFormat.UYVY, PixelFormat.VYUY):
        Y, U, V = unpack_packed_422_to_planes(arr, width, height, fmt)
        rgb = yuv_to_rgb_numpy(Y, U, V)
    elif fmt == PixelFormat.NV12:
        Y, U, V = parse_nvxx(arr, width, height, "UV", is_420=True)
        rgb = yuv_to_rgb_numpy(Y, U, V)
    elif fmt == PixelFormat.NV21:
        Y, U, V = parse_nvxx(arr, width, height, "VU", is_420=True)
        rgb = yuv_to_rgb_numpy(Y, U, V)
    elif fmt == PixelFormat.NV16:
        Y, U, V = parse_nvxx(arr, width, height, "UV", is_420=False)
        rgb = yuv_to_rgb_numpy(Y, U, V)
    elif fmt == PixelFormat.NV61:
        Y, U, V = parse_nvxx(arr, width, height, "VU", is_420=False)
        rgb = yuv_to_rgb_numpy(Y, U, V)
    elif fmt == PixelFormat.GRAY:
        Y, U, V = parse_gray(arr, width, height)
        rgb = yuv_to_rgb_numpy(Y, U, V)
    elif fmt == PixelFormat.YUV444_PLANAR:
        Y, U, V = parse_yuv444_planar(arr, width, height)
        rgb = yuv_to_rgb_numpy(Y, U, V)
    elif fmt == PixelFormat.YUV444_INTERLEAVED_UV:
        Y, U, V = parse_yuv444_interleaved(arr, width, height, "UV")
        rgb = yuv_to_rgb_numpy(Y, U, V)
    elif fmt == PixelFormat.YUV444_INTERLEAVED_VU:
        Y, U, V = parse_yuv444_interleaved(arr, width, height, "VU")
        rgb = yuv_to_rgb_numpy(Y, U, V)
    
    # RGB/BGR格式解析
    elif fmt == PixelFormat.RGB565:
        rgb = parse_rgb565(arr, width, height, is_bgr=False)
    elif fmt == PixelFormat.BGR565:
        rgb = parse_rgb565(arr, width, height, is_bgr=True)
    elif fmt == PixelFormat.RGB888:
        rgb = parse_rgb888(arr, width, height, is_bgr=False)
    elif fmt == PixelFormat.BGR888:
        rgb = parse_rgb888(arr, width, height, is_bgr=True)
    elif fmt == PixelFormat.RGB8888:
        rgb = parse_rgb8888(arr, width, height, is_bgr=False, is_native=False)
    elif fmt == PixelFormat.BGR8888:
        rgb = parse_rgb8888(arr, width, height, is_bgr=True, is_native=False)
    elif fmt == PixelFormat.NRGB8888:
        rgb = parse_rgb8888(arr, width, height, is_bgr=False, is_native=True)
    elif fmt == PixelFormat.NBGR8888:
        rgb = parse_rgb8888(arr, width, height, is_bgr=True, is_native=True)
    else:
        raise ValueError(f"不支持的像素格式: {fmt}")

    return rgb


# -----------------------------
# QImage 显示与栅格绘制
# -----------------------------

class GridPixmap(QtGui.QPixmap):
    def __init__(self, base_image: QtGui.QImage, show_grid: bool, grid_step: int, grid_color: QtGui.QColor):
        super().__init__(QtGui.QPixmap.fromImage(base_image))
        if show_grid and grid_step > 0:
            painter = QtGui.QPainter(self)
            pen = QtGui.QPen(grid_color)
            pen.setWidth(1)
            painter.setPen(pen)
            w = self.width()
            h = self.height()
            # 画垂直线
            x = 0
            while x < w:
                painter.drawLine(x, 0, x, h)
                x += grid_step
            # 画水平线
            y = 0
            while y < h:
                painter.drawLine(0, y, w, y)
                y += grid_step
            painter.end()


def numpy_to_qimage(rgb: np.ndarray) -> QtGui.QImage:
    h, w, _ = rgb.shape
    bytes_per_line = 3 * w
    # 注意：必须保证连续内存
    if not rgb.flags['C_CONTIGUOUS']:
        rgb = np.ascontiguousarray(rgb)
    return QtGui.QImage(rgb.data, w, h, bytes_per_line, QtGui.QImage.Format_RGB888).copy()


# -----------------------------
# 主窗口与UI
# -----------------------------

class ImageViewer(QtWidgets.QLabel):
    def __init__(self):
        super().__init__()
        self.setAlignment(QtCore.Qt.AlignCenter)
        self.setBackgroundRole(QtGui.QPalette.Base)
        self.setSizePolicy(QtWidgets.QSizePolicy.Ignored, QtWidgets.QSizePolicy.Ignored)
        # 修改：设置为不缩放，保持原始大小
        self.setScaledContents(False)

    def set_pixmap(self, pixmap: QtGui.QPixmap):
        self.setPixmap(pixmap)
        # 调整窗口以适应图片大小（可选）
        self.resize(pixmap.size())


class MainWindow(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Raw YUV Viewer")
        self.resize(1000, 700)

        # 创建一个滚动区域来容纳图像查看器（当图像太大时）
        self.scroll_area = QtWidgets.QScrollArea()
        self.scroll_area.setWidgetResizable(True)
        self.scroll_area.setAlignment(QtCore.Qt.AlignCenter)

        self.viewer = ImageViewer()
        self.scroll_area.setWidget(self.viewer)
        self.setCentralWidget(self.scroll_area)

        # 当前状态
        self.image_spec = ImageSpec(640, 480, PixelFormat.NV12)
        self.yuv444_uv_order = "UV"  # 仅对 YUV444 interleaved 使用
        self.current_file: Optional[str] = None
        self.current_rgb: Optional[np.ndarray] = None

        # 栅格设置
        self.show_grid = False
        self.grid_step = 8
        self.grid_color = QtGui.QColor(0, 255, 0, 120)

        self._create_actions()
        self._create_menus()

        # 启用拖放
        self.setAcceptDrops(True)

        # 状态栏
        self.statusBar().showMessage("拖入原始帧文件，或 文件 -> 打开")

    # ---- 菜单和动作 ----

    def _create_actions(self):
        self.act_open = QtWidgets.QAction("打开...", self)
        self.act_open.setShortcut("Ctrl+O")
        self.act_open.triggered.connect(self.open_file_dialog)

        self.act_quit = QtWidgets.QAction("退出", self)
        self.act_quit.setShortcut("Ctrl+Q")
        self.act_quit.triggered.connect(self.close)

        self.act_toggle_grid = QtWidgets.QAction("显示栅格", self, checkable=True)
        self.act_toggle_grid.triggered.connect(self.toggle_grid)

        self.act_set_grid_step = QtWidgets.QAction("设置栅格步长...", self)
        self.act_set_grid_step.triggered.connect(self.set_grid_step)

        self.act_rerender = QtWidgets.QAction("重新渲染", self)
        self.act_rerender.setShortcut("F5")
        self.act_rerender.triggered.connect(self.rerender)

        self.act_fit_to_window = QtWidgets.QAction("适应窗口大小", self, checkable=True)
        self.act_fit_to_window.triggered.connect(self.toggle_fit_to_window)

        self.act_actual_size = QtWidgets.QAction("实际大小", self)
        self.act_actual_size.triggered.connect(self.show_actual_size)

    def _create_menus(self):
        m_file = self.menuBar().addMenu("文件")
        m_file.addAction(self.act_open)
        m_file.addSeparator()
        m_file.addAction(self.act_quit)

        m_settings = self.menuBar().addMenu("设置")
        
        # 创建尺寸子菜单
        m_size = m_settings.addMenu("图像尺寸")
        for size_text in PREDEFINED_SIZES:
            if size_text == "自定义...":
                m_size.addSeparator()
                act_custom = QtWidgets.QAction("自定义...", self)
                act_custom.triggered.connect(self.set_custom_size)
                m_size.addAction(act_custom)
            else:
                act = QtWidgets.QAction(size_text, self)
                act.triggered.connect(lambda checked, s=size_text: self.set_predefined_size(s))
                m_size.addAction(act)
        
        # 创建格式子菜单
        m_format = m_settings.addMenu("像素格式")
        for fmt in SUPPORTED_FORMATS:
            act = QtWidgets.QAction(fmt, self, checkable=True)
            act.setChecked(fmt == self.image_spec.fmt)
            act.triggered.connect(lambda checked, f=fmt: self.set_pixel_format(f))
            m_format.addAction(act)
        
        m_settings.addSeparator()
        m_settings.addAction(self.act_rerender)

        m_view = self.menuBar().addMenu("视图")
        m_view.addAction(self.act_toggle_grid)
        m_view.addAction(self.act_set_grid_step)
        m_view.addSeparator()
        m_view.addAction(self.act_fit_to_window)
        m_view.addAction(self.act_actual_size)

    # ---- 事件：拖放 ----

    def dragEnterEvent(self, event: QtGui.QDragEnterEvent):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()

    def dropEvent(self, event: QtGui.QDropEvent):
        urls = event.mimeData().urls()
        if not urls:
            return
        path = urls[0].toLocalFile()
        if path:
            self.load_file(path)

    # ---- 打开/加载/渲染 ----

    def open_file_dialog(self):
        path, _ = QtWidgets.QFileDialog.getOpenFileName(self, "选择原始帧文件", "", "All Files (*)")
        if path:
            self.load_file(path)

    def load_file(self, path: str):
        try:
            with open(path, "rb") as f:
                raw = f.read()
            self.current_file = path
            self.statusBar().showMessage(f"已加载: {os.path.basename(path)} ({len(raw)} bytes)")
            self.current_rgb = self.decode_with_current_spec(raw)
            self.update_view()
        except Exception as e:
            QtWidgets.QMessageBox.critical(self, "错误", str(e))

    def decode_with_current_spec(self, raw: bytes) -> np.ndarray:
        # 对 YUV444_INTERLEAVED_*，uv 顺序已在格式名中决定；也保留一个变量用于未来扩展
        return decode_raw_frame(raw, self.image_spec, self.yuv444_uv_order)

    def update_view(self):
        if self.current_rgb is None:
            self.viewer.setPixmap(QtGui.QPixmap())
            return
        qimg = numpy_to_qimage(self.current_rgb)
        pm = GridPixmap(qimg, self.show_grid, self.grid_step, self.grid_color)
        self.viewer.set_pixmap(pm)
        
        # 更新窗口标题显示当前信息
        self.setWindowTitle(f"Raw YUV Viewer - {self.image_spec.width}x{self.image_spec.height} {self.image_spec.fmt}")

    def rerender(self):
        if not self.current_file:
            return
        try:
            with open(self.current_file, "rb") as f:
                raw = f.read()
            self.current_rgb = self.decode_with_current_spec(raw)
            self.update_view()
        except Exception as e:
            QtWidgets.QMessageBox.critical(self, "错误", str(e))

    # ---- 设置项 ----

    def set_predefined_size(self, size_text: str):
        """设置预定义的尺寸"""
        size_part = size_text.split()[0]
        w, h = map(int, size_part.split('x'))
        self.image_spec = ImageSpec(w, h, self.image_spec.fmt)
        self.statusBar().showMessage(f"已设置尺寸: {w}x{h}")
        self.rerender()
        
        # 更新格式菜单的选中状态
        self.update_format_menu_checked()

    def set_custom_size(self):
        """设置自定义尺寸"""
        w, ok1 = QtWidgets.QInputDialog.getInt(self, "图像宽度", "宽 (pixels):", self.image_spec.width, 1, 16384, 1)
        if not ok1:
            return
        h, ok2 = QtWidgets.QInputDialog.getInt(self, "图像高度", "高 (pixels):", self.image_spec.height, 1, 16384, 1)
        if not ok2:
            return
        self.image_spec = ImageSpec(w, h, self.image_spec.fmt)
        self.statusBar().showMessage(f"已设置尺寸: {w}x{h}")
        self.rerender()

    def set_pixel_format(self, fmt: str):
        """设置像素格式"""
        self.image_spec = ImageSpec(self.image_spec.width, self.image_spec.height, fmt)
        self.statusBar().showMessage(f"已设置像素格式: {fmt}")
        self.rerender()
        
        # 更新格式菜单的选中状态
        self.update_format_menu_checked()

    def update_format_menu_checked(self):
        """更新格式菜单的选中状态"""
        format_menu = self.menuBar().actions()[1].menu().actions()[1].menu()  # 获取格式子菜单
        for action in format_menu.actions():
            if action.text() == self.image_spec.fmt:
                action.setChecked(True)
            else:
                action.setChecked(False)

    def toggle_grid(self, checked: bool):
        self.show_grid = checked
        self.update_view()

    def set_grid_step(self):
        step, ok = QtWidgets.QInputDialog.getInt(self, "栅格步长", "像素步长:", self.grid_step, 1, 512, 1)
        if not ok:
            return
        self.grid_step = step
        self.update_view()

    def toggle_fit_to_window(self, checked: bool):
        """切换是否适应窗口大小"""
        if checked:
            self.scroll_area.setWidgetResizable(True)
            self.viewer.setScaledContents(True)
        else:
            self.scroll_area.setWidgetResizable(False)
            self.viewer.setScaledContents(False)
            if self.viewer.pixmap():
                self.viewer.resize(self.viewer.pixmap().size())

    def show_actual_size(self):
        """显示实际大小（取消适应窗口）"""
        self.act_fit_to_window.setChecked(False)
        self.scroll_area.setWidgetResizable(False)
        self.viewer.setScaledContents(False)
        if self.viewer.pixmap():
            self.viewer.resize(self.viewer.pixmap().size())


def main():
    app = QtWidgets.QApplication(sys.argv)
    w = MainWindow()
    w.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()