import socket
import cv2
import numpy as np

# 配置
TCP_IP = '0.0.0.0'
TCP_PORT = 5005

def gstreamer_pipeline(
    sensor_id=0,
    capture_width=1920,
    capture_height=1080,
    display_width=960,
    display_height=540,
    framerate=30,
    flip_method=0,
):
    return (
        "nvarguscamerasrc sensor-id=%d ! "
        "video/x-raw(memory:NVMM), width=(int)%d, height=(int)%d, framerate=(fraction)%d/1 ! "
        "nvvidconv flip-method=%d ! "
        "video/x-raw, width=(int)%d, height=(int)%d, format=(string)BGRx ! "
        "videoconvert ! "
        "video/x-raw, format=(string)BGR ! appsink"
        % (
            sensor_id,
            capture_width,
            capture_height,
            framerate,
            flip_method,
            display_width,
            display_height,
        )
    )

# 初始化摄像头
cap = cv2.VideoCapture(gstreamer_pipeline(flip_method=2), cv2.CAP_GSTREAMER)
if not cap.isOpened():
    print("无法打开摄像头，尝试使用默认摄像头设备")
    cap = cv2.VideoCapture(0)

# 创建服务器套接字
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind((TCP_IP, TCP_PORT))
sock.listen(1)
print(f"服务器已启动，监听 {TCP_IP}:{TCP_PORT}")

try:
    while True:
        print("等待客户端连接...")
        conn, addr = sock.accept()
        print(f"客户端已连接：{addr}")
        conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    print("无法读取摄像头画面")
                    break

                # 编码为JPEG
                encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 70]
                result, img_encoded = cv2.imencode('.jpg', frame, encode_param)
                if not result:
                    continue

                data = img_encoded.tobytes()

                # 发送图像大小（4字节）
                conn.sendall(len(data).to_bytes(4, 'big'))

                # 发送图像数据
                conn.sendall(data)

                # 显示本地画面（可选）
                cv2.imshow('Sending Image', frame)
                if cv2.waitKey(1) == 27:  # 按 ESC 退出
                    raise KeyboardInterrupt

        except (BrokenPipeError, ConnectionResetError) as e:
            print(f"[错误] 客户端断开连接: {e}")
        finally:
            conn.close()
            print("客户端连接已关闭，等待新连接...")

except KeyboardInterrupt:
    print("服务器正在关闭...")

finally:
    cap.release()
    cv2.destroyAllWindows()
    sock.close()
    print("所有资源已释放")