import sophon_chakcy.sail as sail
from sophon_chakcy.get_detector import get_detector
from typing import List
import time
import threading

def task(handle:sail.Handle, encoder:sail.Encoder, detector):
    decoder = sail.Decoder("resources/datasets/test_car_person_1080P.mp4", False, 0)
    count = 0
    while(True):
        # 初始化bm_image
        img = sail.BMImage()
        # decode
        decode_ret = decoder.read(handle, img)
        count += 1
        if count % 20 == 0:
            continue
        if(decode_ret!=0):
            decoder.reconnect()
            continue
        
        ## you can use bmcv to do something here
        
        # time.sleep(0.01)
        filtered_detections = detector.detect(img)
        img = detector.draw_detections(img, filtered_detections)

        filtered_detections = detector.detect(img)
        img = detector.draw_detections(img, filtered_detections)

        filtered_detections = detector.detect(img)
        img = detector.draw_detections(img, filtered_detections)

        ## encode   (push rtsp/rtmp stream or write local video file)
        encode_ret = encoder.video_write(img)
        while(encode_ret != 0):
            # time.sleep(0.01)
            encode_ret = encoder.video_write(img)

# 视频推流
def video_push_stream(handle:sail.Handle, encoders:List[sail.Encoder], detector):
    threads = []
    for encoder in encoders:
        t = threading.Thread(target=task, args=(handle, encoder, detector))
        t.start()
        threads.append(t)
    for t in threads:
        t.join()

if __name__ == "__main__":
    handle = sail.Handle(0)

    # 初始化bmcv的控制指令
    bmcv = sail.Bmcv(handle)

    # 初始化解码器
    detector_class, detector_config = get_detector("yolo_example")
    detector = detector_class(handle, detector_config)
   
    enc_params = f"width=1920:height=1080:bitrate=4000:gop=18:gop_preset=2:framerate=1"
    encoders = [
        sail.Encoder(f"rtsp://192.168.137.1:8554/stream{idx}", handle, 'h264_bm', 'I420', enc_params, 10, 1)
        for idx in range(16)
    ]
    video_push_stream(handle, encoders, detector)