import glob
import os
import random
import cv2

# --- CONFIGURATION ---
OUTPUT_DIRECTORY = "dataset_v3"
FRAMES_TO_EXTRACT = 20  # Number of random anchor frames to find targets
SCALE_FACTOR = 0.5  # Resize video to 1/2 size for easier viewing

# Define the EXACT size of the bounding box around your click point
BOX_WIDTH = 64
BOX_HEIGHT = 64

# Defined classification options (YOLO maps these to IDs: 0, 1, 2, 3...)
CLASSIFICATION_OPTIONS = [
    "guardian",
    "interceptor",
    "sentinel",
    "rocket",
    "bomber",
    "background",
    "asteroid",
]
# ---------------------

def get_dataset_counts():
    """Counts the number of existing bounding boxes by parsing YOLO text files."""
    counts = {option: 0 for option in CLASSIFICATION_OPTIONS}
    labels_dir = os.path.join(OUTPUT_DIRECTORY, "labels")

    if os.path.exists(labels_dir):
        # Scan all text annotation files in the YOLO labels directory
        for txt_file in glob.glob(os.path.join(labels_dir, "*.txt")):
            try:
                with open(txt_file, "r") as f:
                    for line in f:
                        parts = line.strip().split()
                        if parts:
                            class_id = int(parts[0])
                            if 0 <= class_id < len(CLASSIFICATION_OPTIONS):
                                option = CLASSIFICATION_OPTIONS[class_id]
                                counts[option] += 1
            except Exception:
                pass  # Skip corrupted or unreadable text files safely
    return counts

def mouse_click_callback(event, x, y, flags, param):
    """Handles mouse clicks.

    Calculates YOLO coordinates, saves the full frame, and appends a line to
    the text file.
    """
    if event == cv2.EVENT_LBUTTONDOWN:
        # Ignore clicks if the user is in "background" capture mode
        if param["session_label"] == "background":
            print(
                " Info: In background mode, do not click targets. Just press SPACE to save frame."
            )
            return

        frame = param["frame"]
        display_frame = param["display_frame"]
        frame_idx = param["frame_idx"]
        window_name = param["window_name"]

        # Calculate top-left and bottom-right corners
        x1 = x - BOX_WIDTH // 2
        y1 = y - BOX_HEIGHT // 2
        x2 = x1 + BOX_WIDTH
        y2 = y1 + BOX_HEIGHT

        # --- Edge Case Handling ---
        if x1 < 0:
            x1 = 0
            x2 = BOX_WIDTH
        if y1 < 0:
            y1 = 0
            y2 = BOX_HEIGHT
        if x2 > frame.shape[1]:
            x2 = frame.shape[1]
            x1 = x2 - BOX_WIDTH
        if y2 > frame.shape[0]:
            y2 = frame.shape[0]
            y1 = y2 - BOX_HEIGHT

        # --- YOLO Normalized Coordinate Conversion ---
        img_h, img_w = frame.shape[0], frame.shape[1]

        true_center_x = (x1 + x2) / 2.0
        true_center_y = (y1 + y2) / 2.0

        x_center_norm = true_center_x / img_w
        y_center_norm = true_center_y / img_h
        width_norm = BOX_WIDTH / img_w
        height_norm = BOX_HEIGHT / img_h

        # Clean video name to ensure a unique, safe filename
        video_base = os.path.splitext(os.path.basename(param["video_name"]))[0]
        filename_base = f"{video_base}_frame_{frame_idx}"

        image_save_path = os.path.join(
            param["images_dir"], f"{filename_base}.jpg"
        )
        label_save_path = os.path.join(
            param["labels_dir"], f"{filename_base}.txt"
        )

        # 1. Save the full frame image (overwriting is fine if clicking multiple items)
        cv2.imwrite(image_save_path, frame)

        # 2. Append the target annotation to the YOLO text file
        class_id = param["class_id"]
        with open(label_save_path, "a") as f:
            f.write(
                f"{class_id} {x_center_norm:.6f} {y_center_norm:.6f} {width_norm:.6f} {height_norm:.6f}\n"
            )

        print(
            f" Logged {param['session_label']} box coordinate to {label_save_path}"
        )
        param["box_drawn"] = True

        # Draw visual feedback box and short class label on screen
        cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(
            display_frame,
            param["session_label"],
            (x1, y1 - 5),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.4,
            (0, 255, 0),
            1,
        )
        cv2.imshow(window_name, display_frame)

def label_video_session():
    print("=========================================")
    print("        AI YOLO DATASET LABELER          ")
    print("=========================================\n")

    # 1. Automatically find and select a random MP4 video
    mp4_files = glob.glob("*.mp4")
    if not mp4_files:
        print(
            "Error: No .mp4 files found in the current directory.\n"
            "Please place this script in the same folder as your video files."
        )
        return

    selected_video = random.choice(mp4_files)
    print(f"-> Selected Video File: '{selected_video}'")

    # 2. Fetch and display current dataset balance dashboard
    counts = get_dataset_counts()
    print("\n-----------------------------------------")
    print("      CURRENT DATASET BALANCE STATUS     ")
    print("-----------------------------------------")
    for option, count in counts.items():
        print(f"  • {option.capitalize():<12} : {count} boxes logged")
    print("-----------------------------------------")
    print(
        "💡 Tip: Try to choose options with lower counts to keep data balanced!"
    )

    # 3. Force selection of a valid label
    print("\nAvailable Classification Targets:")
    for idx, option in enumerate(CLASSIFICATION_OPTIONS, start=1):
        print(f"  [{idx}] {option} (Current Boxes: {counts[option]})")

    while True:
        try:
            choice = input(
                f"\nEnter the number (1-{len(CLASSIFICATION_OPTIONS)}) of what you are labeling: "
            ).strip()
            choice_idx = int(choice) - 1
            if 0 <= choice_idx < len(CLASSIFICATION_OPTIONS):
                session_label = CLASSIFICATION_OPTIONS[choice_idx]
                session_class_id = choice_idx
                break
            else:
                print("Invalid selection. Please choose a number from the list.")
        except ValueError:
            print("Invalid input. Please enter a number.")

    # Open video and check validity
    cap = cv2.VideoCapture(selected_video)
    if not cap.isOpened():
        print(f"Error: Could not open video file {selected_video}")
        return

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total_frames == 0:
        print("Error: Video has 0 frames.")
        return

    # Select random unique anchor frames to jump between
    num_frames = min(FRAMES_TO_EXTRACT, total_frames)
    frame_indices = random.sample(range(total_frames), num_frames)
    frame_indices.sort()

    # Set up standard YOLO dataset directory branches
    images_dir = os.path.join(OUTPUT_DIRECTORY, "images")
    labels_dir = os.path.join(OUTPUT_DIRECTORY, "labels")
    os.makedirs(images_dir, exist_ok=True)
    os.makedirs(labels_dir, exist_ok=True)

    # Comprehensive User Instructions
    print("\n" + "=" * 50)
    print("               HOW TO LABEL")
    print("=" * 50)
    print(f"• TARGET OBJECT: {session_label.upper()}")
    print(f"• SESSION SCOPE: {num_frames} random search zones in '{selected_video}'")
    print(f"• BOUNDING BOX : {BOX_WIDTH}x{BOX_HEIGHT} pixels scaled to YOLO layout")
    print("-" * 50)
    print("► CONTROLS:")
    print("  [Left-Click]     - Click directly on target to tag position (auto-saves frame).")
    print("  [Right Arrow]    - Move FORWARD 1 individual frame (sequential tracking, no save).")
    print("  [Left Arrow]     - Move BACKWARD 1 individual frame (sequential tracking, no save).")
    print("  [SPACE / ENTER]  - Jump completely to the NEXT random video search spot.")
    print("  [Q Key]          - Quit and save all progress up to this point.")
    print("=" * 50)
    input("\nPress ENTER when you are ready to begin...")

    # Cross-platform arrow key mappings for waitKeyEx()
    LEFT_KEYS = [81, 2, 2424832, 65361, 63234]
    RIGHT_KEYS = [83, 3, 2490368, 65363, 63235]

    i = 0
    current_frame = frame_indices[i]

    while i < len(frame_indices):
        cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
        ret, frame = cap.read()
        if not ret:
            print(f" Warning: Could not read frame {current_frame}. Moving to next random anchor.")
            i += 1
            if i < len(frame_indices):
                current_frame = frame_indices[i]
            continue

        # Resize the frame by 1/2 for laptop viewability
        frame_resized = cv2.resize(
            frame, (0, 0), fx=SCALE_FACTOR, fy=SCALE_FACTOR
        )
        frame_display = frame_resized.copy()

        window_name = f"Labeling: {session_label} | Spot {i+1}/{num_frames} (Frame {current_frame})"
        cv2.namedWindow(window_name)

        session_state = {
            "frame": frame_resized,
            "display_frame": frame_display,
            "images_dir": images_dir,
            "labels_dir": labels_dir,
            "frame_idx": current_frame,
            "window_name": window_name,
            "video_name": selected_video,
            "class_id": session_class_id,
            "session_label": session_label,
            "box_drawn": False,
        }

        cv2.setMouseCallback(window_name, mouse_click_callback, session_state)

        # Keep window open until user acts
        cv2.imshow(window_name, frame_display)
        while True:
            cv2.moveWindow(window_name, 0, 0)
            key = cv2.waitKeyEx(1)
            
            # Enter (13) or Space (32) -> Jump to the next random anchor spot
            if key == 13 or key == 32:
                if session_label == "background":
                    video_base = os.path.splitext(
                        os.path.basename(selected_video)
                    )[0]
                    filename_base = f"{video_base}_frame_{current_frame}"
                    image_save_path = os.path.join(
                        images_dir, f"{filename_base}.jpg"
                    )
                    label_save_path = os.path.join(
                        labels_dir, f"{filename_base}.txt"
                    )

                    cv2.imwrite(image_save_path, frame_resized)
                    open(label_save_path, "a").close()  # Creates empty file
                    print(f" Registered negative background frame: {filename_base}")
                
                i += 1
                if i < len(frame_indices):
                    current_frame = frame_indices[i]
                break

            # Right Arrow -> Move forward exactly ONE individual sequential frame
            elif key in RIGHT_KEYS:
                if current_frame < total_frames - 10:
                    current_frame += 10
                else:
                    print(" Already at the very last frame of the video!")
                    continue
                break

            # Left Arrow -> Move backward exactly ONE individual sequential frame
            elif key in LEFT_KEYS:
                if current_frame > 0:
                    current_frame -= 10
                else:
                    print(" Already at the first frame of the video!")
                    continue
                break

            # Quit 
            elif key == ord("q") or key == ord("Q") or (key & 0xFF) in [ord("q"), ord("Q")]:
                print("\nQuitting session early... Saving progress.")
                cap.release()
                cv2.destroyAllWindows()
                return

        cv2.destroyWindow(window_name)

    cap.release()
    cv2.destroyAllWindows()
    print(f"\nSession finished! Dataset updated successfully inside '{OUTPUT_DIRECTORY}'.")

if __name__ == "__main__":
    label_video_session()