ocr-screenshot-gallery/OCR-scripts/ocr_search_rofi.py

#!/usr/bin/env python3

import os
import sqlite3
import subprocess
import sys
import tempfile

# Configuration
DB_PATH = os.path.expanduser("~/screenshot_ocr.db")
SCREENSHOTS_DIR = os.path.expanduser("~/Screenshots")
ROFI_PROMPT = "OCR Search"


def check_dependencies():
    """Check if required dependencies are installed."""
    # Check if rofi is installed
    try:
        subprocess.run(
            ["which", "rofi"],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
    except subprocess.CalledProcessError:
        print(
            "Error: rofi is not installed. Please install it with: sudo pacman -S rofi"
        )
        return False

    # Check if database exists
    if not os.path.exists(DB_PATH):
        print(f"Error: Database not found at {DB_PATH}")
        print("Run the OCR script first to create and populate the database.")
        return False

    return True


def format_text(text, max_length=100):
    """Format OCR text for display in rofi."""
    # Replace newlines with spaces
    text = text.replace("\n", " ")

    # Replace multiple spaces with a single space
    while "  " in text:
        text = text.replace("  ", " ")

    # Truncate if too long
    if len(text) > max_length:
        text = text[:max_length] + "..."

    return text


def get_screenshot_data():
    """Get data from the SQLite database."""
    try:
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()

        # Get all records
        cursor.execute(
            "SELECT filename, ocr_text, full_path FROM ocr_results ORDER BY filename"
        )

        results = []
        for filename, ocr_text, full_path in cursor.fetchall():
            formatted_text = format_text(ocr_text)
            # Format for display: filename | ocr_text
            display_text = f"{filename} | {formatted_text}"
            results.append((display_text, full_path))

        conn.close()
        return results

    except sqlite3.Error as e:
        print(f"Database error: {e}")
        return []


def show_rofi_menu(items):
    """Display rofi menu with the given items and return selection."""
    if not items:
        print("No items to display.")
        return None, None

    # Create temporary file with menu items
    with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file:
        for display_text, _ in items:
            temp_file.write(f"{display_text}\n")
        temp_file_path = temp_file.name

    try:
        # Run rofi command
        cmd = [
            "rofi",
            "-dmenu",
            "-i",  # Case-insensitive matching
            "-p",
            ROFI_PROMPT,
            "-width",
            "80",
            "-lines",
            "15",
            "-font",
            "mono 10",
        ]

        process = subprocess.Popen(
            cmd,
            stdin=open(temp_file_path, "r"),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True,
        )

        stdout, stderr = process.communicate()

        if process.returncode != 0:
            print(f"Rofi error: {stderr}")
            return None, None

        selection = stdout.strip()
        if not selection:
            return None, None

        # Find the matching item
        for i, (display_text, full_path) in enumerate(items):
            if display_text == selection:
                return display_text, full_path

        return None, None

    finally:
        # Clean up temporary file
        os.unlink(temp_file_path)


def open_file(file_path):
    """Open the selected file with the default application."""
    if os.path.exists(file_path):
        try:
            subprocess.Popen(
                ["xdg-open", file_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE
            )
            print(f"Opening: {file_path}")
            return True
        except Exception as e:
            print(f"Error opening file: {e}")
            return False
    else:
        print(f"Error: File not found: {file_path}")
        return False


def main():
    """Main function."""
    if not check_dependencies():
        sys.exit(1)

    # Get screenshot data
    items = get_screenshot_data()
    if not items:
        print("No OCR data found in the database.")
        sys.exit(1)

    # Show rofi menu
    selection, file_path = show_rofi_menu(items)

    # Open selected file
    if selection and file_path:
        open_file(file_path)
    else:
        print("No selection made.")


if __name__ == "__main__":
    main()