whisper-local/whisper_local/tray/_download_progress.py

"""Download-Fortschrittsdialog für den ersten Whisper-Modell-Download (Windows)."""

from __future__ import annotations

import queue
import sys
import threading
from typing import Any

import tqdm as tqdm_module


class TkProgressTqdm(tqdm_module.tqdm):
    """tqdm-Ersatz, der Fortschritts-Updates thread-safe in eine Queue schreibt."""

    _queue: queue.Queue | None = None

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self._desc = kwargs.get("desc", "")
        self._accumulated_n = 0
        super().__init__(*args, **kwargs)

    def update(self, n: int | float | None = 1) -> bool | None:
        if n is None:
            n = 0
        self._accumulated_n += n
        result = super().update(n)
        if self._queue is not None:
            self._queue.put({
                "file": self._desc,
                "n": self._accumulated_n,
                "total": self.total or 0,
            })
        return result


def load_model_with_progress(
    model_name: str,
    compute_type: str,
    download_root: str | None,
) -> Any:
    """Lädt WhisperModel — zeigt bei Bedarf einen Download-Fortschrittsdialog.

    Wenn das Modell bereits gecacht ist (kein tqdm-Update kommt), erscheint
    kein Dialog. Auf Download-Fehler wird ein Fehlerdialog gezeigt und sys.exit(1)
    aufgerufen.
    """
    import tkinter as tk
    from tkinter import messagebox, ttk

    from faster_whisper import WhisperModel

    from whisper_local.tray._theme import apply_system_theme

    q: queue.Queue[dict[str, Any] | None] = queue.Queue()
    result: list[WhisperModel | None] = [None]
    error: list[BaseException | None] = [None]
    original_tqdm = tqdm_module.tqdm

    def worker() -> None:
        TkProgressTqdm._queue = q
        tqdm_module.tqdm = TkProgressTqdm
        try:
            result[0] = WhisperModel(
                model_name,
                compute_type=compute_type,
                download_root=download_root,
            )
        except Exception as exc:
            error[0] = exc
        finally:
            tqdm_module.tqdm = original_tqdm
            TkProgressTqdm._queue = None
            q.put(None)  # Sentinel: signalisiert Fertigstellung

    thread = threading.Thread(target=worker, daemon=True)
    thread.start()

    # --- tkinter-Dialog (lazy: erscheint nur bei echtem Download) ---
    root = tk.Tk()
    root.withdraw()  # zunächst versteckt
    root.title("whisper-local – Modell wird geladen")
    root.resizable(False, False)
    apply_system_theme(root)

    frame = ttk.Frame(root, padding=16)
    frame.pack(fill=tk.BOTH, expand=True)

    ttk.Label(frame, text=f"Lade Whisper-Modell '{model_name}'...").pack(anchor=tk.W)

    file_var = tk.StringVar(value="")
    ttk.Label(frame, textvariable=file_var, foreground="gray").pack(anchor=tk.W, pady=(4, 8))

    progress_var = tk.DoubleVar(value=0.0)
    ttk.Progressbar(
        frame, variable=progress_var, maximum=100, length=320, mode="determinate"
    ).pack(fill=tk.X)

    pct_var = tk.StringVar(value="0 %")
    ttk.Label(frame, textvariable=pct_var).pack(anchor=tk.E, pady=(2, 0))

    _dialog_shown = False

    def poll() -> None:
        nonlocal _dialog_shown
        try:
            while True:
                msg = q.get_nowait()
                if msg is None:  # Sentinel → fertig
                    root.quit()
                    return
                # Erste echte Meldung → Dialog anzeigen
                if not _dialog_shown:
                    root.deiconify()
                    _dialog_shown = True
                # UI aktualisieren
                file_var.set(msg["file"])
                if msg["total"] > 0:
                    pct = 100.0 * msg["n"] / msg["total"]
                    progress_var.set(pct)
                    pct_var.set(f"{pct:.0f} %")
        except queue.Empty:
            pass
        root.after(50, poll)

    root.after(50, poll)
    root.mainloop()

    if error[0] is not None:
        root.withdraw()
        messagebox.showerror("Fehler beim Modell-Download", str(error[0]))
        root.destroy()
        sys.exit(1)

    root.destroy()
    return result[0]