import cv2
import mss
import numpy
import argparse
import time
import threading
from queue import Queue, Empty, Full
from ultralytics import YOLO

# --- Thread 1: Screen Capture ---
def capture_thread(sct, monitor, capture_queue: Queue, stop_event: threading.Event, fps: int):
    """Continuously captures the screen and puts frames into a queue."""
    frame_delay = 1.0 / fps if fps > 0 else 0
    while not stop_event.is_set():
        start_time = time.time()
        
        img = numpy.array(sct.grab(monitor))
        try:
            # Use non-blocking put to avoid filling up the queue and causing lag
            capture_queue.put_nowait(img)
        except Full:
            pass # If the queue is full, just drop the frame

        # Frame rate limiting
        elapsed_time = time.time() - start_time
        sleep_time = frame_delay - elapsed_time
        if sleep_time > 0:
            time.sleep(sleep_time)

# --- Thread 2: YOLO Detection and Plotting ---
def detection_thread(model, capture_queue: Queue, display_queue: Queue, stop_event: threading.Event, conf: float):
    """Gets frames, runs detection and plotting, and puts results in a queue."""
    while not stop_event.is_set():
        try:
            frame = capture_queue.get(timeout=0.1)
        except Empty:
            continue

        # Pre-process frame
        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
        
        # Run detection and plotting
        results = model.predict(frame_bgr, verbose=False, conf=conf)
        annotated_frame = results[0].plot()
        
        try:
            # Use non-blocking put
            display_queue.put_nowait(annotated_frame)
        except Full:
            pass

def main():
    # --- Argument Parsing ---
    parser = argparse.ArgumentParser(description="High-performance, multi-threaded screen object detection.")
    parser.add_argument("--model", type=str, default="yolov8n.pt", help="Path to the YOLO model.")
    parser.add_argument("--region", type=int, nargs=4, metavar=('L', 'T', 'W', 'H'), default=(200, 200, 200, 300), help="Screen region: Left Top Width Height.")
    parser.add_argument("--conf", type=float, default=0.30, help="Confidence threshold.")
    parser.add_argument("--fps", type=int, default=15, help="Capture frame rate.")
    parser.add_argument("--title", type=str, default="yoloscr Live Detection", help="Window title.")
    args = parser.parse_args()

    # --- Initial Setup ---
    stop_event = threading.Event()
    capture_queue = Queue(maxsize=2)
    display_queue = Queue(maxsize=2)

    print("--- yoloscr Settings ---")
    print(f"Model: {args.model}, Confidence: {args.conf}, Target FPS: {args.fps}")
    
    # --- Load Model ---
    print("Loading YOLO model...")
    try:
        model = YOLO(args.model)
        print("Model loaded successfully.")
    except Exception as e:
        print(f"Error loading model: {e}")
        return

    # --- Main Application Logic ---
    with mss.mss() as sct:
        # Determine capture region
        monitor = sct.monitors[1] # Default to primary monitor
        if args.region:
            monitor = {"top": args.region[1], "left": args.region[0], "width": args.region[2], "height": args.region[3]}
        
        print(f"Capture Region: {monitor['left']}x{monitor['top']} at {monitor['width']}x{monitor['height']}")

        # --- Start Threads ---
        threads = [
            threading.Thread(target=capture_thread, args=(sct, monitor, capture_queue, stop_event, args.fps)),
            threading.Thread(target=detection_thread, args=(model, capture_queue, display_queue, stop_event, args.conf))
        ]
        for t in threads:
            t.daemon = True
            t.start()

        print("Threads started. Detection is live. Press 'q' in the window to quit.")

        # --- Main GUI Loop ---
        cv2.namedWindow(args.title, cv2.WINDOW_NORMAL)
        try:
            while not stop_event.is_set():
                try:
                    display_frame = display_queue.get(timeout=0.01) # Wait briefly for a new frame
                    cv2.imshow(args.title, display_frame)
                except Empty:
                    pass # No new frame, just continue and keep the GUI responsive

                if cv2.waitKey(1) & 0xFF == ord('q'):
                    print("'q' pressed, exiting.")
                    break
        except KeyboardInterrupt:
            print("\nCtrl+C detected.")
        finally:
            # --- Cleanup ---
            print("Shutting down threads...")
            stop_event.set()
            for t in threads:
                t.join(timeout=1)
            cv2.destroyAllWindows()
            print("yoloscr has finished.")

if __name__ == "__main__":
    main()


