import os
import sys
import cv2
import time  # Import time module for adding delay
from icecream import ic

env_path = os.path.join(os.path.dirname(__file__), '')
if env_path not in sys.path:
    sys.path.append(env_path)

from pytracking.evaluation import Tracker


def run_webcam(debug=None, refresh_rate=30):
    tracker = Tracker('dimp', 'dimp50', debug)

    class UIControl:
        def __init__(self):
            self.mode = 'init'  # init, select, track
            self.target_tl = (-1, -1)
            self.target_br = (-1, -1)
            self.new_init = False

        def mouse_callback(self, event, x, y, flags, param):
            if event == cv2.EVENT_LBUTTONDOWN and self.mode == 'init':
                self.target_tl = (x, y)
                self.target_br = (x, y)
                self.mode = 'select'
            elif event == cv2.EVENT_MOUSEMOVE and self.mode == 'select':
                self.target_br = (x, y)
            elif event == cv2.EVENT_LBUTTONDOWN and self.mode == 'select':
                self.target_br = (x, y)
                self.mode = 'init'
                self.new_init = True

        def get_tl(self):
            return self.target_tl if self.target_tl[0] < self.target_br[0] else self.target_br

        def get_br(self):
            return self.target_br if self.target_tl[0] < self.target_br[0] else self.target_tl

        def get_bb(self):
            tl = self.get_tl()
            br = self.get_br()
            bb = [min(tl[0], br[0]), min(tl[1], br[1]), abs(br[0] - tl[0]), abs(br[1] - tl[1])]
            return bb

    ui_control = UIControl()
    cap = cv2.VideoCapture(0)
    #cap = cv2.VideoCapture('/home/mht/Downloads/car4.mp4')
    #cap = cv2.VideoCapture('/home/mht/Downloads/messi.mp4')
    display_name = 'tracker'
    cv2.namedWindow(display_name, cv2.WINDOW_NORMAL)
    cv2.setMouseCallback(display_name, ui_control.mouse_callback)

    # Read the first frame before the loop starts
    ret, frame = cap.read()
    if not ret:
        print("Error: Cannot read video.")
        cap.release()
        sys.exit()

    frame_disp = frame.copy()


    # Show the first frame and let the user draw a bounding box
    while not ui_control.new_init:

        temp_disp = frame_disp.copy()
        if ui_control.mode == 'select':
            cv2.rectangle(temp_disp, ui_control.get_tl(), ui_control.get_br(), (0, 255, 0), 2)
        cv2.putText(temp_disp, "Draw bounding box and click to confirm...",
                    (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
        cv2.imshow(display_name, temp_disp)
        key = cv2.waitKey(1)
        if key == ord('q'):  # Quit
            cap.release()
            cv2.destroyAllWindows()
            sys.exit()

    # Initialize the tracker with the bounding box selected
    init_bbox = ui_control.get_bb()
    ic(init_bbox)
    tracker.init(frame, init_bbox)

    #tracker.init(frame, [243, 159, 17, 14]) #fixed starting bb

    # Frame delay for refresh rate
    frame_delay = 1.0 / refresh_rate

    # Main loop starts
    while True:

        start_time = time.time()
        ret, frame = cap.read()
        if not ret:
            break  # Exit if no more frames
        frame_disp = frame.copy()

        bbox, success = tracker.update(frame)

        if bbox is not None:
            color = (0, 255, 0) if success else (0, 0, 255)
            cv2.rectangle(frame_disp, (bbox[0], bbox[1]), (bbox[0] + bbox[2], bbox[1] + bbox[3]),
                          color, 3)

        cv2.imshow(display_name, frame_disp)

        key = cv2.waitKey(1)
        if key == ord('q'):  # Quit
            break
        elif key == ord('r'):  # Reset tracker
            tracker.stop()
            ui_control.mode = 'init'
            ui_control.new_init = False

            # Reinitialize on current frame
            while not ui_control.new_init:
                temp_disp = frame.copy()
                if ui_control.mode == 'select':
                    cv2.rectangle(temp_disp, ui_control.get_tl(), ui_control.get_br(), (0, 255, 0), 2)
                cv2.putText(temp_disp, "Draw bounding box and click to confirm...",
                            (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
                cv2.imshow(display_name, temp_disp)
                key = cv2.waitKey(1)
                if key == ord('q'):  # Quit
                    cap.release()
                    cv2.destroyAllWindows()
                    sys.exit()
            init_bbox = ui_control.get_bb()
            tracker.init(frame, init_bbox)

        # Maintain refresh rate
        elapsed_time = time.time() - start_time
        time.sleep(max(0, frame_delay - elapsed_time))

    cap.release()
    cv2.destroyAllWindows()


def main():
    run_webcam(debug=False, refresh_rate=24)


if __name__ == '__main__':
    main()