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 längerem Laden einen Wartebalken.

    Wenn das Modell in unter 500 ms bereit ist (vollständiger Cache), erscheint
    kein Dialog. Bei Download oder langsamer Initialisierung wird ein indeterminater
    Wartebalken angezeigt. Auf 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

    result: list[Any] = [None]
    error: list[BaseException | None] = [None]
    done_event = threading.Event()

    def worker() -> None:
        try:
            result[0] = WhisperModel(
                model_name,
                compute_type=compute_type,
                download_root=download_root,
            )
        except Exception as exc:
            error[0] = exc
        finally:
            done_event.set()

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

    # Kurz warten – wenn Modell sofort bereit (vollständiger Cache), kein Dialog
    if done_event.wait(timeout=0.5):
        if error[0] is not None:
            root = tk.Tk()
            root.withdraw()
            messagebox.showerror("Fehler beim Modell-Laden", str(error[0]))
            root.destroy()
            sys.exit(1)
        return result[0]

    # Modell braucht länger (Download oder Initialisierung) → Dialog anzeigen
    root = tk.Tk()
    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)
    ttk.Label(frame, text="Bitte warten…", foreground="gray").pack(
        anchor=tk.W, pady=(4, 8)
    )

    pb = ttk.Progressbar(frame, length=320, mode="indeterminate")
    pb.pack(fill=tk.X)
    pb.start(10)

    def poll() -> None:
        if done_event.is_set():
            root.quit()
            return
        root.after(100, poll)

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

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

    root.destroy()
    return result[0]