import cv2
import numpy as np
import time
import threading
from collections import deque
import torch

from config import *
from helpers import clip_box, crop_roi, preprocess_for_yolo, filter_yolo_boxes_with_scores

# Async YOLO worker
# =========================

class YOLOWorker:
    def __init__(self, model):
        self.model = model
        self.req = deque(maxlen=YOLO_QUEUE_MAX)
        self.res = deque(maxlen=1)
        self.lock = threading.Lock()
        self.running = False
        self.thread = threading.Thread(target=self._loop, daemon=True)

    def start(self):
        self.running = True
        self.thread.start()

    def stop(self):
        self.running = False
        self.thread.join(timeout=1.0)

    def submit(self, frame_eff_bgr, roi_box_eff, mode, ts):
        with self.lock:
            self.req.append((frame_eff_bgr, roi_box_eff, mode, ts))

    def try_get(self):
        with self.lock:
            if not self.res:
                return None
            return self.res.pop()

    def _loop(self):
        while self.running:
            item = None
            with self.lock:
                if self.req:
                    item = self.req.pop()
                    self.req.clear()

            if item is None:
                time.sleep(0.001)
                continue

            frame, roi_box, mode, ts = item
            h, w = frame.shape[:2]
            frame_infer = preprocess_for_yolo(frame)
            dets = []
            infer_ms = 0.0
            used_roi = False

            try:
                if roi_box is not None:
                    roi_box = clip_box(roi_box, w, h)
                    crop, ox, oy = crop_roi(frame_infer, roi_box)
                    if crop.size > 0:
                        used_roi = True
                        t0 = time.perf_counter()
                        with torch.inference_mode():
                            r = self.model(
                                crop,
                                conf=YOLO_CONF_EFFECTIVE,
                                imgsz=IMG_SIZE_ROI,
                                verbose=False,
                                max_det=MAX_DET,
                                device=DEVICE,
                                half=USE_HALF
                            )[0]
                        infer_ms = (time.perf_counter() - t0) * 1000.0
                        dets = filter_yolo_boxes_with_scores(
                            r,
                            frame_w=w,
                            frame_h=h,
                            offset_x=ox,
                            offset_y=oy,
                            min_conf=BT_LOW
                        )

                else:
                    sh, sw = frame_infer.shape[:2]
                    short = min(sh, sw)
                    scale = 1.0
                    target = IMG_SIZE_FULL
                    if short > target:
                        scale = target / float(short)
                        small = cv2.resize(
                            frame_infer,
                            (int(sw * scale), int(sh * scale)),
                            interpolation=cv2.INTER_AREA
                        )
                    else:
                        small = frame_infer

                    t0 = time.perf_counter()
                    with torch.inference_mode():
                        r = self.model(
                            small,
                            conf=YOLO_CONF_EFFECTIVE,
                            imgsz=IMG_SIZE_FULL,
                            verbose=False,
                            max_det=MAX_DET,
                            device=DEVICE,
                            half=USE_HALF
                        )[0]
                    infer_ms = (time.perf_counter() - t0) * 1000.0

                    dets_s = filter_yolo_boxes_with_scores(
                        r,
                        frame_w=small.shape[1],
                        frame_h=small.shape[0],
                        offset_x=0,
                        offset_y=0,
                        min_conf=BT_LOW
                    )
                    if scale != 1.0:
                        inv = 1.0 / scale
                        dets = [
                            np.array([d[0] * inv, d[1] * inv, d[2] * inv, d[3] * inv, d[4]], dtype=np.float32)
                            for d in dets_s
                        ]
                    else:
                        dets = dets_s

            except Exception:
                dets = []
                infer_ms = 0.0

            with self.lock:
                self.res.append((dets, ts, mode, infer_ms, used_roi))


# =========================