gui/gui.py

"""
EL HELAL Studio – Professional Photo Workflow
Integrated Pipeline with Memory Optimization and Manual Crop
"""

import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from PIL import Image, ImageTk, ImageOps
import os
import threading
from pathlib import Path
import time
import sys

# Add core directory to python path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'core')))

# Import our processing tools
import crop
import process_images
import color_steal
from layout_engine import generate_layout


class CropDialog(tk.Toplevel):
    """Interactive window to adjust crop manually."""
    def __init__(self, parent, image_path, current_rect, callback):
        super().__init__(parent)
        self.title("Adjust Crop (5:7 Aspect Ratio)")
        self.image_path = image_path
        self.callback = callback
        
        # Load original image for reference
        self.orig_pil = Image.open(image_path)
        self.w, self.h = self.orig_pil.size
        
        # Scale for display
        screen_h = self.winfo_screenheight()
        target_h = int(screen_h * 0.8)
        self.display_scale = min(target_h / self.h, 1.0)
        
        self.display_w = int(self.w * self.display_scale)
        self.display_h = int(self.h * self.display_scale)
        
        self.display_img = self.orig_pil.resize((self.display_w, self.display_h), Image.LANCZOS)
        self.tk_img = ImageTk.PhotoImage(self.display_img)
        
        # Rect in original coordinates (x1, y1, x2, y2)
        self.rect = list(current_rect) if current_rect else [0, 0, 100, 140]
        
        # UI Layout
        self.canvas = tk.Canvas(self, width=self.display_w, height=self.display_h, bg="black", highlightthickness=0)
        self.canvas.pack(pady=10, padx=10)
        self.canvas.create_image(0, 0, image=self.tk_img, anchor=tk.NW)
        
        self.rect_id = self.canvas.create_rectangle(0, 0, 0, 0, outline="yellow", width=3)
        self._update_canvas_rect()
        
        ctrl = tk.Frame(self)
        ctrl.pack(fill=tk.X, padx=20, pady=5)
        tk.Label(ctrl, text="Drag the yellow box to move the crop. The size is fixed to 5:7.", font=("Arial", 10, "italic")).pack()
        
        btn_frame = tk.Frame(self)
        btn_frame.pack(pady=15)
        tk.Button(btn_frame, text=" Cancel ", command=self.destroy, width=10).pack(side=tk.LEFT, padx=10)
        tk.Button(btn_frame, text=" Apply & Reprocess ", bg="#27ae60", fg="white", 
                  command=self._apply, width=20, font=("Arial", 10, "bold")).pack(side=tk.LEFT, padx=10)
        
        # Mouse Events
        self.canvas.bind("<B1-Motion>", self._on_drag)
        self.canvas.bind("<Button-1>", self._on_click)
        
        # Center the window
        self.update_idletasks()
        wx = (self.winfo_screenwidth() - self.winfo_width()) // 2
        wy = (self.winfo_screenheight() - self.winfo_height()) // 2
        self.geometry(f"+{wx}+{wy}")
        
        self.grab_set() # Modal

    def _update_canvas_rect(self):
        x1, y1, x2, y2 = self.rect
        self.canvas.coords(self.rect_id, 
                           x1 * self.display_scale, y1 * self.display_scale,
                           x2 * self.display_scale, y2 * self.display_scale)

    def _on_click(self, event):
        self.start_x = event.x
        self.start_y = event.y

    def _on_drag(self, event):
        dx = (event.x - self.start_x) / self.display_scale
        dy = (event.y - self.start_y) / self.display_scale
        
        rw = self.rect[2] - self.rect[0]
        rh = self.rect[3] - self.rect[1]
        
        nx1 = self.rect[0] + dx
        ny1 = self.rect[1] + dy
        
        if nx1 < 0: nx1 = 0
        if ny1 < 0: ny1 = 0
        if nx1 + rw > self.w: nx1 = self.w - rw
        if ny1 + rh > self.h: ny1 = self.h - rh
        
        self.rect[0] = nx1
        self.rect[1] = ny1
        self.rect[2] = nx1 + rw
        self.rect[3] = ny1 + rh
        
        self.start_x = event.x
        self.start_y = event.y
        self._update_canvas_rect()

    def _apply(self):
        self.callback(tuple(map(int, self.rect)))
        self.destroy()


class StudioApp:
    """Main application with memory-efficient batch handling."""
    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("EL HELAL Studio — Professional Workflow")
        self.root.minsize(1000, 900)
        self.root.configure(bg="#f0f0f0")

        self._image_data: list[dict] = []
        self._current_index = 0
        self._phase = "empty"  # "empty" | "input" | "preview"
        
        self.model = None
        self.transform = None
        self.luts = color_steal.load_trained_curves()
        self.is_model_ready = False

        self._build_ui()
        self.root.bind("<Configure>", self._on_resize)
        
        threading.Thread(target=self._warm_up_model, daemon=True).start()

    def _warm_up_model(self):
        try:
            self._set_status("Initializing AI Engine...")
            self.model, _ = process_images.setup_model()
            self.transform = process_images.get_transform()
            self.is_model_ready = True
            self._set_status("AI Engine Ready.")
        except Exception as e:
            self._set_status(f"Critical Error: AI Model failed to load ({e})")

    def _build_ui(self):
        # Header
        header = tk.Frame(self.root, bg="#1a2634", pady=15)
        header.pack(fill=tk.X)
        tk.Label(header, text="EL HELAL Studio", font=("Arial", 26, "bold"), fg="#e8b923", bg="#1a2634").pack()
        tk.Label(header, text="Memory Optimized Pipeline: Auto-Crop | AI Background | Color Grade | Layout", 
                 font=("Arial", 10), fg="white", bg="#1a2634").pack()

        # Input Area
        input_frame = tk.Frame(self.root, bg="#f0f0f0", pady=10)
        input_frame.pack(fill=tk.X, padx=20)
        tk.Label(input_frame, text="Name (الاسم):", font=("Arial", 11, "bold"), bg="#f0f0f0").pack(side=tk.LEFT, padx=(0, 5))
        self.entry_name = tk.Entry(input_frame, font=("Arial", 14), width=30, justify=tk.RIGHT)
        self.entry_name.pack(side=tk.LEFT, padx=(0, 25))
        tk.Label(input_frame, text="ID (الرقم):", font=("Arial", 11, "bold"), bg="#f0f0f0").pack(side=tk.LEFT, padx=(0, 5))
        self.entry_id = tk.Entry(input_frame, font=("Arial", 14), width=20)
        self.entry_id.pack(side=tk.LEFT)

        # Toolbar
        toolbar = tk.Frame(self.root, bg="#f0f0f0", pady=10)
        toolbar.pack(fill=tk.X, padx=20)
        self.btn_open = tk.Button(toolbar, text=" 📂 Select Photos ", command=self._open_files, 
                                  bg="#3498db", fg="white", relief=tk.FLAT, padx=15, pady=8, font=("Arial", 10, "bold"))
        self.btn_open.pack(side=tk.LEFT, padx=5)
        self.btn_process = tk.Button(toolbar, text=" ⚡ Start All ", command=self._process_all, 
                                     bg="#e67e22", fg="white", relief=tk.FLAT, padx=15, pady=8, 
                                     font=("Arial", 10, "bold"), state=tk.DISABLED)
        self.btn_process.pack(side=tk.LEFT, padx=5)
        self.btn_save = tk.Button(toolbar, text=" 💾 Save All ", command=self._save_all, 
                                  bg="#27ae60", fg="white", relief=tk.FLAT, padx=15, pady=8, 
                                  font=("Arial", 10, "bold"), state=tk.DISABLED)
        self.btn_save.pack(side=tk.LEFT, padx=5)
        self.btn_edit_crop = tk.Button(toolbar, text=" ✂ Edit Crop ", command=self._edit_crop, 
                                       bg="#95a5a6", fg="white", relief=tk.FLAT, padx=15, pady=8, 
                                       font=("Arial", 10, "bold"), state=tk.DISABLED)
        self.btn_edit_crop.pack(side=tk.LEFT, padx=(30, 0))

        # Navigation
        nav_frame = tk.Frame(toolbar, bg="#f0f0f0"); nav_frame.pack(side=tk.RIGHT)
        self.btn_prev = tk.Button(nav_frame, text=" ◀ ", command=self._prev_image, state=tk.DISABLED, width=4)
        self.btn_prev.pack(side=tk.LEFT, padx=2)
        self.lbl_counter = tk.Label(nav_frame, text="", font=("Arial", 11, "bold"), bg="#f0f0f0", width=8)
        self.lbl_counter.pack(side=tk.LEFT)
        self.btn_next = tk.Button(nav_frame, text=" ▶ ", command=self._next_image, state=tk.DISABLED, width=4)
        self.btn_next.pack(side=tk.LEFT, padx=2)

        # Progress
        self.progress = ttk.Progressbar(self.root, orient=tk.HORIZONTAL, mode='determinate')
        self.progress.pack(fill=tk.X, padx=20, pady=(0, 10))

        # Canvas
        canvas_frame = tk.Frame(self.root, bg="#ddd", bd=1, relief=tk.SUNKEN)
        canvas_frame.pack(fill=tk.BOTH, expand=True, padx=20, pady=5)
        self.canvas = tk.Canvas(canvas_frame, bg="white", highlightthickness=0)
        self.canvas.pack(fill=tk.BOTH, expand=True)

        self.status = tk.Label(self.root, text="Ready", font=("Arial", 10), bg="#1a2634", 
                               fg="white", anchor=tk.W, padx=15, pady=8)
        self.status.pack(fill=tk.X, side=tk.BOTTOM)

    # ── Operations ──

    def _open_files(self):
        paths = filedialog.askopenfilenames(
            title="Select Student Photos",
            filetypes=[("Image files", "*.jpg *.jpeg *.png *.bmp *.tiff")]
        )
        if not paths: return
        
        self._image_data = []
        for path in paths:
            # Generate thumbnail for preview to save memory
            try:
                with Image.open(path) as img:
                    img.thumbnail((800, 800), Image.LANCZOS)
                    thumb = img.copy()
                
                self._image_data.append({
                    "path": path,
                    "name": "",
                    "id": "",
                    "thumb": thumb,  # Store small version only
                    "result": None,
                    "crop_rect": None
                })
            except Exception as e:
                print(f"Error loading {path}: {e}")
            
        self._current_index = 0
        self._phase = "input"
        self._load_current_fields()
        self._update_nav()
        self._update_ui_state()
        self._show_current()
        self._set_status(f"Loaded {len(self._image_data)} photos. Memory usage optimized.")

    def _process_all(self):
        self._save_current_fields()
        if not self.is_model_ready:
            messagebox.showwarning("Wait", "AI Model is still loading.")
            return
            
        self.progress['value'] = 0
        self.progress['maximum'] = len(self._image_data)
        self.btn_open.config(state=tk.DISABLED)
        self.btn_process.config(state=tk.DISABLED)
        
        threading.Thread(target=self._run_pipeline_batch, daemon=True).start()

    def _run_pipeline_batch(self):
        total = len(self._image_data)
        for i in range(total):
            self.root.after(0, self._set_status, f"Processing {i+1}/{total}...")
            self._process_single_image(i)
            self.root.after(0, lambda v=i+1: self.progress.configure(value=v))
        self.root.after(0, self._on_batch_done)

    def _process_single_image(self, idx):
        data = self._image_data[idx]
        try:
            # 1. CROP (Face Detection) - Loads full resolution from disk
            temp_crop = f"temp_{idx}_{int(time.time())}.jpg"
            if not data["crop_rect"]:
                data["crop_rect"] = crop.get_auto_crop_rect(data["path"])
            
            if not crop.apply_custom_crop(data["path"], temp_crop, data["crop_rect"]):
                raise Exception("Crop failed")
            
            cropped_img = Image.open(temp_crop)
            
            # 2. BACKGROUND REMOVAL (AI)
            trans = process_images.remove_background(self.model, cropped_img, self.transform)
            
            # 3. COLOR GRADING
            graded = color_steal.apply_to_image(self.luts, trans) if self.luts else trans
            
            # 4. FINAL LAYOUT
            data["result"] = generate_layout(graded, data["name"], data["id"])
            
            if os.path.exists(temp_crop): os.remove(temp_crop)
            
        except Exception as e:
            print(f"Error processing index {idx}: {e}")
            data["result"] = None

    def _on_batch_done(self):
        self._phase = "preview"
        self.btn_open.config(state=tk.NORMAL)
        self.btn_process.config(state=tk.NORMAL)
        for i, d in enumerate(self._image_data):
            if d["result"]:
                self._current_index = i
                break
        self._update_ui_state()
        self._load_current_fields()
        self._update_nav()
        self._show_current()
        self._set_status("Processing complete.")

    def _edit_crop(self):
        data = self._image_data[self._current_index]
        if not data["crop_rect"]:
            data["crop_rect"] = crop.get_auto_crop_rect(data["path"])
        
        def on_apply(new_rect):
            data["crop_rect"] = new_rect
            self._set_status("Reprocessing image...")
            threading.Thread(target=self._reprocess_current, daemon=True).start()

        CropDialog(self.root, data["path"], data["crop_rect"], on_apply)

    def _reprocess_current(self):
        self._process_single_image(self._current_index)
        self.root.after(0, self._show_current)
        self.root.after(0, self._set_status, "Reprocess Done.")

    # ── UI Sync ──

    def _update_ui_state(self):
        if not self._image_data:
            self.btn_process.config(state=tk.DISABLED); self.btn_save.config(state=tk.DISABLED); self.btn_edit_crop.config(state=tk.DISABLED)
            return
        self.btn_process.config(state=tk.NORMAL)
        self.btn_save.config(state=tk.NORMAL if self._phase == "preview" else tk.DISABLED)
        self.btn_edit_crop.config(state=tk.NORMAL)

    def _show_current(self):
        if not self._image_data: return
        data = self._image_data[self._current_index]
        
        if self._phase == "preview" and data["result"]:
            img = data["result"]
        else:
            # Use thumbnail for navigation preview to stay memory-efficient
            img = data["thumb"]
            
        self._show_image_on_canvas(img)

    def _show_image_on_canvas(self, img):
        self.canvas.update_idletasks()
        cw, ch = self.canvas.winfo_width(), self.canvas.winfo_height()
        if cw < 10 or ch < 10: return
        
        iw, ih = img.size
        scale = min(cw/iw, ch/ih, 1.0)
        preview = img.resize((int(iw*scale), int(ih*scale)), Image.LANCZOS)
        self.tk_preview = ImageTk.PhotoImage(preview)
        self.canvas.delete("all")
        self.canvas.create_image(cw//2, ch//2, image=self.tk_preview, anchor=tk.CENTER)

    def _save_current_fields(self):
        if not self._image_data: return
        data = self._image_data[self._current_index]
        data["name"] = self.entry_name.get().strip()
        data["id"] = self.entry_id.get().strip()

    def _load_current_fields(self):
        if not self._image_data: return
        data = self._image_data[self._current_index]
        self.entry_name.delete(0, tk.END); self.entry_name.insert(0, data["name"])
        self.entry_id.delete(0, tk.END); self.entry_id.insert(0, data["id"])

    def _prev_image(self):
        if self._current_index > 0:
            self._save_current_fields(); self._current_index -= 1
            self._load_current_fields(); self._update_nav(); self._show_current()

    def _next_image(self):
        if self._current_index < len(self._image_data) - 1:
            self._save_current_fields(); self._current_index += 1
            self._load_current_fields(); self._update_nav(); self._show_current()

    def _update_nav(self):
        n = len(self._image_data)
        self.btn_prev.config(state=tk.NORMAL if self._current_index > 0 else tk.DISABLED)
        self.btn_next.config(state=tk.NORMAL if self._current_index < n-1 else tk.DISABLED)
        self.lbl_counter.config(text=f"{self._current_index+1} / {n}" if n else "")

    def _save_all(self):
        results = [d for d in self._image_data if d["result"]]
        if not results: return
        folder = filedialog.askdirectory(title="Select Save Folder")
        if not folder: return
        count = 0
        for d in results:
            try:
                name = Path(d["path"]).stem + "_layout.jpg"
                d["result"].save(os.path.join(folder, name), "JPEG", quality=95, dpi=(300, 300))
                count += 1
            except: pass
        messagebox.showinfo("Success", f"Saved {count} layouts.")

    def _on_resize(self, event):
        if self._image_data: self._show_current()

    def _set_status(self, msg):
        self.status.config(text=msg); self.root.update_idletasks()


if __name__ == "__main__":
    root = tk.Tk()
    w, h = 1100, 950
    root.geometry(f"{w}x{h}+{(root.winfo_screenwidth()-w)//2}+{(root.winfo_screenheight()-h)//2}")
    StudioApp(root); root.mainloop()