#!/usr/bin/env python3
"""
Simple RTSP stream with YOLOv8 detection using OpenCV + subprocess for RTSP.
"""

import cv2
import subprocess
import time
from ultralytics import YOLO

BALL_CLASS_ID = 32  # sports ball in COCO
WIDTH = 1280
HEIGHT = 720
FPS = 25


def main():
    print("Loading YOLOv8...")
    model = YOLO("yolov8n.pt")
    model.to("cuda")
    print("Model loaded on CUDA")

    # Open video source
    source = "/opt/nvidia/deepstream/deepstream/samples/streams/sample_1080p_h264.mp4"
    cap = cv2.VideoCapture(source)

    if not cap.isOpened():
        print("ERROR: Cannot open video")
        return

    # Start ffmpeg to serve RTSP via tcp
    ffmpeg_cmd = [
        'ffmpeg',
        '-y',
        '-f', 'rawvideo',
        '-vcodec', 'rawvideo',
        '-pix_fmt', 'bgr24',
        '-s', f'{WIDTH}x{HEIGHT}',
        '-r', str(FPS),
        '-i', '-',
        '-c:v', 'libx264',
        '-preset', 'ultrafast',
        '-tune', 'zerolatency',
        '-f', 'rtsp',
        '-rtsp_transport', 'tcp',
        'rtsp://localhost:8554/live'
    ]

    print("Starting ffmpeg...")
    proc = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE)

    print(f"\nRTSP stream: rtsp://pickle:8554/live")
    print("Press Ctrl+C to stop\n")

    frame_count = 0
    start = time.time()

    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                continue

            frame = cv2.resize(frame, (WIDTH, HEIGHT))

            # Detection
            results = model(frame, verbose=False, classes=[BALL_CLASS_ID], conf=0.3)
            for r in results:
                for box in r.boxes:
                    x1, y1, x2, y2 = map(int, box.xyxy[0])
                    conf = float(box.conf[0])
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(frame, f"Ball {conf:.2f}", (x1, y1-10),
                               cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

            # FPS
            frame_count += 1
            fps = frame_count / (time.time() - start)
            cv2.putText(frame, f"FPS: {fps:.1f}", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)

            # Write to ffmpeg
            proc.stdin.write(frame.tobytes())

            if frame_count % 100 == 0:
                print(f"Frame {frame_count}, FPS: {fps:.1f}")

    except KeyboardInterrupt:
        print("\nStopping...")
    except BrokenPipeError:
        print("ffmpeg pipe broken")
    finally:
        cap.release()
        proc.stdin.close()
        proc.wait()


if __name__ == "__main__":
    main()