import cv2
from websocket import create_connection
import win32file
import struct
import numpy as np
from ultralytics import YOLO

# object classes
classNames = ["person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
              "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
              "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
              "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
              "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
              "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
              "carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", "bed",
              "diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
              "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
              "teddy bear", "hair drier", "toothbrush"
              ]


def handle_connection(pipe):
    # Load YOLO model
    model = YOLO('yolov8s.pt')  # Load an official Detect model

    print("Pipe connected!")
    try:
        while True:
            data = win32file.ReadFile(pipe, 640*480*3)[1]

            if not data:
                return

            nparr = np.frombuffer(data, np.uint8).reshape((480, 640, 3))

            # frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            frame = cv2.cvtColor(nparr, cv2.COLOR_BGR2RGB)

            # Perform object detection
            results = model.track(frame, persist=True)

            for r in results:
                lines = ""
                boxes = r.boxes
                for box in boxes:
                    if box.cls[0].item() == 0 and not box.id is None:
                        # bounding box
                        id = box.id.int().cpu().tolist()
                        # x1, y1, x2, y2 = box.xyxy[0]
                        # #  2.08333 = 1/480 * 1000 or normalize, then save 4 sig-figures and cast to int
                        # #  1.5625 = 1/640 * 1000 or normalize, then save 4 sig-figures and cast to int

                        # x1, x2, y1, y2 = int(x1), int(x2), int(y1), int(y2)

                        # # put box in cam
                        # cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 255), 3)

                        # # class name
                        # cls = int(box.cls[0])

                        # # object details
                        # org = [x1, y1]
                        # org2 = [x1, y1+50]
                        # font = cv2.FONT_HERSHEY_SIMPLEX
                        # fontScale = 1
                        # color = (255, 0, 0)
                        # color_w = (255, 255, 255)
                        # thickness = 2

                        # cv2.putText(frame, classNames[cls], org, font, fontScale, color, thickness)
                        # cv2.putText(frame, str(id), org2, font, fontScale, color_w, thickness)

                        x1, y1, x2, y2 = box.xyxyn[0]
                        x1, y1, x2, y2 = int(x1 * 1000), int(y1 * 1000), int(x2 * 1000), int(y2 * 1000)
                        lines += f"{id} {x1}:{y1} {x2}:{y2}\n"

                # cv2.imshow('Webcam', frame)
                # if cv2.waitKey(1) == ord('q'):
                #     break

            ret = lines.encode('utf-8')
            length_prefix = struct.pack('I', len(ret))

            win32file.WriteFile(pipe, length_prefix+ret)
    except Exception as e:
        print("Pipe closed with error: " + str(e))
        return
    # finally:
    #     cv2.destroyAllWindows()