from collections import defaultdict
import cv2
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator, colors

def detect_fire_in_video(
    input_video_path: str,
    output_video_path: str,
    model_path: str,
    device: str = "cpu"
):
    """
    Belirtilen videoda YOLO modelini kullanarak yangın tespiti yapar ve 
    etiketlenmiş çıktıyı output_video_path'e kaydeder.

    :param input_video_path: Girdi videosunun dosya yolu
    :param output_video_path: İşlenen (annotate edilmiş) videonun kaydedileceği dosya yolu
    :param model_path: YOLO model dosyasının (ör. .pt uzantılı) dosya yolu
    :param device: İşlem için kullanılacak donanım ('cpu', 'cuda', 'mps' vb.)
    :return: İşlenen video yolunu döndürür.
    """

    # Takip geçmişi için defaultdict oluştur
    track_history = defaultdict(lambda: [])

    # YOLO modelini yükle
    model = YOLO(model_path, device=device)

    # Video dosyasını aç
    cap = cv2.VideoCapture(input_video_path)

    # Video özelliklerini al
    w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH,
                                           cv2.CAP_PROP_FRAME_HEIGHT,
                                           cv2.CAP_PROP_FPS))

    # Çıktı videosu kaydı için VideoWriter ayarla
    out = cv2.VideoWriter(
        output_video_path,
        cv2.VideoWriter_fourcc(*"MJPG"),
        fps,
        (w, h)
    )

    while True:
        ret, im0 = cap.read()
        if not ret:
            print("Video bitti veya çerçeve okunamadı.")
            break

        # Annotator objesi ile çizim yapma
        annotator = Annotator(im0, line_width=2)

        # model.track ile takip (tracking) sonuçlarını al
        results = model.track(im0, persist=True)

        # Track ID'ler ve maskeler varsa işleme al
        if results[0].boxes.id is not None and results[0].masks is not None:
            masks = results[0].masks.xy
            track_ids = results[0].boxes.id.int().cpu().tolist()

            for mask, track_id in zip(masks, track_ids):
                annotator.seg_bbox(
                    mask=mask,
                    mask_color=colors(int(track_id), True),
                    label=str(track_id)
                )

        # Annotate edilmiş frame'i kaydet
        out.write(im0)

    out.release()
    cap.release()
    cv2.destroyAllWindows()

    return output_video_path