from sophon_chakcy.utils.multi_decoder_manager import MultiDecoderManager, TaskArgs
import sophon_chakcy.sail as sail


def main():
    hand = sail.Handle(0)
    input_output_path = [
        # (f"resources/datasets/test_car_person_1080P_{i}.mp4", 
        (f"rtsp://192.168.137.1:8554/stream", 
        #  f"./results/output{i}.mp4") 
         f"rtsp://192.168.137.1:8554/stream{i}") 
         for i in range(8)
    ]
    with MultiDecoderManager() as decoder_manager:
        decoder_manager.set_local_flag(True)
        try:
            for input_path, output_path in input_output_path:
                decoder_manager.add_channel(
                    input_path, 
                    frame_skip_num=20, 
                    loopnum=0, 
                    output=output_path,
                    task_args=TaskArgs(
                        task_name=f"task_{input_path.split('/')[-1]}",
                        algorithms=["yolo_example"],
                        handle=hand
                    )
                )
            while True:
                for bmimg, channel_idx in decoder_manager.read_frame_all_channels():
                    encoder = decoder_manager.get_encoder(channel_idx)
                    algorithms_list = decoder_manager.get_algorithms(channel_idx)
                    if encoder is None or bmimg is None:
                        continue
                    # if decoder_manager.is_channel_eof(channel_idx):
                    #     decoder_manager.remove_channel(channel_idx)
                    
                    if algorithms_list is not None:
                        for algorithm in algorithms_list:
                            filtered_detections = algorithm.detect(bmimg)
                            bmimg = algorithm.draw_detections(bmimg, filtered_detections, None)
                    
                    encoder.video_write(bmimg.data())
                if len(decoder_manager.get_active_channels()) == 0:
                    break
                
        except:
            pass
        finally:
            decoder_manager.cleanup()

if __name__ == "__main__":
    main()