A!leys Galerie

#!/usr/bin/env python3
"""
Ailey OCR Screenshot Organizer

Organizes screenshots based on OCR text content with custom folder structure and dark mode support.
"""

import os
import sys
import json
import tempfile
import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from PIL import Image, ImageTk, ImageGrab
import pytesseract
import pyperclip
from typing_extensions import Annotated

# Configuration
DEFAULT_CONFIG = {
    "screenshot_dir": "~/screenshots",
    "ocr_language": "eng",
    "theme": "dark",
    "folder_pattern": "{text_line1[:30]}/{text_line2[:30]}",
    "max_characters": 30,
    "sorting_key": "text",
    "case_sensitive": False
}

LOG_FORMAT = "%(asctime)s - %(levelname)s - %(message)s"

@dataclass
class ScreenshotData:
    """Data structure to hold OCR results and metadata"""
    filename: str
    text_lines: List[str]
    first_line: str
    second_line: str
    full_text: str
    folder_path: str

class ThemeManager:
    """Handles dark/light theme switching"""
    def __init__(self, root: tk.Tk):
        self.root = root
        self.themes = {
            "dark": {
                "bg": "#2b2b2b",
                "fg": "#ffffff",
                "button_bg": "#3d3d3d",
                "button_fg": "#ffffff",
                "frame_bg": "#3d3d3d",
                "label_fg": "#ffffff",
                "treeview_bg": "#2b2b2b",
                "treeview_fg": "#ffffff",
                "highlight": "#4d4d4d"
            },
            "light": {
                "bg": "#ffffff",
                "fg": "#000000",
                "button_bg": "#e0e0e0",
                "button_fg": "#000000",
                "frame_bg": "#e0e0e0",
                "label_fg": "#000000",
                "treeview_bg": "#ffffff",
                "treeview_fg": "#000000",
                "highlight": "#c0c0c0"
            }
        }

    def apply_theme(self, theme_name: str):
        """Apply selected theme to all widgets"""
        theme = self.themes.get(theme_name, self.themes["dark"])
        self.root.configure(bg=theme["bg"])

        # Configure all widgets with the theme
        for widget in self.root.winfo_children():
            if isinstance(widget, ttk.Frame):
                widget.configure(style="TFrame")
            elif isinstance(widget, ttk.Button):
                widget.configure(style="TButton")
            elif isinstance(widget, ttk.Label):
                widget.configure(style="TLabel")
            elif isinstance(widget, ttk.Treeview):
                widget.configure(style="Treeview")

        # Configure styles
        style = ttk.Style()
        style.theme_use("clam")
        style.configure("TFrame", background=theme["frame_bg"])
        style.configure("TButton",
                       background=theme["button_bg"],
                       foreground=theme["button_fg"],
                       focuscolor=theme["highlight"])
        style.configure("TLabel",
                       background=theme["bg"],
                       foreground=theme["label_fg"])
        style.configure("Treeview",
                       background=theme["treeview_bg"],
                       foreground=theme["treeview_fg"],
                       fieldbackground=theme["treeview_bg"],
                       rowbackground=theme["treeview_bg"],
                       bordercolor=theme["highlight"])
        style.map("Treeview",
                 background=[("selected", theme["highlight"])],
                 foreground=[("selected", theme["treeview_fg"])])
        style.configure("Dark.TFrame", background=theme["frame_bg"])
        style.configure("Dark.TLabel", background=theme["bg"], foreground=theme["label_fg"])

class ScreenshotOrganizer:
    """Main class for organizing screenshots by OCR content"""
    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("Ailey OCR Screenshot Organizer")
        self.root.geometry("800x600")
        self.root.resizable(True, True)

        # Initialize logging
        logging.basicConfig(level=logging.INFO, format=LOG_FORMAT)
        self.logger = logging.getLogger("ScreenshotOrganizer")

        # Configuration management
        self.config_file = Path.home() / ".screenshot_organizer_config.json"
        self.config = self._load_config()
        self.theme_manager = ThemeManager(root)
        self.theme_manager.apply_theme(self.config["theme"])

        # Data storage
        self.screenshots: List[ScreenshotData] = []
        self.current_screenshot_index = -1

        # GUI Setup
        self.setup_ui()

    def setup_ui(self):
        """Set up the user interface"""
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)

        # Control frame
        control_frame = ttk.Frame(main_frame)
        control_frame.pack(fill=tk.X, pady=(0, 10))

        # Theme selection
        theme_frame = ttk.Frame(control_frame)
        theme_frame.pack(side=tk.RIGHT, padx=(10, 0))

        ttk.Label(theme_frame, text="Theme:").pack(side=tk.LEFT)
        self.theme_var = tk.StringVar(value=self.config["theme"])
        theme_menu = ttk.OptionMenu(theme_frame, self.theme_var, *["dark", "light"])
        theme_menu.pack(side=tk.LEFT, padx=(0, 5))
        theme_menu.configure(width=8)

        # Mode selection
        mode_frame = ttk.Frame(control_frame)
        mode_frame.pack(side=tk.LEFT)

        ttk.Label(mode_frame, text="Mode:").pack(side=tk.LEFT)
        self.mode_var = tk.StringVar(value="batch")
        mode_menu = ttk.OptionMenu(mode_frame, self.mode_var, "batch", "single")
        mode_menu.pack(side=tk.LEFT, padx=(0, 5))
        mode_menu.configure(width=8)

        # Buttons
        button_frame = ttk.Frame(control_frame)
        button_frame.pack(side=tk.LEFT, padx=(10, 0))

        ttk.Button(button_frame, text="Capture Screenshot", command=self.capture_screenshot).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(button_frame, text="Process All", command=self.process_all_screenshots).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(button_frame, text="Save Config", command=self.save_config).pack(side=tk.LEFT, padx=(0, 5))

        # Preview frame
        preview_frame = ttk.Frame(main_frame)
        preview_frame.pack(fill=tk.BOTH, expand=True)

        # Treeview for file list
        self.tree = ttk.Treeview(preview_frame, columns=("Filename", "Folder", "First Line", "Second Line"), show="headings")
        self.tree.heading("Filename", text="Filename")
        self.tree.heading("Folder", text="Folder")
        self.tree.heading("First Line", text="First Line")
        self.tree.heading("Second Line", text="Second Line")
        self.tree.column("Filename", width=200)
        self.tree.column("Folder", width=200)
        self.tree.column("First Line", width=200)
        self.tree.column("Second Line", width=200)

        self.tree.pack(fill=tk.BOTH, expand=True, side=tk.LEFT, padx=(0, 10))

        # Scrollbar
        scrollbar = ttk.Scrollbar(preview_frame, orient=tk.VERTICAL, command=self.tree.yview)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.tree.configure(yscrollcommand=scrollbar.set)

        # Preview image
        self.preview_label = ttk.Label(preview_frame)
        self.preview_label.pack(fill=tk.BOTH, expand=True, side=tk.RIGHT)
        self.preview_label.config(borderwidth=2, relief=tk.SOLID)

        # Status bar
        self.status_var = tk.StringVar()
        self.status_var.set("Ready")
        status_bar = ttk.Label(main_frame, textvariable=self.status_var, relief=tk.SUNKEN)
        status_bar.pack(fill=tk.X, pady=(10, 0))

        # Bind treeview selection
        self.tree.bind("<<TreeviewSelect>>", self.on_tree_select)

        # Bind theme change
        self.theme_var.trace_add("write", lambda *args: self.on_theme_change())

    def on_theme_change(self, *args):
        """Handle theme changes"""
        new_theme = self.theme_var.get()
        self.config["theme"] = new_theme
        self.theme_manager.apply_theme(new_theme)

    def on_tree_select(self, event):
        """Handle treeview selection"""
        selected = self.tree.selection()
        if selected:
            index = int(selected[0][1:])  # Extract index from "I001"
            if index < len(self.screenshots):
                self.current_screenshot_index = index
                self.update_preview()

    def update_preview(self):
        """Update the image preview"""
        if self.current_screenshot_index >= 0 and self.current_screenshot_index < len(self.screenshots):
            data = self.screenshots[self.current_screenshot_index]
            if data.filename:
                try:
                    img = Image.open(data.filename)
                    img.thumbnail((300, 300))
                    photo = ImageTk.PhotoImage(img)
                    self.preview_label.config(image=photo)
                    self.preview_label.image = photo
                    self.status_var.set(f"Preview: {data.filename} - {len(data.text_lines)} lines")
                except Exception as e:
                    self.logger.error(f"Error loading preview: {e}")
                    self.preview_label.config(image=None)
                    self.preview_label.image = None
                    self.status_var.set("Error loading preview")
        else:
            self.preview_label.config(image=None)
            self.preview_label.image = None
            self.status_var.set("Select a screenshot to preview")

    def capture_screenshot(self):
        """Capture a screenshot from the screen"""
        try:
            self.status_var.set("Capturing screenshot...")
            self.root.update()

            # Get screen dimensions
            screen = ImageGrab.grab()
            width, height = screen.size

            # Ask for coordinates
            x = self.root.winfo_rootx() + self.root.winfo_width() + 10
            y = self.root.winfo_rooty()

            # Capture the screenshot
            screenshot = ImageGrab.grab(bbox=(x, y, x + width, y + height))

            # Save temporarily
            temp_file = Path(tempfile.mkstemp(suffix=".png")[1])
            screenshot.save(temp_file, "PNG")

            # Process the new screenshot
            self.process_screenshot(temp_file)

        except Exception as e:
            self.logger.error(f"Error capturing screenshot: {e}")
            messagebox.showerror("Error", f"Failed to capture screenshot: {e}")

    def process_screenshot(self, image_path: Path):
        """Process a single screenshot with OCR"""
        try:
            self.status_var.set(f"Processing: {image_path.name}...")
            self.root.update()

            # Perform OCR
            text = pytesseract.image_to_string(image_path, lang=self.config["ocr_language"])

            # Clean and parse text
            lines = [line.strip() for line in text.split("\n") if line.strip()]
            first_line = lines[0] if len(lines) > 0 else ""
            second_line = lines[1] if len(lines) > 1 else ""

            # Truncate lines if needed
            if len(first_line) > self.config["max_characters"]:
                first_line = first_line[:self.config["max_characters"]] + "..."
            if len(second_line) > self.config["max_characters"]:
                second_line = second_line[:self.config["max_characters"]] + "..."

            # Create folder path
            folder_pattern = self.config["folder_pattern"]
            folder_path = folder_pattern.format(
                text_line1=first_line,
                text_line2=second_line
            ).strip("/")

            # Add to data
            data = ScreenshotData(
                filename=str(image_path),
                text_lines=lines,
                first_line=first_line,
                second_line=second_line,
                full_text=text,
                folder_path=folder_path
            )

            self.screenshots.append(data)
            self.current_screenshot_index = len(self.screenshots) - 1

            # Update UI
            self.update_tree()
            self.update_preview()
            self.status_var.set(f"Processed: {image_path.name}")

            # Copy to clipboard if in single mode
            if self.mode_var.get() == "single":
                pyperclip.copy(image_path)
                self.status_var.set(f"Copied to clipboard: {image_path.name}")

        except Exception as e:
            self.logger.error(f"Error processing screenshot {image_path}: {e}")
            messagebox.showerror("Error", f"Failed to process screenshot: {e}")

    def process_all_screenshots(self):
        """Process all screenshots in the configured directory"""
        try:
            self.status_var.set("Processing all screenshots...")
            self.root.update()

            screenshot_dir = Path(self.config["screenshot_dir"]).expanduser()
            if not screenshot_dir.exists():
                screenshot_dir.mkdir(parents=True)

            # Process all PNG files in the directory
            for png_file in screenshot_dir.glob("*.png"):
                if png_file.name.startswith("screenshot_"):
                    self.process_screenshot(png_file)

            self.status_var.set(f"Processed all screenshots in {screenshot_dir}")
            messagebox.showinfo("Success", f"Processed all screenshots in {screenshot_dir}")

        except Exception as e:
            self.logger.error(f"Error processing all screenshots: {e}")
            messagebox.showerror("Error", f"Failed to process all screenshots: {e}")

    def update_tree(self):
        """Update the treeview with current screenshots"""
        self.tree.delete(*self.tree.get_children())
        for i, data in enumerate(self.screenshots):
            self.tree.insert("", "end", iid=f"I{i}", text="",
                           values=(data.filename, data.folder_path, data.first_line, data.second_line))

    def save_config(self):
        """Save current configuration"""
        try:
            self.config_file.write_text(json.dumps(self.config, indent=2))
            self.logger.info(f"Configuration saved to {self.config_file}")
            messagebox.showinfo("Success", f"Configuration saved to {self.config_file}")
        except Exception as e:
            self.logger.error(f"Error saving configuration: {e}")
            messagebox.showerror("Error", f"Failed to save configuration: {e}")

    def _load_config(self) -> Dict:
        """Load configuration from file"""
        if self.config_file.exists():
            try:
                with open(self.config_file, "r") as f:
                    config = json.load(f)
                    # Validate config
                    for key in DEFAULT_CONFIG:
                        if key not in config:
                            config[key] = DEFAULT_CONFIG[key]
                    return config
            except Exception as e:
                self.logger.warning(f"Error loading config: {e}. Using defaults.")
        return DEFAULT_CONFIG.copy()

def main():
    """Main entry point"""
    try:
        # Check if pytesseract is installed
        pytesseract.get_tesseract_version()
    except Exception as e:
        print(f"Error: Tesseract OCR not installed. Please install it from https://github.com/tesseract-ocr/tesseract")
        print(f"Then run: pip install pytesseract")
        sys.exit(1)

    root = tk.Tk()
    app = ScreenshotOrganizer(root)
    root.mainloop()

if __name__ == "__main__":
    main()