#!/usr/bin/env python3
"""
Simple example: Decode 4 RTSP streams and display with OpenCV using callbacks

Usage:
    python stream_decoder_test.py
"""

import os

import cv2
from dotenv import load_dotenv

from services import StreamDecoderFactory

load_dotenv()

# Frame storage for each camera
frames = {1: None, 2: None, 3: None, 4: None}


def make_callback(cam_id):
    """Create callback for specific camera"""

    def callback(frame_ref):
        # Transfer to CPU and convert RGB to BGR
        frame = frame_ref.rgb_tensor.cpu().permute(1, 2, 0).numpy()
        frames[cam_id] = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        frame_ref.free()

    return callback


# Initialize factory and decoders
factory = StreamDecoderFactory(gpu_id=0)
decoders = []

for i in range(1, 5):
    url = os.getenv(f"CAMERA_URL_{i}")
    decoder = factory.create_decoder(url, buffer_size=5)
    decoder.register_frame_callback(make_callback(i))
    decoder.start()
    decoders.append(decoder)
    print(f"Camera {i} started")

# Display loop
print("Press 'q' to quit")
while True:
    # Show each camera in separate window
    for cam_id, frame in frames.items():
        if frame is not None:
            cv2.imshow(f"Camera {cam_id}", frame)

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

# Cleanup
for decoder in decoders:
    decoder.stop()
cv2.destroyAllWindows()