import sys
import json
import threading
import queue
import subprocess
import tkinter as tk
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont, ImageTk
from pypdf import PdfReader

# --- Configuration & Globals ---
padding = 60 # White margin to the right


image_queue = queue.Queue(maxsize=5) # Buffer a few images ahead
try:
    font = ImageFont.truetype("DejaVuSans.ttf", size=30)
except OSError:
    font = ImageFont.load_default()

# --- Processing Logic (Worker Thread) ---

def page_number(b, nb_pages):
    column_width = 1000 // nb_pages
    center_x = (b[1] + b[3]) // 2
    return center_x // column_width

def prepare_image(image_path: str, bounding_boxes, all_labels, nb_pages):
    """
    Draws boxes on the image and returns the PIL Image object.
    Does NOT display it.
    """
    im = Image.open(image_path)
    # Ensure image is loaded so we can pass it between threads safely
    im.load()

    width, height = im.size

    # Add white padding to the right
    new_im = Image.new(im.mode, (width + padding, height), "white")
    new_im.paste(im, (0, 0))

    draw = ImageDraw.Draw(new_im)

    bounding_boxes.sort(key=lambda b: (page_number(b["box_2d"], nb_pages), b["box_2d"][0]))

    last_label_index = -1

    for bbox in bounding_boxes:
        raw_y_min = int(bbox["box_2d"][0] * height / 1000)
        raw_x_min = int(bbox["box_2d"][1] * width  / 1000)
        raw_y_max = int(bbox["box_2d"][2] * height / 1000)
        raw_x_max = int(bbox["box_2d"][3] * width  / 1000)

        abs_y_min = max(0, raw_y_min - 10)
        abs_x_min = max(0, raw_x_min - 10)
        abs_y_max = min(height, raw_y_max + 10)
        abs_x_max = min(width, raw_x_max + 10)

        color = "black"
        label = bbox.get("label")

        if label and label in all_labels:
            current_index = all_labels.index(label)
            if current_index < last_label_index:
                color = "red"
            last_label_index = current_index

        draw.rectangle(
            ((abs_x_min, abs_y_min), (abs_x_max, abs_y_max)),
            outline=color,
            width=4,
        )
        if label:
            # draw.text((abs_x_min + 8, abs_y_min + 6), label, fill=color, font=font)
            if abs_y_min > 80:
                draw.text((abs_x_min + 8, abs_y_min - 30), label, fill=color, font=font)
            else:
                draw.text((abs_x_min + 8, abs_y_max + 6), label, fill=color, font=font)

    return new_im

def worker_thread(base_dir, files_to_process, all_labels):
    """
    Iterates through files, processes them, and puts them in the queue.
    """
    for img_path in files_to_process:
        json_path = base_dir / f"{img_path.stem}.json"
        pdf_path = base_dir / f"{img_path.stem}.pdf"

        nb_pages = 1
        if pdf_path.exists():
            try:
                nb_pages = len(PdfReader(pdf_path).pages)
            except Exception:
                pass

        if json_path.exists():
            try:
                with open(json_path, 'r') as f:
                    json_result = json.load(f)

                bb_list = json_result.get("list", [])
                print(f"Processing {img_path.name}...")

                # Draw boxes
                pil_image = prepare_image(str(img_path), bb_list, all_labels, nb_pages)

                # Block if queue is full (waiting for user to view)
                image_queue.put((pil_image, json_path))
            except Exception as e:
                print(f"Error processing {img_path.name}: {e}")

    # Sentinel to indicate finished
    image_queue.put((None, None))

# --- GUI Logic (Main Thread) ---

class ImageViewer:
    def __init__(self, root, base_dir):
        self.root = root
        self.base_dir = base_dir
        self.root.title("Bounding Box Viewer")

        # UI Elements
        self.label = tk.Label(root, text="Waiting for images...")
        self.label.pack(expand=True, fill="both")

        # State
        self.current_image = None
        self.current_json_path = None
        self.is_viewing = False
        self.scale_factor = 1.0        # To track resizing
        self.orig_size = (1, 1)        # To track original dimensions

        # Input Bindings
        self.root.bind('<Return>', self.on_enter)
        self.root.bind('e', self.on_edit)
        self.root.bind('o', self.on_open_pdf) # <--- 3. Add Key Binding
        self.root.bind('<Escape>', lambda e: self.root.quit())
        self.label.bind('<Button-1>', self.on_click) # Bind left mouse click

        # Start polling queue
        self.poll_queue()

    def poll_queue(self):
        if not self.is_viewing:
            try:
                pil_image, json_path = image_queue.get_nowait()

                if pil_image is None:
                    print("All images processed.")
                    self.root.quit() # Stop the program
                    return

                self.display_image(pil_image, json_path)
            except queue.Empty:
                pass
        self.root.after(100, self.poll_queue)

    def on_open_pdf(self, event):
        if self.is_viewing and self.current_json_path:
            # Replace .json extension with .pdf
            pdf_path = self.current_json_path.with_suffix(".pdf")

            print(f"Opening {pdf_path}")
            # Use subprocess to run xdg-open without blocking
            subprocess.Popen(['xdg-open', str(pdf_path)])

    def display_image(self, pil_image, json_path):
        self.orig_size = pil_image.size
        self.scale_factor = 1.0

        # Resize if too large for screen
        screen_h = self.root.winfo_screenheight() - 100
        if pil_image.height > screen_h:
            self.scale_factor = screen_h / pil_image.height
            pil_image = pil_image.resize((int(pil_image.width * self.scale_factor),
                                          int(pil_image.height * self.scale_factor)))

        self.tk_image = ImageTk.PhotoImage(pil_image)
        self.label.config(image=self.tk_image, text="")
        self.current_json_path = json_path
        self.is_viewing = True
        self.root.lift()

    def on_enter(self, event):
        if self.is_viewing:
            print("Next...")
            self.is_viewing = False
            self.label.config(image="", text="Loading next...")

    def on_edit(self, event):
        if self.is_viewing and self.current_json_path:
            print(f"Opening {self.current_json_path}")
            subprocess.Popen(['xdg-open', str(self.current_json_path)])

    def on_click(self, event):
        if not self.is_viewing: return

        # Map click to original image coordinates
        x = int(event.x / self.scale_factor)
        y = int(event.y / self.scale_factor)
        w, h = self.orig_size

        # Create 10px box (5px radius)
        # Coordinate format: [y_min, x_min, y_max, x_max] (0-1000 scale)
        box = [
            int(max(0, y - 5) / h * 1000),
            int(max(0, x - 5) / (w- padding) * 1000),
            int(min(h, y + 5) / h * 1000),
            int(min(w, x + 5) / (w - padding) * 1000),
        ]

        box_str = "{ \"box_2d\": " + str(box) + ", \"label\": \"\" },"
        print(f"Copied box at ({x},{y}): {box_str}")

        self.root.clipboard_clear()
        self.root.clipboard_append(box_str)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python plotting_gui.py <directory_or_file>")
        sys.exit(1)

    input_path = Path(sys.argv[1])
    files_to_process = []

    if input_path.is_file():
        # File mode
        base_dir = input_path.parent
        stem = input_path.stem

        # Try to locate the image in Cutleft directory
        img_path = base_dir / "Cutleft" / f"{stem}.jpg"

        # Fallback: Check if user provided the jpg inside Cutleft directly
        if not img_path.exists() and input_path.parent.name == "Cutleft" and input_path.suffix.lower() == ".jpg":
             base_dir = input_path.parent.parent
             img_path = input_path

        if not img_path.exists():
            print(f"Error: Could not find image at {img_path}")
            sys.exit(1)

        files_to_process = [img_path]

    else:
        # Directory mode
        base_dir = input_path
        cutleft_dir = base_dir / "Cutleft"

        if not cutleft_dir.exists():
            print(f"Error: {cutleft_dir} does not exist.")
            sys.exit(1)

        files_to_process = sorted(cutleft_dir.glob("*.jpg"))

    try:
        all_labels = list(filter(None, (base_dir / "labels").read_text().splitlines()))
    except FileNotFoundError:
        all_labels = []

    # Start Processing Thread
    t = threading.Thread(target=worker_thread, args=(base_dir, files_to_process, all_labels))
    t.daemon = True # Kill thread if main app closes
    t.start()

    # Start GUI
    root = tk.Tk()
    app = ImageViewer(root, base_dir)
    root.mainloop()