import os
import cv2
import re
from surround_view import CaptureThread, CameraProcessingThread
from surround_view import FisheyeCameraModel, BirdView
from surround_view import MultiBufferManager, ProjectedImageBuffer
import surround_view.param_settings as settings


def is_rtsp_url(input_str):
    """检查输入是否为RTSP URL"""
    return input_str.startswith('rtsp://')


def main():
    # 配置四个RTSP摄像头地址（前、后、左、右）
    camera_ids = [
        "rtsp://admin:@192.168.112.153:554/video",  # 前
        "rtsp://admin:@192.168.112.152:554/video",  # 后
        "rtsp://admin:@192.168.112.150:554/video",  # 左
        "rtsp://admin:@192.168.112.151:554/video"   # 右
    ]

    # 摄像头ID映射（使用索引作为设备ID）
    camera_id_mapping = {i: cam_id for i, cam_id in enumerate(camera_ids)}

    # 摄像头翻转参数（根据实际安装方向调整）
    flip_methods = [0, 0, 0, 0]  # 0:不翻转, 1:水平翻转, 2:垂直翻转, 3:水平+垂直

    # 加载相机内参模型
    yamls_dir = os.path.join(os.getcwd(), "yaml")
    camera_files = [os.path.join(yamls_dir, name + ".yaml") for name in settings.camera_names]
    camera_models = [FisheyeCameraModel(camera_file, name)
                     for camera_file, name in zip(camera_files, settings.camera_names)]

    # 初始化捕获线程（针对RTSP特殊处理）
    capture_tds = []
    for idx, (cam_id, flip) in enumerate(zip(camera_ids, flip_methods)):
        # 使用索引作为设备ID，解决KeyError问题
        if is_rtsp_url(cam_id):
            # RTSP流配置（使用FFmpeg后端）
            capture_tds.append(CaptureThread(
                device_id=idx,  # 使用索引作为设备ID
                flip_method=flip,
                use_gst=False,  # 关闭GStreamer，使用FFmpeg
                api_preference=cv2.CAP_FFMPEG,  # 强制使用FFmpeg
                resolution=(960, 640)
            ))
            # 手动打开RTSP流
            if not capture_tds[-1].cap.open(cam_id, cv2.CAP_FFMPEG):
                print(f"无法打开RTSP流: {cam_id}")
                return
            # 设置RTSP参数（兼容不同OpenCV版本）
            try:
                # 尝试通用参数设置方式
                capture_tds[-1].cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # 减少缓冲
                capture_tds[-1].cap.set(cv2.CAP_PROP_FRAME_WIDTH, 960)
                capture_tds[-1].cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 640)
                # 处理FFMPEG选项（兼容旧版本）
                if hasattr(cv2, 'CAP_PROP_FFMPEG_OPTION'):
                    capture_tds[-1].cap.set(cv2.CAP_PROP_FFMPEG_OPTION, "rtsp_transport", "tcp")
                    capture_tds[-1].cap.set(cv2.CAP_PROP_FFMPEG_OPTION, "stimeout", "5000000")
                else:
                    print(f"RTSP流 {cam_id} 使用兼容模式，部分参数可能无法设置")
            except Exception as e:
                print(f"设置RTSP参数警告: {e}")
        else:
            # 普通USB/CSI相机配置
            capture_tds.append(CaptureThread(
                device_id=idx,  # 使用索引作为设备ID
                flip_method=flip,
                use_gst=True,
                resolution=(960, 640)
            ))
            if not capture_tds[-1].connect_camera():
                print(f"启动摄像头 {cam_id} 失败，退出程序")
                return

    # 绑定捕获缓冲区管理器
    capture_buffer_manager = MultiBufferManager()
    for td in capture_tds:
        capture_buffer_manager.bind_thread(td, buffer_size=8)
        td.start()

    # 初始化处理线程（使用索引作为device_id）
    proc_buffer_manager = ProjectedImageBuffer()
    process_tds = [
        CameraProcessingThread(
            capture_buffer_manager,
            device_id=idx,  # 使用索引作为设备ID
            camera_model=model
        )
        for idx, model in enumerate(camera_models)
    ]

    # 启动处理线程
    for td in process_tds:
        proc_buffer_manager.bind_thread(td)
        td.start()

    # 初始化环视拼接
    birdview = BirdView(proc_buffer_manager)
    birdview.load_weights_and_masks("./weights.png", "./masks.png")
    birdview.start()

    # 主循环显示
    try:
        while True:
            birdview_img = birdview.get()
            if birdview_img is not None:
                display_img = cv2.resize(birdview_img, (800, 600))
                cv2.imshow("Surround View", display_img)

            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                break

            # 打印帧率信息
            status = []
            for td in capture_tds:
                cam_id = camera_id_mapping[td.device_id]
                status.append(f"捕获 {cam_id.split('@')[-1].split(':')[0]} FPS: {td.stat_data.average_fps:.1f}")
            for td in process_tds:
                cam_name = settings.camera_names[td.device_id]
                status.append(f"处理 {cam_name} FPS: {td.stat_data.average_fps:.1f}")
            status.append(f"全景 FPS: {birdview.stat_data.average_fps:.1f}")
            print(" | ".join(status), end="\r")

    finally:
        for td in process_tds:
            td.stop()
        for td in capture_tds:
            td.stop()
            td.disconnect_camera()
        birdview.stop()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()