Trixy7/trixy_core/utils/download.py

# -*- coding: utf-8 -*-
"""
Download-Utilities fuer Modelle und Ressourcen.

Unterstuetzt Downloads von:
- Hugging Face Hub
- GitHub Releases
- Direkte URLs
- Mit Fortschrittsanzeige und typisierten EventData-Events

Events (wenn EventManager verfuegbar):
- before_download: Vor dem Download (cancellable, URL/Pfad aenderbar)
- download_progress: Fortschritts-Update (alle 5%)
- download_completed: Download erfolgreich
- download_failed: Download fehlgeschlagen
- before_extract: Vor dem Entpacken (cancellable, Ziel aenderbar)
- extract_completed: Entpacken abgeschlossen
"""

import asyncio
import hashlib
import shutil
import time
from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import urlparse

from trixy_core.utils.debug import pinfo, pdebug, perror, pprogress

if TYPE_CHECKING:
    from trixy_core.events.eventmanager import EventManager


class DownloadProgressTracker:
    """Fortschritts-Tracker fuer Downloads mit visueller Progressbar."""

    def __init__(
        self,
        total_size: int = 0,
        event_manager: "EventManager | None" = None,
        download_id: str = "",
        filename: str = "",
    ):
        self.total_size = total_size
        self.downloaded = 0
        self.event_manager = event_manager
        self.download_id = download_id
        self.filename = filename
        self._last_percent = -1
        self._start_time = time.time()

    def update(self, chunk_size: int) -> None:
        """Aktualisiert den Fortschritt mit visueller Progressbar."""
        self.downloaded += chunk_size

        if self.total_size > 0:
            percent = (self.downloaded / self.total_size) * 100

            # Geschwindigkeit + ETA berechnen
            elapsed = time.time() - self._start_time
            speed = self.downloaded / elapsed if elapsed > 0 else 0
            remaining = (self.total_size - self.downloaded) / speed if speed > 0 else 0

            suffix = (
                f"{format_size(self.downloaded)}/{format_size(self.total_size)} | "
                f"{format_size(int(speed))}/s | {format_eta(int(remaining))}"
            )

            # Visuelle Progressbar anzeigen
            pprogress(self.filename, self.downloaded, self.total_size, suffix)

            # Event emittieren (nur alle 5%)
            current_percent = int(percent / 5) * 5
            if self.event_manager and current_percent > self._last_percent:
                self._last_percent = current_percent
                asyncio.create_task(self._emit_progress(percent, int(speed), int(remaining)))

    async def _emit_progress(self, percent: float, speed: int, eta_seconds: int) -> None:
        """Emittiert Progress-Event ueber trigger()."""
        if self.event_manager:
            from trixy_core.events.event_data.download import DownloadProgress as DownloadProgressEvent
            event = DownloadProgressEvent(
                download_id=self.download_id,
                filename=self.filename,
                downloaded=self.downloaded,
                total=self.total_size,
                percent=percent,
                speed=speed,
                eta_seconds=eta_seconds,
            )
            await self.event_manager.trigger("download_progress", event)

    @property
    def percent(self) -> float:
        """Fortschritt in Prozent."""
        if self.total_size == 0:
            return 0.0
        return (self.downloaded / self.total_size) * 100


# Rueckwaertskompatibilitaet: Alias fuer die umbenannte Klasse
DownloadProgress = DownloadProgressTracker


async def download_file(
    url: str,
    dest_path: Path,
    chunk_size: int = 65536,
    timeout: int = 3600,
    event_manager: "EventManager | None" = None,
    download_id: str | None = None,
    display_name: str | None = None,
) -> bool:
    """
    Laedt eine Datei von einer URL herunter.

    Args:
        url: Download-URL
        dest_path: Ziel-Pfad
        chunk_size: Chunk-Groesse in Bytes
        timeout: Timeout in Sekunden
        event_manager: Optional EventManager fuer Events
        download_id: Eindeutige Download-ID fuer Events
        display_name: Anzeigename fuer Progressbar (statt Dateiname)

    Returns:
        True bei Erfolg
    """
    filename = display_name or dest_path.name
    dl_id = download_id or hashlib.md5(url.encode()).hexdigest()[:8]

    # Event: Vor dem Download (cancellable)
    if event_manager:
        from trixy_core.events.event_data.download import BeforeDownload
        before = BeforeDownload(
            download_id=dl_id,
            url=url,
            filename=filename,
            dest_path=str(dest_path),
        )
        before = await event_manager.trigger("before_download", before)
        if before.is_cancelled():
            pinfo(f"Download abgebrochen durch Plugin: {filename}")
            return False
        # Uebernahme von Aenderungen (Redirect, anderer Pfad)
        url = before.url
        dest_path = Path(before.dest_path)
        filename = dest_path.name

    temp_path = None
    start_time = time.time()
    total_size = 0

    try:
        import aiohttp

        dest_path.parent.mkdir(parents=True, exist_ok=True)
        temp_path = dest_path.with_suffix(dest_path.suffix + ".tmp")

        # Timeout nur fuer Socket-Reads, nicht gesamt — grosse Downloads brechen sonst ab
        client_timeout = aiohttp.ClientTimeout(
            total=None,
            sock_read=60,
            sock_connect=30,
        )

        async with aiohttp.ClientSession() as session:
            async with session.get(url, timeout=client_timeout) as response:
                if response.status != 200:
                    error_msg = f"HTTP {response.status}"
                    perror(f"Download fehlgeschlagen: {error_msg}")

                    if event_manager:
                        from trixy_core.events.event_data.download import DownloadFailed
                        failed = DownloadFailed(
                            download_id=dl_id,
                            url=url,
                            filename=filename,
                            error=error_msg,
                        )
                        await event_manager.trigger("download_failed", failed)
                    return False

                total_size = int(response.headers.get("content-length", 0))
                progress = DownloadProgressTracker(
                    total_size,
                    event_manager,
                    dl_id,
                    filename,
                )

                with open(temp_path, "wb") as f:
                    async for chunk in response.content.iter_chunked(chunk_size):
                        f.write(chunk)
                        progress.update(len(chunk))

        # Temp-Datei umbenennen
        temp_path.rename(dest_path)

        # Event: Download abgeschlossen
        elapsed = time.time() - start_time
        if event_manager:
            from trixy_core.events.event_data.download import DownloadCompleted
            completed = DownloadCompleted(
                download_id=dl_id,
                url=url,
                filename=filename,
                dest_path=str(dest_path),
                size=total_size,
                duration_seconds=elapsed,
                speed_avg=int(total_size / elapsed) if elapsed > 0 else 0,
            )
            await event_manager.trigger("download_completed", completed)

        return True

    except ImportError:
        # Fallback ohne aiohttp
        return await _download_file_urllib(
            url, dest_path, chunk_size,
            event_manager, dl_id
        )
    except (KeyboardInterrupt, asyncio.CancelledError):
        # Ctrl+C oder Task-Abbruch — sauber aufraumen
        print()  # Newline nach Progressbar
        perror(f"Download abgebrochen: {filename}")

        if event_manager:
            from trixy_core.events.event_data.download import DownloadFailed
            failed = DownloadFailed(
                download_id=dl_id,
                url=url,
                filename=filename,
                error="Download abgebrochen (Keyboard Interrupt)",
            )
            await event_manager.trigger("download_failed", failed)

        try:
            if temp_path and temp_path.exists():
                temp_path.unlink()
        except:
            pass

        return False

    except Exception as e:
        error_msg = str(e)
        perror(f"Download-Fehler: {error_msg}")

        if event_manager:
            from trixy_core.events.event_data.download import DownloadFailed
            failed = DownloadFailed(
                download_id=dl_id,
                url=url,
                filename=filename,
                error=error_msg,
            )
            await event_manager.trigger("download_failed", failed)

        try:
            if temp_path and temp_path.exists():
                temp_path.unlink()
        except:
            pass

        return False


async def _download_file_urllib(
    url: str,
    dest_path: Path,
    chunk_size: int = 8192,
    event_manager: "EventManager | None" = None,
    download_id: str = "",
) -> bool:
    """Fallback-Download mit urllib."""
    import urllib.request

    filename = dest_path.name
    temp_path = None
    start_time = time.time()

    try:
        dest_path.parent.mkdir(parents=True, exist_ok=True)
        temp_path = dest_path.with_suffix(dest_path.suffix + ".tmp")

        loop = asyncio.get_event_loop()

        def do_download():
            with urllib.request.urlopen(url) as response:
                total_size = int(response.headers.get("content-length", 0))
                progress = DownloadProgressTracker(
                    total_size,
                    None,  # Kein async in sync context
                    download_id,
                    filename,
                )

                with open(temp_path, "wb") as f:
                    while True:
                        chunk = response.read(chunk_size)
                        if not chunk:
                            break
                        f.write(chunk)
                        progress.update(len(chunk))

                return total_size

        total_size = await loop.run_in_executor(None, do_download)
        temp_path.rename(dest_path)

        # Event: Download abgeschlossen
        elapsed = time.time() - start_time
        if event_manager:
            from trixy_core.events.event_data.download import DownloadCompleted
            completed = DownloadCompleted(
                download_id=download_id,
                url=url,
                filename=filename,
                dest_path=str(dest_path),
                size=total_size,
                duration_seconds=elapsed,
                speed_avg=int(total_size / elapsed) if elapsed > 0 else 0,
            )
            await event_manager.trigger("download_completed", completed)

        return True

    except (KeyboardInterrupt, asyncio.CancelledError):
        print()
        perror(f"Download abgebrochen: {filename}")

        if event_manager:
            from trixy_core.events.event_data.download import DownloadFailed
            failed = DownloadFailed(
                download_id=download_id,
                url=url,
                filename=filename,
                error="Download abgebrochen (Keyboard Interrupt)",
            )
            await event_manager.trigger("download_failed", failed)

        if temp_path and temp_path.exists():
            temp_path.unlink()
        return False

    except Exception as e:
        error_msg = str(e)
        perror(f"Download-Fehler (urllib): {error_msg}")

        if event_manager:
            from trixy_core.events.event_data.download import DownloadFailed
            failed = DownloadFailed(
                download_id=download_id,
                url=url,
                filename=filename,
                error=error_msg,
            )
            await event_manager.trigger("download_failed", failed)

        if temp_path and temp_path.exists():
            temp_path.unlink()
        return False


async def download_from_huggingface(
    repo_id: str,
    filename: str,
    dest_path: Path,
    revision: str = "main",
    event_manager: "EventManager | None" = None,
    download_id: str | None = None,
) -> bool:
    """
    Laedt eine Datei von Hugging Face Hub herunter.

    Args:
        repo_id: Repository-ID (z.B. "rhasspy/piper-voices")
        filename: Dateiname im Repository
        dest_path: Ziel-Pfad
        revision: Branch/Tag/Commit
        event_manager: Optional EventManager fuer Events
        download_id: Eindeutige Download-ID

    Returns:
        True bei Erfolg
    """
    dl_id = download_id or hashlib.md5(f"{repo_id}/{filename}".encode()).hexdigest()[:8]

    try:
        from huggingface_hub import hf_hub_download

        dest_path.parent.mkdir(parents=True, exist_ok=True)

        loop = asyncio.get_event_loop()

        def do_download():
            # huggingface_hub laedt in Cache, wir kopieren dann
            cached_path = hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                revision=revision,
                local_dir=None,  # Nutzt Cache
            )
            # Kopieren zum Ziel
            shutil.copy2(cached_path, dest_path)
            return Path(cached_path).stat().st_size

        pinfo(f"Lade von Hugging Face: {repo_id}/{filename}")
        start_time = time.time()
        file_size = await loop.run_in_executor(None, do_download)
        elapsed = time.time() - start_time

        pinfo(f"Download abgeschlossen: {dest_path.name}")

        # Event: Download abgeschlossen
        if event_manager:
            from trixy_core.events.event_data.download import DownloadCompleted
            completed = DownloadCompleted(
                download_id=dl_id,
                url=f"hf://{repo_id}/{filename}",
                filename=dest_path.name,
                dest_path=str(dest_path),
                size=file_size,
                duration_seconds=elapsed,
                speed_avg=int(file_size / elapsed) if elapsed > 0 else 0,
            )
            await event_manager.trigger("download_completed", completed)

        return True

    except ImportError:
        # Fallback: Direkte URL
        url = f"https://huggingface.co/{repo_id}/resolve/{revision}/{filename}"
        pinfo(f"Hugging Face Hub nicht installiert, nutze direkte URL")
        return await download_file(
            url, dest_path,
            event_manager=event_manager, download_id=dl_id
        )

    except Exception as e:
        error_msg = str(e)
        perror(f"Hugging Face Download-Fehler: {error_msg}")

        if event_manager:
            from trixy_core.events.event_data.download import DownloadFailed
            failed = DownloadFailed(
                download_id=dl_id,
                url=f"hf://{repo_id}/{filename}",
                filename=dest_path.name,
                error=error_msg,
            )
            await event_manager.trigger("download_failed", failed)

        return False


async def download_from_github(
    repo: str,
    tag: str,
    asset_name: str,
    dest_path: Path,
    event_manager: "EventManager | None" = None,
    download_id: str | None = None,
) -> bool:
    """
    Laedt ein Release-Asset von GitHub herunter.

    Args:
        repo: Repository (z.B. "mozilla/DeepSpeech")
        tag: Release-Tag (z.B. "v0.9.3")
        asset_name: Asset-Dateiname
        dest_path: Ziel-Pfad
        event_manager: Optional EventManager fuer Events
        download_id: Eindeutige Download-ID

    Returns:
        True bei Erfolg
    """
    url = f"https://github.com/{repo}/releases/download/{tag}/{asset_name}"
    dl_id = download_id or hashlib.md5(url.encode()).hexdigest()[:8]

    pinfo(f"Lade von GitHub: {repo} @ {tag} / {asset_name}")

    return await download_file(
        url, dest_path,
        event_manager=event_manager, download_id=dl_id
    )


async def download_and_extract(
    url: str,
    dest_dir: Path,
    event_manager: "EventManager | None" = None,
    download_id: str | None = None,
) -> bool:
    """
    Laedt ein Archiv herunter und extrahiert es.

    Unterstuetzt: .zip, .tar.gz, .tgz

    Args:
        url: Download-URL
        dest_dir: Ziel-Verzeichnis
        event_manager: Optional EventManager fuer Events
        download_id: Eindeutige Download-ID

    Returns:
        True bei Erfolg
    """
    import tempfile

    # Dateiname aus URL
    parsed = urlparse(url)
    filename = Path(parsed.path).name
    dl_id = download_id or hashlib.md5(url.encode()).hexdigest()[:8]

    # Temporaere Datei — suffix mit _ damit Temp-Prefix nicht im Dateinamen erscheint
    with tempfile.NamedTemporaryFile(suffix="_" + filename, delete=False) as tmp:
        temp_path = Path(tmp.name)

    try:
        # Download — display_name zeigt echten Dateinamen statt Temp-Pfad
        if not await download_file(
            url, temp_path,
            event_manager=event_manager, download_id=dl_id,
            display_name=filename,
        ):
            return False

        # Event: Vor dem Entpacken (cancellable)
        if event_manager:
            from trixy_core.events.event_data.download import BeforeExtract
            before = BeforeExtract(
                download_id=dl_id,
                filename=filename,
                archive_path=str(temp_path),
                dest_dir=str(dest_dir),
            )
            before = await event_manager.trigger("before_extract", before)
            if before.is_cancelled():
                pinfo(f"Entpacken abgebrochen/uebernommen durch Plugin: {filename}")
                return True  # Plugin hat sich drum gekuemmert
            dest_dir = Path(before.dest_dir)  # Plugin kann Ziel aendern

        # Extrahieren
        dest_dir.mkdir(parents=True, exist_ok=True)

        loop = asyncio.get_event_loop()

        def do_extract():
            if filename.endswith(".zip"):
                import zipfile
                with zipfile.ZipFile(temp_path, "r") as zf:
                    zf.extractall(dest_dir)
                    return len(zf.namelist())
            elif filename.endswith((".tar.gz", ".tgz")):
                import tarfile
                with tarfile.open(temp_path, "r:gz") as tf:
                    tf.extractall(dest_dir)
                    return len(tf.getnames())
            elif filename.endswith(".tar"):
                import tarfile
                with tarfile.open(temp_path, "r") as tf:
                    tf.extractall(dest_dir)
                    return len(tf.getnames())
            else:
                raise ValueError(f"Unbekanntes Archiv-Format: {filename}")

        file_count = await loop.run_in_executor(None, do_extract)
        pinfo(f"Extrahiert nach: {dest_dir} ({file_count} Dateien)")

        # Event: Entpacken abgeschlossen
        if event_manager:
            from trixy_core.events.event_data.download import ExtractCompleted
            completed = ExtractCompleted(
                download_id=dl_id,
                filename=filename,
                dest_dir=str(dest_dir),
                file_count=file_count,
            )
            await event_manager.trigger("extract_completed", completed)

        return True

    except Exception as e:
        error_msg = str(e)
        perror(f"Extrahieren fehlgeschlagen: {error_msg}")

        if event_manager:
            from trixy_core.events.event_data.download import DownloadFailed
            failed = DownloadFailed(
                download_id=dl_id,
                url=url,
                filename=filename,
                error=error_msg,
            )
            await event_manager.trigger("download_failed", failed)

        return False

    finally:
        if temp_path.exists():
            temp_path.unlink()


def format_size(size_bytes: int) -> str:
    """Formatiert Bytes als lesbare Groesse."""
    for unit in ["B", "KB", "MB", "GB"]:
        if size_bytes < 1024:
            return f"{size_bytes:.1f} {unit}"
        size_bytes /= 1024
    return f"{size_bytes:.1f} TB"


def format_eta(seconds: int) -> str:
    """
    Formatiert ETA mit automatischer Einheit, max. 2-stellig.

    <= 99s → "XXs", <= 99min → "XXm", sonst → "XXh"
    """
    if seconds < 0:
        return "~"
    if seconds <= 99:
        return f"{seconds}s"
    minutes = seconds // 60
    if minutes <= 99:
        return f"{minutes}m"
    hours = min(seconds // 3600, 99)
    return f"{hours}h"