first commit

download_manager.py

@@ -0,0 +1,582 @@
+import os
+import requests
+import threading
+import time
+import tkinter as tk
+from tkinter import ttk, messagebox
+from typing import Optional, Dict, Any, List, Callable
+from dataclasses import dataclass
+from enum import Enum
+import json
+from datetime import datetime
+import queue
+
+
+class DownloadStatus(Enum):
+    """Download status enumeration."""
+    QUEUED = "Queued"
+    DOWNLOADING = "Downloading"
+    PAUSED = "Paused"
+    COMPLETED = "Completed"
+    FAILED = "Failed"
+    CANCELLED = "Cancelled"
+    AUTH_REQUIRED = "Auth Required"
+
+
+@dataclass
+class DownloadItem:
+    """Represents a download item."""
+    id: str
+    repo_id: str
+    filename: str
+    url: str
+    save_path: str
+    total_size: int = 0
+    downloaded_size: int = 0
+    status: DownloadStatus = DownloadStatus.QUEUED
+    error_message: str = ""
+    start_time: Optional[float] = None
+    end_time: Optional[float] = None
+    speed: float = 0.0
+    eta: int = 0
+    headers: Dict[str, str] = None
+    resume_position: int = 0
+    
+    def __post_init__(self):
+        if self.headers is None:
+            self.headers = {}
+    
+    @property
+    def progress(self) -> float:
+        """Calculate download progress percentage."""
+        if self.total_size > 0:
+            return (self.downloaded_size / self.total_size) * 100
+        return 0.0
+    
+    @property
+    def is_resumable(self) -> bool:
+        """Check if download can be resumed."""
+        return self.status in [DownloadStatus.PAUSED, DownloadStatus.FAILED]
+
+
+class DownloadManager:
+    """Manages multiple downloads with pause/resume support."""
+    
+    def __init__(self, max_concurrent: int = 3):
+        self.downloads: Dict[str, DownloadItem] = {}
+        self.download_queue: queue.Queue = queue.Queue()
+        self.active_downloads: Dict[str, threading.Thread] = {}
+        self.max_concurrent = max_concurrent
+        self.callbacks: Dict[str, List[Callable]] = {
+            'on_progress': [],
+            'on_status_change': [],
+            'on_complete': [],
+            'on_error': []
+        }
+        self._stop_flags: Dict[str, threading.Event] = {}
+        self._pause_flags: Dict[str, threading.Event] = {}
+        self._worker_thread = threading.Thread(target=self._process_queue, daemon=True)
+        self._worker_thread.start()
+    
+    def add_download(self, repo_id: str, filename: str, url: str, save_path: str, 
+                     headers: Optional[Dict[str, str]] = None) -> str:
+        """Add a new download to the queue."""
+        download_id = f"{repo_id}_{filename}_{int(time.time())}"
+        
+        # Create download item
+        item = DownloadItem(
+            id=download_id,
+            repo_id=repo_id,
+            filename=filename,
+            url=url,
+            save_path=save_path,
+            headers=headers or {}
+        )
+        
+        self.downloads[download_id] = item
+        self.download_queue.put(download_id)
+        self._trigger_callback('on_status_change', item)
+        
+        return download_id
+    
+    def pause_download(self, download_id: str):
+        """Pause a download."""
+        if download_id in self._pause_flags:
+            self._pause_flags[download_id].set()
+            if download_id in self.downloads:
+                self.downloads[download_id].status = DownloadStatus.PAUSED
+                self._trigger_callback('on_status_change', self.downloads[download_id])
+    
+    def resume_download(self, download_id: str):
+        """Resume a paused download."""
+        item = self.downloads.get(download_id)
+        if item and item.is_resumable:
+            item.status = DownloadStatus.QUEUED
+            item.resume_position = item.downloaded_size
+            self.download_queue.put(download_id)
+            self._trigger_callback('on_status_change', item)
+    
+    def cancel_download(self, download_id: str):
+        """Cancel a download."""
+        if download_id in self._stop_flags:
+            self._stop_flags[download_id].set()
+        
+        if download_id in self.downloads:
+            self.downloads[download_id].status = DownloadStatus.CANCELLED
+            self._trigger_callback('on_status_change', self.downloads[download_id])
+            
+            # Remove partial file
+            item = self.downloads[download_id]
+            if os.path.exists(item.save_path):
+                try:
+                    os.remove(item.save_path)
+                except Exception:
+                    pass
+    
+    def retry_download(self, download_id: str):
+        """Retry a failed download."""
+        item = self.downloads.get(download_id)
+        if item and item.status in [DownloadStatus.FAILED, DownloadStatus.AUTH_REQUIRED]:
+            item.status = DownloadStatus.QUEUED
+            item.downloaded_size = 0
+            item.resume_position = 0
+            item.error_message = ""
+            self.download_queue.put(download_id)
+            self._trigger_callback('on_status_change', item)
+    
+    def _process_queue(self):
+        """Process download queue."""
+        while True:
+            # Check if we can start more downloads
+            if len(self.active_downloads) < self.max_concurrent:
+                try:
+                    download_id = self.download_queue.get(timeout=1)
+                    if download_id in self.downloads:
+                        thread = threading.Thread(
+                            target=self._download_file,
+                            args=(download_id,),
+                            daemon=True
+                        )
+                        self.active_downloads[download_id] = thread
+                        thread.start()
+                except queue.Empty:
+                    pass
+            
+            # Clean up finished downloads
+            finished = []
+            for download_id, thread in self.active_downloads.items():
+                if not thread.is_alive():
+                    finished.append(download_id)
+            
+            for download_id in finished:
+                del self.active_downloads[download_id]
+            
+            time.sleep(0.5)
+    
+    def _download_file(self, download_id: str):
+        """Download a file with resume support."""
+        item = self.downloads[download_id]
+        
+        # Create flags for this download
+        self._stop_flags[download_id] = threading.Event()
+        self._pause_flags[download_id] = threading.Event()
+        
+        try:
+            # Update status
+            item.status = DownloadStatus.DOWNLOADING
+            item.start_time = time.time()
+            self._trigger_callback('on_status_change', item)
+            
+            # Create directory if needed
+            os.makedirs(os.path.dirname(item.save_path), exist_ok=True)
+            
+            # Setup headers for resume
+            headers = item.headers.copy()
+            if item.resume_position > 0:
+                headers['Range'] = f'bytes={item.resume_position}-'
+            
+            # Make request
+            response = requests.get(item.url, headers=headers, stream=True, timeout=30)
+            
+            # Check for authentication issues
+            if response.status_code == 401 or response.status_code == 403:
+                item.status = DownloadStatus.AUTH_REQUIRED
+                item.error_message = f"Authentication failed: {response.status_code}"
+                self._trigger_callback('on_error', item)
+                return
+            
+            response.raise_for_status()
+            
+            # Get total size
+            if item.resume_position == 0:
+                item.total_size = int(response.headers.get('content-length', 0))
+            
+            # Open file for writing (append if resuming)
+            mode = 'ab' if item.resume_position > 0 else 'wb'
+            
+            with open(item.save_path, mode) as f:
+                chunk_size = 8192
+                last_update = time.time()
+                bytes_since_update = 0
+                
+                for chunk in response.iter_content(chunk_size=chunk_size):
+                    # Check stop flag
+                    if self._stop_flags[download_id].is_set():
+                        return
+                    
+                    # Check pause flag
+                    if self._pause_flags[download_id].is_set():
+                        item.status = DownloadStatus.PAUSED
+                        self._trigger_callback('on_status_change', item)
+                        return
+                    
+                    if chunk:
+                        f.write(chunk)
+                        item.downloaded_size += len(chunk)
+                        bytes_since_update += len(chunk)
+                        
+                        # Calculate speed and ETA
+                        current_time = time.time()
+                        time_diff = current_time - last_update
+                        
+                        if time_diff >= 1.0:  # Update every second
+                            item.speed = bytes_since_update / time_diff
+                            
+                            if item.speed > 0:
+                                remaining = item.total_size - item.downloaded_size
+                                item.eta = int(remaining / item.speed)
+                            
+                            self._trigger_callback('on_progress', item)
+                            last_update = current_time
+                            bytes_since_update = 0
+            
+            # Download completed
+            item.status = DownloadStatus.COMPLETED
+            item.end_time = time.time()
+            self._trigger_callback('on_complete', item)
+            
+        except requests.exceptions.RequestException as e:
+            item.status = DownloadStatus.FAILED
+            item.error_message = str(e)
+            self._trigger_callback('on_error', item)
+            
+        except Exception as e:
+            item.status = DownloadStatus.FAILED
+            item.error_message = f"Unexpected error: {e}"
+            self._trigger_callback('on_error', item)
+        
+        finally:
+            # Clean up flags
+            if download_id in self._stop_flags:
+                del self._stop_flags[download_id]
+            if download_id in self._pause_flags:
+                del self._pause_flags[download_id]
+            
+            self._trigger_callback('on_status_change', item)
+    
+    def register_callback(self, event: str, callback: Callable):
+        """Register a callback for download events."""
+        if event in self.callbacks:
+            self.callbacks[event].append(callback)
+    
+    def _trigger_callback(self, event: str, item: DownloadItem):
+        """Trigger callbacks for an event."""
+        for callback in self.callbacks.get(event, []):
+            try:
+                callback(item)
+            except Exception as e:
+                print(f"Callback error: {e}")
+    
+    def get_all_downloads(self) -> List[DownloadItem]:
+        """Get all download items."""
+        return list(self.downloads.values())
+    
+    def clear_completed(self):
+        """Clear completed downloads from the list."""
+        to_remove = []
+        for download_id, item in self.downloads.items():
+            if item.status in [DownloadStatus.COMPLETED, DownloadStatus.CANCELLED]:
+                to_remove.append(download_id)
+        
+        for download_id in to_remove:
+            del self.downloads[download_id]
+
+
+class DownloadManagerTab:
+    """Download Manager GUI tab."""
+    
+    def __init__(self, parent: ttk.Frame, download_manager: DownloadManager):
+        self.parent = parent
+        self.manager = download_manager
+        self.item_widgets: Dict[str, Dict[str, Any]] = {}
+        
+        # Register callbacks
+        self.manager.register_callback('on_progress', self._on_progress)
+        self.manager.register_callback('on_status_change', self._on_status_change)
+        self.manager.register_callback('on_complete', self._on_complete)
+        self.manager.register_callback('on_error', self._on_error)
+        
+        self._build_ui()
+        
+        # Start update timer
+        self._update_display()
+    
+    def _build_ui(self):
+        """Build the download manager UI."""
+        # Top controls
+        controls_frame = ttk.Frame(self.parent)
+        controls_frame.pack(fill=tk.X, padx=5, pady=5)
+        
+        ttk.Button(controls_frame, text="Clear Completed", 
+                  command=self._clear_completed).pack(side=tk.LEFT, padx=5)
+        ttk.Button(controls_frame, text="Pause All", 
+                  command=self._pause_all).pack(side=tk.LEFT, padx=5)
+        ttk.Button(controls_frame, text="Resume All", 
+                  command=self._resume_all).pack(side=tk.LEFT, padx=5)
+        
+        # Status summary
+        self.status_label = ttk.Label(controls_frame, text="Downloads: 0 active, 0 queued")
+        self.status_label.pack(side=tk.RIGHT, padx=5)
+        
+        # Downloads list with scrollbar
+        list_frame = ttk.Frame(self.parent)
+        list_frame.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
+        
+        # Canvas for scrollable content
+        self.canvas = tk.Canvas(list_frame, bg='white')
+        scrollbar = ttk.Scrollbar(list_frame, orient="vertical", command=self.canvas.yview)
+        self.scrollable_frame = ttk.Frame(self.canvas)
+        
+        self.scrollable_frame.bind(
+            "<Configure>",
+            lambda e: self.canvas.configure(scrollregion=self.canvas.bbox("all"))
+        )
+        
+        self.canvas.create_window((0, 0), window=self.scrollable_frame, anchor="nw")
+        self.canvas.configure(yscrollcommand=scrollbar.set)
+        
+        self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
+        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
+    
+    def add_download_widget(self, item: DownloadItem):
+        """Add a download widget to the display."""
+        if item.id in self.item_widgets:
+            return
+        
+        # Create frame for this download
+        frame = ttk.LabelFrame(self.scrollable_frame, text=f"{item.repo_id} / {item.filename}")
+        frame.pack(fill=tk.X, padx=5, pady=5)
+        
+        # Info row
+        info_frame = ttk.Frame(frame)
+        info_frame.pack(fill=tk.X, padx=5, pady=5)
+        
+        status_label = ttk.Label(info_frame, text=f"Status: {item.status.value}")
+        status_label.pack(side=tk.LEFT, padx=5)
+        
+        size_label = ttk.Label(info_frame, text="Size: -")
+        size_label.pack(side=tk.LEFT, padx=5)
+        
+        speed_label = ttk.Label(info_frame, text="Speed: -")
+        speed_label.pack(side=tk.LEFT, padx=5)
+        
+        eta_label = ttk.Label(info_frame, text="ETA: -")
+        eta_label.pack(side=tk.LEFT, padx=5)
+        
+        # Progress bar
+        progress_bar = ttk.Progressbar(frame, length=400, mode='determinate')
+        progress_bar.pack(fill=tk.X, padx=5, pady=5)
+        
+        # Error label (hidden by default)
+        error_label = ttk.Label(frame, text="", foreground="red")
+        
+        # Control buttons
+        controls_frame = ttk.Frame(frame)
+        controls_frame.pack(fill=tk.X, padx=5, pady=5)
+        
+        pause_btn = ttk.Button(controls_frame, text="Pause", width=10,
+                               command=lambda: self.manager.pause_download(item.id))
+        pause_btn.pack(side=tk.LEFT, padx=2)
+        
+        resume_btn = ttk.Button(controls_frame, text="Resume", width=10,
+                               command=lambda: self.manager.resume_download(item.id))
+        resume_btn.pack(side=tk.LEFT, padx=2)
+        
+        cancel_btn = ttk.Button(controls_frame, text="Cancel", width=10,
+                               command=lambda: self.manager.cancel_download(item.id))
+        cancel_btn.pack(side=tk.LEFT, padx=2)
+        
+        retry_btn = ttk.Button(controls_frame, text="Retry", width=10,
+                              command=lambda: self.manager.retry_download(item.id))
+        retry_btn.pack(side=tk.LEFT, padx=2)
+        
+        # Store widgets
+        self.item_widgets[item.id] = {
+            'frame': frame,
+            'status_label': status_label,
+            'size_label': size_label,
+            'speed_label': speed_label,
+            'eta_label': eta_label,
+            'progress_bar': progress_bar,
+            'error_label': error_label,
+            'pause_btn': pause_btn,
+            'resume_btn': resume_btn,
+            'cancel_btn': cancel_btn,
+            'retry_btn': retry_btn
+        }
+        
+        # Initial update
+        self._update_download_widget(item)
+    
+    def _update_download_widget(self, item: DownloadItem):
+        """Update a download widget."""
+        if item.id not in self.item_widgets:
+            self.add_download_widget(item)
+            return
+        
+        widgets = self.item_widgets[item.id]
+        
+        # Update status
+        widgets['status_label'].config(text=f"Status: {item.status.value}")
+        
+        # Update size
+        if item.total_size > 0:
+            size_text = f"Size: {self._format_size(item.downloaded_size)} / {self._format_size(item.total_size)}"
+        else:
+            size_text = f"Size: {self._format_size(item.downloaded_size)}"
+        widgets['size_label'].config(text=size_text)
+        
+        # Update speed
+        if item.speed > 0:
+            widgets['speed_label'].config(text=f"Speed: {self._format_size(item.speed)}/s")
+        else:
+            widgets['speed_label'].config(text="Speed: -")
+        
+        # Update ETA
+        if item.eta > 0:
+            eta_text = self._format_time(item.eta)
+            widgets['eta_label'].config(text=f"ETA: {eta_text}")
+        else:
+            widgets['eta_label'].config(text="ETA: -")
+        
+        # Update progress bar
+        widgets['progress_bar']['value'] = item.progress
+        
+        # Update error message
+        if item.error_message:
+            widgets['error_label'].config(text=f"Error: {item.error_message}")
+            widgets['error_label'].pack(fill=tk.X, padx=5, pady=2)
+        else:
+            widgets['error_label'].pack_forget()
+        
+        # Update button states
+        if item.status == DownloadStatus.DOWNLOADING:
+            widgets['pause_btn'].config(state="normal")
+            widgets['resume_btn'].config(state="disabled")
+            widgets['cancel_btn'].config(state="normal")
+            widgets['retry_btn'].config(state="disabled")
+        elif item.status == DownloadStatus.PAUSED:
+            widgets['pause_btn'].config(state="disabled")
+            widgets['resume_btn'].config(state="normal")
+            widgets['cancel_btn'].config(state="normal")
+            widgets['retry_btn'].config(state="disabled")
+        elif item.status in [DownloadStatus.FAILED, DownloadStatus.AUTH_REQUIRED]:
+            widgets['pause_btn'].config(state="disabled")
+            widgets['resume_btn'].config(state="disabled")
+            widgets['cancel_btn'].config(state="disabled")
+            widgets['retry_btn'].config(state="normal")
+        elif item.status == DownloadStatus.COMPLETED:
+            widgets['pause_btn'].config(state="disabled")
+            widgets['resume_btn'].config(state="disabled")
+            widgets['cancel_btn'].config(state="disabled")
+            widgets['retry_btn'].config(state="disabled")
+        else:
+            widgets['pause_btn'].config(state="disabled")
+            widgets['resume_btn'].config(state="disabled")
+            widgets['cancel_btn'].config(state="normal")
+            widgets['retry_btn'].config(state="disabled")
+    
+    def _format_size(self, bytes_size: float) -> str:
+        """Format bytes to human readable size."""
+        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
+            if bytes_size < 1024.0:
+                return f"{bytes_size:.1f} {unit}"
+            bytes_size /= 1024.0
+        return f"{bytes_size:.1f} PB"
+    
+    def _format_time(self, seconds: int) -> str:
+        """Format seconds to human readable time."""
+        if seconds < 60:
+            return f"{seconds}s"
+        elif seconds < 3600:
+            minutes = seconds // 60
+            secs = seconds % 60
+            return f"{minutes}m {secs}s"
+        else:
+            hours = seconds // 3600
+            minutes = (seconds % 3600) // 60
+            return f"{hours}h {minutes}m"
+    
+    def _on_progress(self, item: DownloadItem):
+        """Handle progress update."""
+        self.parent.after(0, lambda: self._update_download_widget(item))
+    
+    def _on_status_change(self, item: DownloadItem):
+        """Handle status change."""
+        self.parent.after(0, lambda: self._update_download_widget(item))
+    
+    def _on_complete(self, item: DownloadItem):
+        """Handle download completion."""
+        self.parent.after(0, lambda: self._update_download_widget(item))
+    
+    def _on_error(self, item: DownloadItem):
+        """Handle download error."""
+        self.parent.after(0, lambda: self._update_download_widget(item))
+        
+        # Show error notification for auth issues
+        if item.status == DownloadStatus.AUTH_REQUIRED:
+            self.parent.after(0, lambda: messagebox.showerror(
+                "Authentication Required",
+                f"Authentication failed for {item.filename}.\n"
+                f"Please check your API key in Settings."
+            ))
+    
+    def _clear_completed(self):
+        """Clear completed downloads."""
+        # Remove widgets for completed downloads
+        for download_id in list(self.item_widgets.keys()):
+            item = self.manager.downloads.get(download_id)
+            if item and item.status in [DownloadStatus.COMPLETED, DownloadStatus.CANCELLED]:
+                self.item_widgets[download_id]['frame'].destroy()
+                del self.item_widgets[download_id]
+        
+        # Clear from manager
+        self.manager.clear_completed()
+    
+    def _pause_all(self):
+        """Pause all active downloads."""
+        for item in self.manager.get_all_downloads():
+            if item.status == DownloadStatus.DOWNLOADING:
+                self.manager.pause_download(item.id)
+    
+    def _resume_all(self):
+        """Resume all paused downloads."""
+        for item in self.manager.get_all_downloads():
+            if item.status == DownloadStatus.PAUSED:
+                self.manager.resume_download(item.id)
+    
+    def _update_display(self):
+        """Update the display periodically."""
+        # Update status summary
+        all_downloads = self.manager.get_all_downloads()
+        active = sum(1 for d in all_downloads if d.status == DownloadStatus.DOWNLOADING)
+        queued = sum(1 for d in all_downloads if d.status == DownloadStatus.QUEUED)
+        self.status_label.config(text=f"Downloads: {active} active, {queued} queued")
+        
+        # Add any new downloads
+        for item in all_downloads:
+            if item.id not in self.item_widgets:
+                self.add_download_item(item)
+        
+        # Schedule next update
+        self.parent.after(500, self._update_display)