from ultralytics import YOLO
import subprocess
import threading
import os
import time
import cv2
from PIL import Image
import threading

class ObjectTracker:
    def __init__(self, video_source=None, model_path='best.onnx'):
        # Initialize the YOLO model and video capture
        self.model = YOLO(model_path, task="detect")
        if video_source is not None:
            self.cap = cv2.VideoCapture(video_source)

        # Object classes and tracking information
        self.class_names = ["head", "helmet", "motorbike"]
        self.passing_ids = set()
        self.previous_centers = {}
        self.threads = []
        
        self.start = True

    # 
    def process_object(self, buffer):
        """Run subprocess to process object asynchronously."""
        process = subprocess.Popen(['python', 'recognition.py'], stdin=subprocess.PIPE)
        process.communicate(input=buffer.tobytes())
        print(2)

    def detect_and_track(self):
        """Main loop for detecting and tracking objects."""
        while True:
            ret, frame = self.cap.read()
            if not ret:
                break  # Stop if no frame is returned

            frame = cv2.resize(frame, (600, 600))
            motorbike_detect = list(self.model.track(frame, persist=True, conf=0.65, classes=[2]))
            annotated_frame = motorbike_detect[0].plot()

            self.handle_detections(motorbike_detect, frame)
            cv2.imshow('Webcam', annotated_frame)

            if cv2.waitKey(1) == ord('q'):
                break

        self.cap.release()
        cv2.destroyAllWindows()


    def detect_and_track_2(self, video):
        self.cap = cv2.VideoCapture(video)
        self.start = True
        """Main loop for detecting and tracking objects."""
        while self.start:
            ret, frame = self.cap.read()
            if not ret:
                break  # Stop if no frame is returned

            frame = cv2.resize(frame, (600, 600))
            motorbike_detect = list(self.model.track(frame, persist=True, conf=0.65, classes=[2]))
            annotated_frame = motorbike_detect[0].plot()

            self.handle_detections(motorbike_detect, frame)
            yield annotated_frame

            if cv2.waitKey(1) == ord('q'):
                break

        self.cap.release()
        cv2.destroyAllWindows()


    def handle_detections(self, detections, frame):
        """Process detected objects and track their movement."""
        for result in detections[0].boxes.data.tolist():
            if len(result) == 7:  # Ensure detection has all necessary attributes
                x1, y1, x2, y2, obj_id, conf, label = result

                if obj_id not in self.previous_centers:
                    self.previous_centers[obj_id] = 0

                center_x = int((x1 + x2) / 2)
                print(f"id: {obj_id} {self.previous_centers[obj_id]} {center_x}")

                if self.crossed_line(obj_id, center_x):
                    self.save_and_process_object(frame, x1, y1, x2, y2, obj_id)

                self.previous_centers[obj_id] = center_x

    def crossed_line(self, obj_id, center_x):
        """Check if the object has crossed the predefined line."""
        return (
            self.previous_centers[obj_id] >= 150 and
            center_x <= 150 and
            obj_id not in self.passing_ids
        )

    def save_and_process_object(self, frame, x1, y1, x2, y2, obj_id):
        """Save the object image and process it in a new thread."""
        self.passing_ids.add(obj_id)
        print(f"id:{obj_id} crossed the line")

        img_crop = frame[int(y1 - 150):int(y2), int(x1):int(x2 + 20), :]
        # cv2.imwrite(f'img/object_crossed{obj_id}.jpg', img_crop)
        _, buffer = cv2.imencode('.jpg',img_crop)
        with open(f'img/object_crossed{obj_id}.jpg', "wb") as f:
            f.write(buffer)

        print(1)
        thread = threading.Thread(target=self.process_object, args=(buffer,))
        thread.start()
        self.threads.append(thread)
        print(3)

def test_camera():
    cap = cv2.VideoCapture('http://100.86.155.229:8080/video')

    if not cap.isOpened():
        print("Error: Could not open the webcam.")
    else:
        print("Press 's' to save a frame, or 'q' to quit.")

    # Loop to continuously capture frames
    while True:
        ret, frame = cap.read()  # Capture a frame
        if not ret:
            print("Failed to grab frame.")
            break

        # Display the frame in a window
        cv2.imshow("Webcam Stream", frame)

        # Check for key press events
        key = cv2.waitKey(1) & 0xFF  # Mask for cross-platform compatibility

        if key == ord('s'):  # Press 's' to save a frame
            cv2.imwrite("saved_frame.jpg", frame)  # Save the frame as a JPEG
            print("Frame saved as 'saved_frame.jpg'.")

        elif key == ord('q'):  # Press 'q' to quit
            print("Exiting...")
            break

    # Release the webcam and close the window
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    tracker = ObjectTracker()
    tracker.detect_and_track()
    #test_camera()