"""
Simple Screenshot Tool — cross-platform, pip-installable.

After installing:
    pip install screenshot-saver

Run with:
    screenshot-saver          # console entry-point
    python -m screenshot_saver  # module entry-point
"""

# ---- stdlib ----
import os
import sys
import time
import json
import signal
import atexit
import threading
import queue
import base64
import hashlib
import html as _html_mod
from io import BytesIO
from datetime import datetime
from pathlib import Path

# ---- third-party ----
try:
    from PIL import ImageGrab
    _HAS_IMAGEGRAB = True
except Exception:
    ImageGrab = None
    _HAS_IMAGEGRAB = False

try:
    from pynput import keyboard
    _HAS_PYNPUT = True
except Exception:
    keyboard = None
    _HAS_PYNPUT = False

# ---- tkinter ----
import tkinter as tk
from tkinter import scrolledtext, messagebox, filedialog


# ---------------------------------------------------------------------------
# Platform-aware paths
# ---------------------------------------------------------------------------

def _app_data_dir() -> Path:
    """Return the platform-appropriate directory for app config/data."""
    if sys.platform == "win32":
        base = Path(os.environ.get("APPDATA", Path.home()))
    elif sys.platform == "darwin":
        base = Path.home() / "Library" / "Application Support"
    else:
        base = Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config"))
    return base / "screenshot_saver"


def _default_screenshots_dir() -> Path:
    """Return a sensible default folder for saving screenshots."""
    pictures = Path.home() / "Pictures"
    if pictures.exists():
        return pictures / "ScreenshotSaver"
    return Path.home() / "Screenshots"


_CONFIG_DIR = _app_data_dir()
_CONFIG_FILE = _CONFIG_DIR / "config.json"


def _load_config() -> dict:
    try:
        if _CONFIG_FILE.exists():
            with open(_CONFIG_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)
            return data if isinstance(data, dict) else {}
    except Exception:
        pass
    return {}


def _save_config(cfg: dict) -> bool:
    try:
        _CONFIG_DIR.mkdir(parents=True, exist_ok=True)
        with open(_CONFIG_FILE, "w", encoding="utf-8") as f:
            json.dump(cfg, f, indent=2)
        return True
    except Exception:
        return False


def _ensure_imagegrab():
    if not _HAS_IMAGEGRAB:
        raise RuntimeError(
            "Pillow ImageGrab is not available in this environment.\n"
            "Fix:\n"
            "  pip uninstall pillow -y\n"
            "  pip install --upgrade pillow\n"
        )


def _ensure_pynput():
    if not _HAS_PYNPUT:
        raise RuntimeError(
            "Missing dependency: pynput\n"
            "Install it with:\n"
            "  pip install pynput\n"
        )


def _load_or_prompt_paths():
    """
    First run: ask once where to save screenshots + HTML; persist to config.
    Later runs: reuse saved values (no prompt).
    """
    cfg = _load_config()
    try:
        if cfg.get("screenshots_folder") and cfg.get("html_file"):
            return os.path.abspath(cfg["screenshots_folder"]), os.path.abspath(cfg["html_file"])
    except Exception:
        pass

    screenshots_folder = str(_default_screenshots_dir())
    html_file = str(_default_screenshots_dir() / "screenshots_document.html")

    try:
        r = tk.Tk()
        r.withdraw()
        r.attributes("-topmost", True)

        chosen_dir = filedialog.askdirectory(
            title="Choose folder to save screenshots",
            initialdir=screenshots_folder,
            mustexist=False,
        )
        if chosen_dir:
            screenshots_folder = os.path.abspath(chosen_dir)

        chosen_html = filedialog.asksaveasfilename(
            title="Choose HTML document file",
            initialdir=screenshots_folder,
            initialfile="screenshots_document.html",
            defaultextension=".html",
            filetypes=[("HTML files", "*.html"), ("All files", "*.*")],
        )
        if chosen_html:
            html_file = os.path.abspath(chosen_html)
        else:
            html_file = os.path.abspath(os.path.join(screenshots_folder, "screenshots_document.html"))

        r.destroy()
    except Exception:
        screenshots_folder = str(_default_screenshots_dir())
        html_file = str(_default_screenshots_dir() / "screenshots_document.html")

    cfg["screenshots_folder"] = screenshots_folder
    cfg["html_file"] = html_file
    cfg.setdefault("ask_annotation_after_capture", True)
    _save_config(cfg)

    return screenshots_folder, html_file


# ---------------------------------------------------------------------------
# Main application class
# ---------------------------------------------------------------------------

class SimpleScreenshotTool:
    def __init__(self):
        _ensure_imagegrab()
        _ensure_pynput()

        self.screenshots_folder, self.html_file = _load_or_prompt_paths()
        self.latest_html_file = self.html_file

        cfg = _load_config()
        self.ask_annotation_after_capture = bool(cfg.get("ask_annotation_after_capture", True))

        self.screenshot_count = 0
        self.screenshots_data = []

        os.makedirs(self.screenshots_folder, exist_ok=True)

        self.current_keys = set()
        self.listener = None
        self.is_monitoring = False

        self.clipboard_thread = None
        self.clipboard_stop_event = threading.Event()
        self.last_clipboard_hash = None
        self._pending_clipboard_image = None

        self.clipboard_poll_interval_active = 0.8
        self.clipboard_poll_interval_bg = 3.0
        self.low_cpu_background = True
        self.minimized_to_background = False
        self.window_is_iconic = False

        self._ui_queue = queue.Queue()
        self.data_lock = threading.Lock()
        self._html_write_lock = threading.Lock()

        self.html_dirty = False
        self._processing_cond = threading.Condition()
        self._inflight_processing = 0

        self._active_annotation = None
        self._active_annotation_dialog = None
        self._pending_annotations = []
        self._emergency_saved = False

        self.setup_gui()
        self._start_listener()
        self._register_exit_handlers()

    # ---- Color palette ----
    _CLR_BG        = "#E6E2D8"
    _CLR_SURFACE   = "#D8D4CA"
    _CLR_ACCENT    = "#2E473B"
    _CLR_GREEN     = "#2F4F3F"
    _CLR_GREEN_H   = "#254030"
    _CLR_RED       = "#5C3A3A"
    _CLR_RED_H     = "#4A2E2E"
    _CLR_ORANGE    = "#C2A878"
    _CLR_ORANGE_H  = "#A89060"
    _CLR_TEAL      = "#2F4F3F"
    _CLR_TEAL_H    = "#254030"
    _CLR_TEXT       = "#1A1A1A"
    _CLR_MUTED     = "#555555"
    _CLR_INFO_BG   = "#D8D4CA"
    _CLR_COUNTER   = "#1A1A1A"
    _CLR_BORDER    = "#2E473B"

    def setup_gui(self):
        self.root = tk.Tk()
        self.root.title("Simple Screenshot Tool")
        self.root.geometry("640x580")
        self.root.configure(bg=self._CLR_BG)
        self.root.resizable(True, True)

        outer_frame = tk.Frame(self.root, bg=self._CLR_BORDER, bd=10)
        outer_frame.pack(fill="both", expand=True, padx=10, pady=10)

        content_frame = tk.Frame(outer_frame, bg=self._CLR_BG, bd=2, relief="groove")
        content_frame.pack(fill="both", expand=True)

        tk.Label(
            content_frame, text="\U0001f4f7  Simple ScreenShot Tool",
            font=("Arial", 22, "bold"), fg=self._CLR_TEXT, bg=self._CLR_BG,
        ).pack(pady=(20, 5))

        self.status_label = tk.Label(
            content_frame, text="Status: Ready to Start",
            fg=self._CLR_ACCENT, bg=self._CLR_BG, font=("Arial", 14),
        )
        self.status_label.pack(pady=(0, 15))

        info_frame = tk.Frame(content_frame, bg=self._CLR_INFO_BG, bd=1, relief="solid")
        info_frame.pack(pady=10, padx=30, fill="x")
        info_text = (
            f"\U0001f4c1  Screenshots \u2192 {os.path.abspath(self.screenshots_folder)}\n"
            f"\U0001f4c4  HTML doc     \u2192 {os.path.abspath(self.html_file)}"
        )
        tk.Label(
            info_frame, text=info_text, justify="left",
            font=("Arial", 9), fg=self._CLR_TEXT, bg=self._CLR_INFO_BG,
            padx=10, pady=8,
        ).pack(anchor="w")

        self.counter_label = tk.Label(
            content_frame, text="Screenshots captured: 0",
            font=("Arial", 12, "bold"), fg=self._CLR_COUNTER, bg=self._CLR_BG,
        )
        self.counter_label.pack(pady=10)

        def _btn(parent, text, command, bg, hover, fg="white", **kw):
            b = tk.Button(
                parent, text=text, command=command,
                bg=bg, fg=fg, activebackground=hover, activeforeground=fg,
                font=("Arial", 10, "bold"), relief="raised", bd=2,
                cursor="hand2", padx=kw.pop("padx", 15), pady=kw.pop("pady", 6),
                **kw,
            )
            b.bind("<Enter>", lambda e, h=hover: b.configure(bg=h) if str(b['state']) != 'disabled' else None)
            b.bind("<Leave>", lambda e, n=bg: b.configure(bg=n) if str(b['state']) != 'disabled' else None)
            return b

        btn_frame = tk.Frame(content_frame, bg=self._CLR_BG)
        btn_frame.pack(pady=5)
        self.start_btn = _btn(btn_frame, "\u25b7 Start Monitoring", self.start_monitoring,
                              self._CLR_GREEN, self._CLR_GREEN_H)
        self.start_btn.pack(side="left", padx=10)
        self.stop_btn = _btn(btn_frame, "\u25fb Stop Monitoring", self.stop_monitoring,
                             self._CLR_GREEN, self._CLR_GREEN_H, state="disabled")
        self.stop_btn.pack(side="left", padx=10)

        low_cpu_frame = tk.Frame(content_frame, bg=self._CLR_BG)
        low_cpu_frame.pack(pady=4)
        self.low_cpu_var = tk.BooleanVar(value=self.low_cpu_background)
        tk.Checkbutton(
            low_cpu_frame, text="Low-CPU background mode",
            variable=self.low_cpu_var, onvalue=True, offvalue=False,
            command=self.on_low_cpu_toggle,
            fg=self._CLR_TEXT, bg=self._CLR_BG, selectcolor="#ffffff",
            activebackground=self._CLR_BG, font=("Arial", 9),
        ).pack()

        ann_frame = tk.Frame(content_frame, bg=self._CLR_BG)
        ann_frame.pack(pady=4)
        self.ask_annotation_var = tk.BooleanVar(value=self.ask_annotation_after_capture)
        tk.Checkbutton(
            ann_frame, text="Ask for text after each screenshot",
            variable=self.ask_annotation_var, onvalue=True, offvalue=False,
            command=self.on_annotation_toggle,
            fg=self._CLR_TEXT, bg=self._CLR_BG, selectcolor="#ffffff",
            activebackground=self._CLR_BG, font=("Arial", 9),
        ).pack()

        minimize_frame = tk.Frame(content_frame, bg=self._CLR_BG)
        minimize_frame.pack(pady=6)
        _btn(minimize_frame, "\U0001f5d4 Minimize to Background", self.hide_to_background,
             self._CLR_TEAL, self._CLR_TEAL_H, padx=10, pady=3).pack()

        open_frame = tk.Frame(content_frame, bg=self._CLR_BG)
        open_frame.pack(pady=10)
        _btn(open_frame, "\U0001f4c2 Open Folder", self.open_folder,
             self._CLR_ORANGE, self._CLR_ORANGE_H, fg="#1A1A1A").pack(side="left", padx=10)
        _btn(open_frame, "\U0001f4c4 Open HTML Doc", self.open_html,
             self._CLR_ORANGE, self._CLR_ORANGE_H, fg="#1A1A1A").pack(side="left", padx=10)

        tk.Label(
            content_frame,
            text="Press Win+Shift+S (Snip) to capture \u2014 snips auto-saved from clipboard\n"
                 "Text prompt can be toggled above",
            justify="center", font=("Arial", 9), fg=self._CLR_MUTED, bg=self._CLR_BG,
        ).pack(pady=10)

        log_frame = tk.Frame(content_frame, bg=self._CLR_SURFACE, bd=1, relief="solid")
        log_frame.pack(pady=(5, 15), padx=30, fill="both", expand=True)
        tk.Label(log_frame, text="Activity Log", font=("Arial", 10, "bold"),
                 fg=self._CLR_TEXT, bg=self._CLR_SURFACE).pack(pady=(5, 0))
        self.log_text = scrolledtext.ScrolledText(
            log_frame, height=5, font=("Consolas", 9), bg=self._CLR_SURFACE,
            fg=self._CLR_TEXT, relief="flat", bd=0, wrap="word",
        )
        self.log_text.pack(pady=(2, 5), padx=(10, 0), fill="both", expand=True)

        self.log("Welcome! Click 'Start Monitoring' to begin.")
        self.log(f"Screenshots will be saved to: {self.screenshots_folder}/")

        self.root.bind('<Unmap>', self._on_window_state_change)
        self.root.bind('<Map>', self._on_window_state_change)
        self.root.bind('<Activate>', self._on_window_activated)
        self.root.after(100, self._process_ui_queue)

    # ---- logging ----

    def log(self, message):
        timestamp = datetime.now().strftime("%H:%M:%S")
        line = f"[{timestamp}] {message}\n"
        if hasattr(self, "root") and self.root:
            if threading.current_thread() is threading.main_thread():
                self._append_log_line(line)
            else:
                try:
                    self._ui_queue.put(lambda l=line: self._append_log_line(l))
                except Exception:
                    print(line, end="")
        else:
            print(line, end="")

    def _append_log_line(self, line):
        try:
            self.log_text.insert(tk.END, line)
            self.log_text.see(tk.END)
        except Exception:
            pass

    def _process_ui_queue(self):
        try:
            while True:
                func = self._ui_queue.get_nowait()
                try:
                    func()
                finally:
                    self._ui_queue.task_done()
        except queue.Empty:
            pass
        finally:
            try:
                self.root.after(100, self._process_ui_queue)
            except Exception:
                pass

    # ---- monitoring controls ----

    def start_monitoring(self):
        self._start_listener()
        self.start_clipboard_monitor()
        self.is_monitoring = True
        self.start_btn.config(state='disabled')
        self.stop_btn.config(state='normal')
        self.status_label.config(text="Status: Monitoring Active \u2713", fg="#2E8B57")
        self.log("\U0001f3af Monitoring started! Press Win+Shift+S to capture.")
        self.log("\u26a1 Clipboard monitor active \u2014 snips will be auto-saved.")

    def stop_monitoring(self):
        if self.listener:
            self.stop_clipboard_monitor()
            self.is_monitoring = False
            self.start_btn.config(state='normal')
            self.stop_btn.config(state='disabled')
            self.status_label.config(text="Status: Monitoring Stopped", fg="#C0392B")
            self.log("\u23f8 Monitoring stopped.")
        if self.html_dirty:
            self.flush_html(reason="Stop Monitoring", timeout=5.0)

    # ---- settings toggles ----

    def on_annotation_toggle(self):
        try:
            self.ask_annotation_after_capture = bool(self.ask_annotation_var.get())
        except Exception:
            self.ask_annotation_after_capture = True
        cfg = _load_config()
        cfg.setdefault("screenshots_folder", os.path.abspath(self.screenshots_folder))
        cfg.setdefault("html_file", os.path.abspath(self.html_file))
        cfg["ask_annotation_after_capture"] = self.ask_annotation_after_capture
        _save_config(cfg)
        state = "ON" if self.ask_annotation_after_capture else "OFF"
        self.log(f"\U0001f4dd Text prompt after capture: {state}")

    def on_low_cpu_toggle(self):
        try:
            self.low_cpu_background = bool(self.low_cpu_var.get())
        except Exception:
            self.low_cpu_background = True

    # ---- keyboard handling ----

    def on_press(self, key):
        try:
            if len(self.current_keys) > 10:
                self.current_keys.clear()
            if hasattr(key, 'vk'):
                self.current_keys.add(key.vk)
            elif hasattr(key, 'char'):
                self.current_keys.add(key.char.lower() if key.char else None)
            else:
                self.current_keys.add(str(key))

            has_shift = 160 in self.current_keys or 161 in self.current_keys

            # Restore hotkey: Ctrl+Shift+R
            has_ctrl = 162 in self.current_keys or 163 in self.current_keys
            has_r = 'r' in self.current_keys or 82 in self.current_keys
            if has_ctrl and has_shift and has_r:
                self.restore_from_background()
                return

            if not self.is_monitoring:
                return

            # Win+Shift+S
            has_win = 91 in self.current_keys or 92 in self.current_keys
            has_s = 's' in self.current_keys or 83 in self.current_keys
            if has_win and has_shift and has_s:
                self.on_activate()
        except Exception:
            pass

    def on_release(self, key):
        try:
            if hasattr(key, 'vk'):
                self.current_keys.discard(key.vk)
            elif hasattr(key, 'char'):
                self.current_keys.discard(key.char.lower() if key.char else None)
            else:
                self.current_keys.discard(str(key))
        except Exception:
            pass

    def on_activate(self):
        if self.is_monitoring:
            self.log("\U0001f4f8 Win+Shift+S detected \u2014 clipboard monitor will capture the snip.")

    # ---- window management ----

    def hide_to_background(self):
        try:
            self.minimized_to_background = True
            self.root.iconify()
            self.log("\U0001f6cc App minimized to background. Click taskbar or press Ctrl+Shift+R to restore.")
        except Exception as e:
            self.log(f"\u2717 Error hiding to background: {e}")

    def restore_from_background(self):
        if threading.current_thread() is not threading.main_thread():
            try:
                self._ui_queue.put(self._restore_from_background_ui)
            except Exception:
                pass
            return
        self._restore_from_background_ui()

    def _restore_from_background_ui(self):
        try:
            self._dismiss_active_annotation()
            self.minimized_to_background = False
            self.window_is_iconic = False
            self.root.deiconify()
            self.root.update_idletasks()
            self.root.state('normal')
            self.root.update_idletasks()
            self.root.lift()
            self.root.attributes('-topmost', True)
            self.root.focus_force()
            self.root.after(300, lambda: self.root.attributes('-topmost', False))
            self.log("\U0001f53c App restored from background.")
        except Exception as e:
            self.log(f"\u2717 Error restoring from background: {e}")

    def _on_window_state_change(self, event=None):
        if event and event.widget != self.root:
            return
        try:
            state = str(self.root.state())
        except Exception:
            return
        was_iconic = self.window_is_iconic
        self.window_is_iconic = (state == 'iconic')
        if was_iconic and not self.window_is_iconic:
            self.minimized_to_background = False
        if self.window_is_iconic != was_iconic:
            if self.window_is_iconic or self.minimized_to_background:
                self.log("\u26a0\ufe0f Window minimized/hidden \u2014 using low CPU background polling.")
            else:
                self.log("\u26a1 Window active \u2014 using normal polling.")

    def _on_window_activated(self, event=None):
        if event and event.widget != self.root:
            return
        if self.minimized_to_background:
            self.minimized_to_background = False
            self.window_is_iconic = False

    def _dismiss_active_annotation(self):
        if self._active_annotation_dialog is not None:
            try:
                self._active_annotation_dialog.destroy()
            except Exception:
                pass
            self._active_annotation_dialog = None
            self._active_annotation = None

    # ---- clipboard monitoring ----

    def start_clipboard_monitor(self):
        if self.clipboard_thread and self.clipboard_thread.is_alive():
            return
        self.clipboard_stop_event.clear()
        self.clipboard_thread = threading.Thread(target=self._clipboard_monitor_loop, daemon=True)
        self.clipboard_thread.start()
        self.log("\u26a1 Clipboard monitor started (captures Win+Shift+S/snips).")

    def stop_clipboard_monitor(self):
        if self.clipboard_thread:
            self.clipboard_stop_event.set()
            try:
                self.clipboard_thread.join(timeout=1)
            except Exception:
                pass
            self.clipboard_thread = None
            self.log("\U0001f6d1 Clipboard monitor stopped.")

    def _clipboard_monitor_loop(self):
        while not self.clipboard_stop_event.is_set():
            try:
                if self.low_cpu_background and (self.minimized_to_background or self.window_is_iconic):
                    timeout = self.clipboard_poll_interval_bg
                else:
                    timeout = self.clipboard_poll_interval_active

                if self.clipboard_stop_event.wait(timeout):
                    break

                img = ImageGrab.grabclipboard()
                if img is not None:
                    buffered = BytesIO()
                    img.save(buffered, format='PNG')
                    img_hash = hashlib.md5(buffered.getvalue()).hexdigest()
                    if img_hash != self.last_clipboard_hash:
                        self.last_clipboard_hash = img_hash
                        self._pending_clipboard_image = img
                        self.log("\U0001f4cb New clipboard image detected \u2014 saving...")
                        try:
                            self._ui_queue.put(self._on_clipboard_image)
                        except Exception:
                            pass
            except Exception:
                self.clipboard_stop_event.wait(1)

    def _on_clipboard_image(self, event=None):
        img = self._pending_clipboard_image
        if img is None:
            return
        try:
            self._pending_clipboard_image = None
            ask = bool(self.ask_annotation_var.get()) and (not self.minimized_to_background)
            self.capture_screenshot(image=img, ask_annotation=ask)
        except Exception as e:
            self.log(f"\u2717 Error saving clipboard image: {e}")

    # ---- screenshot capture ----

    def capture_screenshot(self, image=None, ask_annotation=True):
        try:
            screenshot = image if image is not None else ImageGrab.grab()
            annotation = ''
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            filename = f"screenshot_{timestamp}.png"
            filepath = os.path.join(self.screenshots_folder, filename)

            screenshot.save(filepath)

            screenshot_data = {
                'filename': filename,
                'filepath': filepath,
                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                'annotation': annotation,
                'base64': None,
            }
            with self.data_lock:
                self.screenshots_data.append(screenshot_data)
            self.html_dirty = True

            def _update_ui():
                self.screenshot_count += 1
                self.counter_label.config(text=f"Screenshots captured: {self.screenshot_count}")
                self.log(f"\u2713 Screenshot saved: {filename}")
                self.log(f"  \U0001f4cd Saved to: {filepath}")
                self.log("  \u23f3 Processing image for HTML in background...")
                self.log("  \U0001f4c4 HTML will be saved when you Stop Monitoring or close the app.")

            try:
                self.root.after(0, _update_ui)
            except Exception:
                _update_ui()

            with self._processing_cond:
                self._inflight_processing += 1
            threading.Thread(target=self._process_image_for_html, args=(filepath, filename), daemon=True).start()

            if ask_annotation and not self.minimized_to_background:
                try:
                    self.root.after(0, lambda: self.show_annotation_prompt(filename))
                except Exception:
                    self.show_annotation_prompt(filename)
        except Exception as e:
            self.log(f"\u2717 Error: {e}")
            try:
                messagebox.showerror("Capture Error", str(e))
            except Exception:
                pass

    def _process_image_for_html(self, filepath, filename):
        try:
            with open(filepath, 'rb') as f:
                img_bytes = f.read()
            img_base64 = base64.b64encode(img_bytes).decode()
            with self.data_lock:
                for s in self.screenshots_data:
                    if s['filename'] == filename:
                        s['base64'] = img_base64
                        break
            self.html_dirty = True
        except Exception as e:
            self.log(f"\u2717 Error processing image for HTML: {e}")
        finally:
            with self._processing_cond:
                self._inflight_processing = max(0, self._inflight_processing - 1)
                self._processing_cond.notify_all()

        try:
            if self.html_dirty:
                self.generate_html()
                self.html_dirty = False
                self.log(f"\U0001f4be HTML auto-saved after processing: {filename}")
        except Exception as e:
            self.log(f"\u2717 Auto-save HTML failed: {e}")

    # ---- annotation ----

    def show_annotation_prompt(self, filename):
        if self._active_annotation is not None:
            self._pending_annotations.append(filename)
            return
        self._active_annotation = filename
        try:
            self._create_annotation_dialog(filename)
        except Exception as e:
            self.log(f"\u26a0\ufe0f Annotation dialog failed: {e}")
            self._active_annotation = None
            self._show_next_queued_annotation()

    def _show_next_queued_annotation(self):
        if self._pending_annotations and self._active_annotation is None:
            next_file = self._pending_annotations.pop(0)
            self.root.after(50, lambda: self.show_annotation_prompt(next_file))

    def _create_annotation_dialog(self, filename):
        dialog = tk.Toplevel(self.root)
        self._active_annotation_dialog = dialog
        dialog.title("Add Annotation")
        dialog.geometry("520x220")
        dialog.configure(bg=self._CLR_BG)
        dialog.transient(self.root)

        dialog.update_idletasks()
        x = (dialog.winfo_screenwidth() // 2) - (dialog.winfo_width() // 2)
        y = (dialog.winfo_screenheight() // 2) - (dialog.winfo_height() // 2)
        dialog.geometry(f"+{x}+{y}")

        tk.Label(dialog, text="Add text for this screenshot (optional):",
                 font=("Arial", 11), fg=self._CLR_TEXT, bg=self._CLR_BG).pack(pady=12)

        text_entry = tk.Entry(dialog, width=62, font=("Arial", 10),
                              bg="#ffffff", fg=self._CLR_TEXT, insertbackground=self._CLR_TEXT)
        text_entry.pack(pady=6, padx=14)
        text_entry.focus()

        tk.Label(dialog, text="Tip: Press Enter to save, Esc to skip.",
                 fg=self._CLR_MUTED, bg=self._CLR_BG).pack(pady=4)

        def _close_guard():
            try:
                self._active_annotation = None
                self._active_annotation_dialog = None
            except Exception:
                pass
            try:
                self.root.after(50, self._show_next_queued_annotation)
            except Exception:
                pass

        def submit():
            text = text_entry.get().strip()
            try:
                self.update_annotation(filename, text)
            finally:
                _close_guard()
                dialog.destroy()

        def skip():
            try:
                self.update_annotation(filename, "")
            finally:
                _close_guard()
                dialog.destroy()

        btn_frame = tk.Frame(dialog, bg=self._CLR_BG)
        btn_frame.pack(pady=14)
        tk.Button(btn_frame, text="\u2713 Save", command=submit,
                  bg=self._CLR_GREEN, fg="white", activebackground=self._CLR_GREEN_H,
                  activeforeground="white", font=("Arial", 10, "bold"),
                  relief="raised", bd=2, padx=18, pady=8, cursor="hand2").pack(side='left', padx=6)
        tk.Button(btn_frame, text="Skip", command=skip,
                  bg=self._CLR_SURFACE, fg=self._CLR_TEXT, activebackground=self._CLR_INFO_BG,
                  font=("Arial", 10), relief="raised", bd=2, padx=18, pady=8,
                  cursor="hand2").pack(side='left', padx=6)

        text_entry.bind('<Return>', lambda e: submit())
        text_entry.bind('<Escape>', lambda e: skip())
        dialog.protocol("WM_DELETE_WINDOW", skip)

    def update_annotation(self, filename, text):
        with self.data_lock:
            updated = False
            for s in self.screenshots_data:
                if s['filename'] == filename:
                    s['annotation'] = text
                    updated = True
                    break
        if updated:
            self.html_dirty = True
            try:
                self.root.after(0, lambda: self.log(f"\U0001f4dd Annotation {'added' if text else 'skipped'}: {text}"))
            except Exception:
                self.log(f"\U0001f4dd Annotation {'added' if text else 'skipped'}")
        else:
            self.log(f"\u26a0\ufe0f Could not find screenshot to annotate: {filename}")

    # ---- HTML generation ----

    def generate_html(self):
        with self._html_write_lock:
            self._generate_html_inner()

    def _generate_html_inner(self):
        with self.data_lock:
            screenshots_copy = list(self.screenshots_data)

        total_captured = len(screenshots_copy)
        processed_count = len([s for s in screenshots_copy if s.get('base64')])

        html_content = f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Screenshots Document - {datetime.now().strftime("%Y-%m-%d")}</title>
    <style>
        body {{ font-family: Arial, sans-serif; max-width: 900px; margin: 0 auto; padding: 20px; background-color: #f5f5f5; }}
        h1 {{ color: #2c3e50; text-align: center; border-bottom: 3px solid #3498db; padding-bottom: 10px; }}
        .screenshot-container {{ background: white; margin: 20px 0; padding: 20px; border-radius: 8px; box-shadow: 0 2px 5px rgba(0,0,0,0.1); }}
        .timestamp {{ color: #7f8c8d; font-size: 14px; margin-bottom: 10px; }}
        .annotation {{ background-color: #fff3cd; border-left: 4px solid #ffc107; padding: 10px 15px; margin: 10px 0; font-style: italic; color: #856404; }}
        .screenshot {{ max-width: 100%; height: auto; border: 1px solid #ddd; border-radius: 4px; margin-top: 10px; }}
        .count {{ text-align: center; color: #7f8c8d; margin: 20px 0; }}
        .header-info {{ background: #ecf0f1; padding: 15px; border-radius: 5px; text-align: center; margin-bottom: 20px; }}
    </style>
</head>
<body>
    <h1>\U0001f4f8 Screenshots Document</h1>
    <div class="header-info">
        <strong>Generated:</strong> {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}<br>
        <strong>Total Captured:</strong> {total_captured} &nbsp; <strong>Processed:</strong> {processed_count}
    </div>
"""
        processed_list = [s for s in screenshots_copy if s.get('base64')]
        for i, screenshot in enumerate(processed_list, 1):
            html_content += f"""
    <div class="screenshot-container">
        <div class="timestamp"><strong>Screenshot #{i}</strong> - {screenshot['timestamp']}</div>
"""
            if screenshot['annotation']:
                safe_ann = _html_mod.escape(screenshot['annotation'])
                html_content += f"""        <div class="annotation">\U0001f4dd {safe_ann}</div>\n"""
            html_content += f"""        <img src="data:image/png;base64,{screenshot['base64']}" alt="Screenshot {i}" class="screenshot">
    </div>
"""
        html_content += "</body>\n</html>\n"

        out_path = os.path.abspath(self.html_file)
        os.makedirs(os.path.dirname(out_path), exist_ok=True)
        with open(out_path, "w", encoding="utf-8") as f:
            f.write(html_content)
        self.latest_html_file = out_path

    def flush_html(self, reason="manual", timeout=5.0):
        if not self.html_dirty:
            return False
        with self._processing_cond:
            end = time.time() + float(timeout)
            while self._inflight_processing > 0 and time.time() < end:
                remaining = end - time.time()
                self._processing_cond.wait(timeout=max(0.05, min(0.5, remaining)))
        try:
            self.generate_html()
            self.html_dirty = False
            try:
                self.root.after(0, lambda: self.log(f"\U0001f4c4 HTML saved ({reason}) -> {os.path.abspath(self.latest_html_file)}"))
            except Exception:
                self.log(f"\U0001f4c4 HTML saved ({reason}) -> {os.path.abspath(self.latest_html_file)}")
            return True
        except Exception as e:
            self.log(f"\u2717 Error saving HTML ({reason}): {e}")
            return False

    def flush_html_async(self, reason="manual", timeout=5.0):
        threading.Thread(target=self.flush_html, args=(reason, timeout), daemon=True).start()

    # ---- open file / folder ----

    def open_folder(self):
        try:
            import subprocess
            if sys.platform == 'win32':
                os.startfile(self.screenshots_folder)
            elif sys.platform == 'darwin':
                subprocess.run(['open', self.screenshots_folder])
            else:
                subprocess.run(['xdg-open', self.screenshots_folder])
            self.log(f"\U0001f4c2 Opened folder: {self.screenshots_folder}")
        except Exception as e:
            self.log(f"\u2717 Error opening folder: {e}")

    def open_html(self):
        try:
            if self.html_dirty:
                self.flush_html(reason="Open HTML", timeout=5.0)
            target = os.path.abspath(self.latest_html_file)
            if not os.path.exists(target):
                base = os.path.abspath(self.html_file)
                if os.path.exists(base):
                    target = base
            if not os.path.exists(target):
                self.log("\u26a0\ufe0f No HTML document yet \u2014 capture some screenshots first.")
                return
            import webbrowser
            webbrowser.open(f"file://{target}")
            self.log("\U0001f4c4 Opened HTML document in browser")
        except Exception as e:
            self.log(f"\u2717 Error opening HTML: {e}")

    # ---- exit / crash safety ----

    def _register_exit_handlers(self):
        atexit.register(self._emergency_flush_html)
        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                signal.signal(sig, self._signal_handler)
            except (OSError, ValueError):
                pass

    def _signal_handler(self, signum, frame):
        self._emergency_flush_html()
        sys.exit(0)

    def _emergency_flush_html(self):
        if self._emergency_saved:
            return
        self._emergency_saved = True
        try:
            if self.html_dirty or self.screenshots_data:
                self.generate_html()
                self.html_dirty = False
        except Exception:
            pass

    def _periodic_autosave(self):
        try:
            if self.html_dirty:
                self.flush_html_async(reason="periodic auto-save", timeout=3.0)
        except Exception:
            pass
        try:
            self.root.after(30000, self._periodic_autosave)
        except Exception:
            pass

    # ---- run / close ----

    def run(self):
        self.root.protocol("WM_DELETE_WINDOW", self.on_close)
        self.root.after(30000, self._periodic_autosave)
        self.root.mainloop()

    def on_close(self):
        try:
            self.stop_monitoring()
        except Exception:
            pass
        try:
            if self.listener:
                self.listener.stop()
        except Exception:
            pass
        try:
            if self.html_dirty:
                self.flush_html(reason="App Close", timeout=5.0)
        except Exception:
            pass
        if self.screenshot_count > 0:
            self.log(f"\u2713 Session complete: {self.screenshot_count} screenshots captured")
        self.root.destroy()

    def _start_listener(self):
        if self.listener is None or not self.listener.running:
            self.listener = keyboard.Listener(
                on_press=self.on_press,
                on_release=self.on_release)
            self.listener.start()


def main():
    """Entry point for the console script."""
    _ensure_imagegrab()
    _ensure_pynput()
    print("Screenshot Tool - Starting...")
    app = SimpleScreenshotTool()
    app.run()


if __name__ == "__main__":
    main()
