#!/usr/bin/env python3
"""System tray app: recognize handwriting from a clipboard image or a file,
via the local ocr-daemon, and put the result on the clipboard and on screen.
"""

import json
import os
import shutil
import subprocess
import urllib.error
import urllib.request

import gi

gi.require_version("Gtk", "3.0")
gi.require_version("AyatanaAppIndicator3", "0.1")
from gi.repository import AyatanaAppIndicator3 as AppIndicator3
from gi.repository import Gdk, GdkPixbuf, Gtk

OCR_URL = os.environ.get("OCR_DAEMON_URL", "http://127.0.0.1:8734/ocr")

CONTENT_TYPE_BY_EXT = {
    "png": "image/png",
    "jpg": "image/jpeg",
    "jpeg": "image/jpeg",
    "webp": "image/webp",
    "bmp": "image/bmp",
    "gif": "image/gif",
}

# GTK3's Gtk.Clipboard fails to see image data that GNOME Shell's screenshot
# tool puts on the clipboard under Wayland (confirmed: wait_for_image()
# returns nothing for a screenshot that pastes fine elsewhere, and that
# wl-paste reads correctly). wl-clipboard talks to the compositor directly
# and doesn't have this gap, so it's used here when available, with GTK3's
# clipboard as a fallback for non-Wayland sessions that lack wl-clipboard.
IMAGE_MIME_PRIORITY = ["image/png", "image/jpeg", "image/bmp", "image/gif", "image/webp"]


def guess_content_type(path):
    ext = path.rsplit(".", 1)[-1].lower() if "." in path else ""
    return CONTENT_TYPE_BY_EXT.get(ext, "image/png")


def pixbuf_to_png_bytes(pixbuf):
    success, buf = pixbuf.save_to_bufferv("png", [], [])
    if not success:
        raise RuntimeError("failed to encode image as PNG")
    return bytes(buf)


def read_clipboard_image():
    """Returns (data: bytes, content_type: str), or None if no image is on the clipboard."""
    if shutil.which("wl-paste"):
        try:
            listing = subprocess.run(
                ["wl-paste", "--list-types"], capture_output=True, text=True, timeout=5
            )
        except subprocess.SubprocessError:
            listing = None
        types = listing.stdout.splitlines() if listing and listing.returncode == 0 else []
        for mime in IMAGE_MIME_PRIORITY:
            if mime in types:
                result = subprocess.run(["wl-paste", "--type", mime], capture_output=True, timeout=10)
                if result.returncode == 0 and result.stdout:
                    return result.stdout, mime
        return None

    clipboard = Gtk.Clipboard.get_default(Gdk.Display.get_default())
    pixbuf = clipboard.wait_for_image()
    if pixbuf is None:
        return None
    return pixbuf_to_png_bytes(pixbuf), "image/png"


def write_clipboard_text(text):
    if shutil.which("wl-copy"):
        try:
            subprocess.run(["wl-copy"], input=text.encode("utf-8"), timeout=5)
            return
        except subprocess.SubprocessError:
            pass
    Gtk.Clipboard.get_default(Gdk.Display.get_default()).set_text(text, -1)


def post_to_ocr_daemon(data, content_type):
    req = urllib.request.Request(
        OCR_URL, data=data, headers={"Content-Type": content_type}, method="POST"
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        return json.loads(resp.read().decode("utf-8"))


class OcrTrayApp:
    def __init__(self):
        self.result_window = None

        self.indicator = AppIndicator3.Indicator.new(
            "ocr-tray",
            "edit-find-symbolic",
            AppIndicator3.IndicatorCategory.APPLICATION_STATUS,
        )
        self.indicator.set_status(AppIndicator3.IndicatorStatus.ACTIVE)
        self.indicator.set_menu(self._build_menu())

    def _build_menu(self):
        menu = Gtk.Menu()

        paste_item = Gtk.MenuItem(label="Recognize from Clipboard")
        paste_item.connect("activate", self.on_paste_clicked)
        menu.append(paste_item)

        open_item = Gtk.MenuItem(label="Recognize from File...")
        open_item.connect("activate", self.on_open_clicked)
        menu.append(open_item)

        menu.append(Gtk.SeparatorMenuItem())

        quit_item = Gtk.MenuItem(label="Quit")
        quit_item.connect("activate", lambda _: Gtk.main_quit())
        menu.append(quit_item)

        menu.show_all()
        return menu

    def on_paste_clicked(self, _widget):
        found = read_clipboard_image()
        if found is None:
            self._show_error("No image found on the clipboard.")
            return
        data, content_type = found
        self._recognize(data, content_type)

    def on_open_clicked(self, _widget):
        dialog = Gtk.FileChooserDialog(
            title="Select an image", parent=None, action=Gtk.FileChooserAction.OPEN
        )
        dialog.add_buttons(
            Gtk.STOCK_CANCEL, Gtk.ResponseType.CANCEL, Gtk.STOCK_OPEN, Gtk.ResponseType.OK
        )
        image_filter = Gtk.FileFilter()
        image_filter.set_name("Images")
        for pattern in ("*.png", "*.jpg", "*.jpeg", "*.webp", "*.bmp", "*.gif"):
            image_filter.add_pattern(pattern)
        dialog.add_filter(image_filter)

        response = dialog.run()
        path = dialog.get_filename()
        dialog.destroy()

        if response != Gtk.ResponseType.OK or not path:
            return

        try:
            with open(path, "rb") as f:
                data = f.read()
        except OSError as e:
            self._show_error(f"Could not read {path}:\n{e}")
            return

        self._recognize(data, guess_content_type(path))

    def _recognize(self, data, content_type):
        try:
            result = post_to_ocr_daemon(data, content_type)
        except urllib.error.URLError as e:
            self._show_error(f"Could not reach the OCR daemon at {OCR_URL}.\n\n{e}")
            return
        except Exception as e:  # malformed response, timeout, etc.
            self._show_error(f"Unexpected error talking to the OCR daemon:\n{e}")
            return

        if result.get("error"):
            self._show_error(f"OCR daemon error: {result['error']}")
            return

        text = result.get("text", "")
        if not text:
            self._show_error("No handwriting detected in that image.")
            return

        write_clipboard_text(text)

        self._show_result(text)

    def _show_result(self, text):
        if self.result_window is not None:
            self.result_window.destroy()

        window = Gtk.Window(title="Recognized Text (copied to clipboard)")
        window.set_default_size(500, 350)
        window.connect("destroy", self._on_result_window_destroy)

        vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
        vbox.set_margin_start(10)
        vbox.set_margin_end(10)
        vbox.set_margin_top(10)
        vbox.set_margin_bottom(10)
        window.add(vbox)

        scrolled = Gtk.ScrolledWindow()
        scrolled.set_vexpand(True)
        textview = Gtk.TextView()
        textview.set_wrap_mode(Gtk.WrapMode.WORD)
        textview.get_buffer().set_text(text)
        scrolled.add(textview)
        vbox.pack_start(scrolled, True, True, 0)

        button_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=6)
        copy_button = Gtk.Button(label="Copy to Clipboard")
        copy_button.connect("clicked", self._on_copy_clicked, textview)
        button_box.pack_end(copy_button, False, False, 0)
        vbox.pack_start(button_box, False, False, 0)

        window.show_all()
        window.present()
        self.result_window = window

    def _on_result_window_destroy(self, _window):
        self.result_window = None

    def _on_copy_clicked(self, _button, textview):
        buf = textview.get_buffer()
        text = buf.get_text(buf.get_start_iter(), buf.get_end_iter(), True)
        write_clipboard_text(text)

    def _show_error(self, message):
        dialog = Gtk.MessageDialog(
            transient_for=None,
            flags=0,
            message_type=Gtk.MessageType.ERROR,
            buttons=Gtk.ButtonsType.OK,
            text=message,
        )
        dialog.run()
        dialog.destroy()


def main():
    OcrTrayApp()
    Gtk.main()


if __name__ == "__main__":
    main()
