import os
import cv2
import datetime
import threading
from ultralytics import YOLO

class YoloPredict(threading.Thread):
    """ YOLO 推論執行緒 """

    def __init__(self, image, model_path, save_dir="results", callback=None):
        super(YoloPredict, self).__init__()
        self.image = image
        self.model_path = model_path
        self.save_dir = save_dir
        self.callback = callback  # ✅ 設定回呼函數
        self.model = None

        # ✅ 確保 `results/` 資料夾存在
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)

    def run(self):
        """ 執行 YOLO 推論 """
        if self.image is None:
            print("⚠️ 無影像可進行推論")
            return

        try:
            # ✅ 載入 YOLO 模型
            if self.model is None:
                self.model = YOLO(self.model_path)
                print("✅ YOLO 模型載入成功")

            # ✅ 轉換影像格式 (BGR → RGB)
            image_rgb = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)

            # ✅ 使用 YOLO 模型進行推論
            results = self.model.predict(image_rgb, imgsz=640, conf=0.5)
            print("✅ YOLO 推論完成")

            # ✅ 取得標註結果
            result_image = results[0].plot()

            # ✅ 儲存影像
            timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
            save_path = os.path.join(self.save_dir, f"{timestamp}.jpg")
            result_image_bgr = cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR)
            cv2.imwrite(save_path, result_image_bgr)
            print(f"✅ 推論結果儲存至: {save_path}")

            # ✅ 回傳結果到主視窗 (`callback` 函數)
            if self.callback:
                self.callback(result_image, save_path)

        except Exception as e:
            print(f"⚠️ 推論時發生錯誤: {e}")