shelfmark/backend.py

"""Backend logic for the book download application."""

import threading, time
import shutil
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple
import subprocess
import os
from concurrent.futures import ThreadPoolExecutor, Future
from threading import Event

from logger import setup_logger
from config import CUSTOM_SCRIPT
from env import (INGEST_DIR, DOWNLOAD_PATHS, TMP_DIR, MAIN_LOOP_SLEEP_TIME, USE_BOOK_TITLE,
                 MAX_CONCURRENT_DOWNLOADS, DOWNLOAD_PROGRESS_UPDATE_INTERVAL)
from models import book_queue, BookInfo, QueueStatus, SearchFilters
import book_manager

logger = setup_logger(__name__)

# Import WebSocket manager (will be initialized by app.py)
try:
    from websocket_manager import ws_manager
except ImportError:
    logger.warning("WebSocket manager not available")
    ws_manager = None

def _sanitize_filename(filename: str) -> str:
    """Sanitize a filename by replacing spaces with underscores and removing invalid characters."""
    keepcharacters = (' ','.','_')
    return "".join(c for c in filename if c.isalnum() or c in keepcharacters).rstrip()

def search_books(query: str, filters: SearchFilters) -> List[Dict[str, Any]]:
    """Search for books matching the query.
    
    Args:
        query: Search term
        filters: Search filters object
        
    Returns:
        List[Dict]: List of book information dictionaries
    """
    try:
        books = book_manager.search_books(query, filters)
        return [_book_info_to_dict(book) for book in books]
    except Exception as e:
        logger.error_trace(f"Error searching books: {e}")
        return []

def get_book_info(book_id: str) -> Optional[Dict[str, Any]]:
    """Get detailed information for a specific book.
    
    Args:
        book_id: Book identifier
        
    Returns:
        Optional[Dict]: Book information dictionary if found
    """
    try:
        book = book_manager.get_book_info(book_id)
        return _book_info_to_dict(book)
    except Exception as e:
        logger.error_trace(f"Error getting book info: {e}")
        return None

def queue_book(book_id: str, priority: int = 0) -> bool:
    """Add a book to the download queue with specified priority.
    
    Args:
        book_id: Book identifier
        priority: Priority level (lower number = higher priority)
        
    Returns:
        bool: True if book was successfully queued
    """
    try:
        book_info = book_manager.get_book_info(book_id)
        book_queue.add(book_id, book_info, priority)
        logger.info(f"Book queued with priority {priority}: {book_info.title}")
        
        # Broadcast status update via WebSocket
        if ws_manager and ws_manager.is_enabled():
            ws_manager.broadcast_status_update(queue_status())
        
        return True
    except Exception as e:
        logger.error_trace(f"Error queueing book: {e}")
        return False

def queue_status() -> Dict[str, Dict[str, Any]]:
    """Get current status of the download queue.
    
    Returns:
        Dict: Queue status organized by status type with serialized book data
    """
    status = book_queue.get_status()
    for _, books in status.items():
        for _, book_info in books.items():
            if book_info.download_path:
                if not os.path.exists(book_info.download_path):
                    book_info.download_path = None

    # Convert Enum keys to strings and BookInfo objects to dicts for JSON serialization
    return {
        status_type.value: {
            book_id: _book_info_to_dict(book_info)
            for book_id, book_info in books.items()
        }
        for status_type, books in status.items()
    }

def get_book_data(book_id: str) -> Tuple[Optional[bytes], BookInfo]:
    """Get book data for a specific book, including its title.
    
    Args:
        book_id: Book identifier
        
    Returns:
        Tuple[Optional[bytes], str]: Book data if available, and the book title
    """
    try:
        book_info = book_queue._book_data[book_id]
        path = book_info.download_path
        with open(path, "rb") as f:
            return f.read(), book_info
    except Exception as e:
        logger.error_trace(f"Error getting book data: {e}")
        if book_info:
            book_info.download_path = None
        return None, book_info if book_info else BookInfo(id=book_id, title="Unknown")

def _book_info_to_dict(book: BookInfo) -> Dict[str, Any]:
    """Convert BookInfo object to dictionary representation."""
    return {
        key: value for key, value in book.__dict__.items()
        if value is not None
    }

def _prepare_download_folder(book_info: BookInfo) -> Path:
    """Prepare final content-type subdir"""
    content = book_info.content
    content_dir = DOWNLOAD_PATHS.get(content) if content and content in DOWNLOAD_PATHS else INGEST_DIR
    os.makedirs(content_dir, exist_ok=True)
    return content_dir

def _download_book_with_cancellation(book_id: str, cancel_flag: Event) -> Optional[str]:
    """Download and process a book with cancellation support.
    
    Args:
        book_id: Book identifier
        cancel_flag: Threading event to signal cancellation
        
    Returns:
        str: Path to the downloaded book if successful, None otherwise
    """
    try:
        # Check for cancellation before starting
        if cancel_flag.is_set():
            logger.info(f"Download cancelled before starting: {book_id}")
            return None
            
        book_info = book_queue._book_data[book_id]
        logger.info(f"Starting download: {book_info.title}")

        if USE_BOOK_TITLE:
            book_name = _sanitize_filename(book_info.title)
        else:
            book_name = book_id
        # If format is not set, use the format of the first download URL
        if book_info.format == "":
            book_info.format = book_info.download_urls[0].split(".")[-1]
        book_name += f".{book_info.format}"
        book_path = TMP_DIR / book_name

        # Check cancellation before download
        if cancel_flag.is_set():
            logger.info(f"Download cancelled before book manager call: {book_id}")
            return None
        
        progress_callback = lambda progress: update_download_progress(book_id, progress)
        status_callback = lambda status: update_download_status(book_id, status)
        success_download_url = book_manager.download_book(book_info, book_path, progress_callback, cancel_flag, status_callback)
        
        # Stop progress updates
        cancel_flag.wait(0.1)  # Brief pause for progress thread cleanup
        
        if cancel_flag.is_set():
            logger.info(f"Download cancelled during download: {book_id}")
            # Clean up partial download
            if book_path.exists():
                book_path.unlink()
            return None
            
        if not success_download_url:
            raise Exception("Unknown error downloading book")

        # Check cancellation before post-processing
        if cancel_flag.is_set():
            logger.info(f"Download cancelled before post-processing: {book_id}")
            if book_path.exists():
                book_path.unlink()
            return None

        # Update status to verifying
        book_queue.update_status(book_id, QueueStatus.VERIFYING)
        if ws_manager and ws_manager.is_enabled():
            ws_manager.broadcast_status_update(queue_status())
        logger.info(f"Verifying download: {book_info.title}")

        if CUSTOM_SCRIPT:
            logger.info(f"Running custom script: {CUSTOM_SCRIPT}")
            subprocess.run([CUSTOM_SCRIPT, book_path])
            
        # Update status to ingesting
        book_queue.update_status(book_id, QueueStatus.INGESTING)
        if ws_manager and ws_manager.is_enabled():
            ws_manager.broadcast_status_update(queue_status())
        
        if success_download_url and book_info.format == "":
            book_info.format = success_download_url.split(".")[-1]
            book_name += f".{book_info.format}"

        final_dir = _prepare_download_folder(book_info)
        intermediate_path = final_dir / f"{book_id}.crdownload"
        final_path = final_dir / book_name
        
        if os.path.exists(book_path):
            logger.info(f"Moving book to ingest directory: {book_path} -> {final_path}")
            try:
                shutil.move(book_path, intermediate_path)
            except Exception as e:
                try:
                    logger.debug(f"Error moving book: {e}, will try copying instead")
                    shutil.move(book_path, intermediate_path)
                except Exception as e:
                    logger.debug(f"Error copying book: {e}, will try copying without permissions instead")
                    shutil.copyfile(book_path, intermediate_path)
                os.remove(book_path)
            
            # Final cancellation check before completing
            if cancel_flag.is_set():
                logger.info(f"Download cancelled before final rename: {book_id}")
                if intermediate_path.exists():
                    intermediate_path.unlink()
                return None
                
            os.rename(intermediate_path, final_path)
            logger.info(f"Download completed successfully: {book_info.title}")
            
        return str(final_path)
    except Exception as e:
        if cancel_flag.is_set():
            logger.info(f"Download cancelled during error handling: {book_id}")
        else:
            logger.error_trace(f"Error downloading book: {e}")
        return None

def update_download_progress(book_id: str, progress: float) -> None:
    """Update download progress."""
    book_queue.update_progress(book_id, progress)
    
    # Broadcast progress via WebSocket
    if ws_manager and ws_manager.is_enabled():
        ws_manager.broadcast_download_progress(book_id, progress, 'downloading')

def update_download_status(book_id: str, status: str) -> None:
    """Update download status."""
    # Map string status to QueueStatus enum
    status_map = {
        'queued': QueueStatus.QUEUED,
        'resolving': QueueStatus.RESOLVING,
        'bypassing': QueueStatus.BYPASSING,
        'downloading': QueueStatus.DOWNLOADING,
        'verifying': QueueStatus.VERIFYING,
        'ingesting': QueueStatus.INGESTING,
        'complete': QueueStatus.COMPLETE,
        'available': QueueStatus.AVAILABLE,
        'error': QueueStatus.ERROR,
        'done': QueueStatus.DONE,
        'cancelled': QueueStatus.CANCELLED,
    }
    
    queue_status_enum = status_map.get(status.lower())
    if queue_status_enum:
        book_queue.update_status(book_id, queue_status_enum)
        
        # Broadcast status update via WebSocket
        if ws_manager and ws_manager.is_enabled():
            ws_manager.broadcast_status_update(queue_status())

def cancel_download(book_id: str) -> bool:
    """Cancel a download.
    
    Args:
        book_id: Book identifier to cancel
        
    Returns:
        bool: True if cancellation was successful
    """
    result = book_queue.cancel_download(book_id)
    
    # Broadcast status update via WebSocket
    if result and ws_manager and ws_manager.is_enabled():
        ws_manager.broadcast_status_update(queue_status())
    
    return result

def set_book_priority(book_id: str, priority: int) -> bool:
    """Set priority for a queued book.
    
    Args:
        book_id: Book identifier
        priority: New priority level (lower = higher priority)
        
    Returns:
        bool: True if priority was successfully changed
    """
    return book_queue.set_priority(book_id, priority)

def reorder_queue(book_priorities: Dict[str, int]) -> bool:
    """Bulk reorder queue.
    
    Args:
        book_priorities: Dict mapping book_id to new priority
        
    Returns:
        bool: True if reordering was successful
    """
    return book_queue.reorder_queue(book_priorities)

def get_queue_order() -> List[Dict[str, any]]:
    """Get current queue order for display."""
    return book_queue.get_queue_order()

def get_active_downloads() -> List[str]:
    """Get list of currently active downloads."""
    return book_queue.get_active_downloads()

def clear_completed() -> int:
    """Clear all completed downloads from tracking."""
    return book_queue.clear_completed()

def _process_single_download(book_id: str, cancel_flag: Event) -> None:
    """Process a single download job."""
    try:
        # Status will be updated through callbacks during download process
        # (resolving -> bypassing -> downloading -> verifying -> ingesting -> complete)
        download_path = _download_book_with_cancellation(book_id, cancel_flag)
        
        if cancel_flag.is_set():
            book_queue.update_status(book_id, QueueStatus.CANCELLED)
            # Broadcast cancellation
            if ws_manager and ws_manager.is_enabled():
                ws_manager.broadcast_status_update(queue_status())
            return
            
        if download_path:
            book_queue.update_download_path(book_id, download_path)
            new_status = QueueStatus.COMPLETE
        else:
            new_status = QueueStatus.ERROR
            
        book_queue.update_status(book_id, new_status)
        
        # Broadcast final status (completed or error)
        if ws_manager and ws_manager.is_enabled():
            ws_manager.broadcast_status_update(queue_status())
        
        logger.info(
            f"Book {book_id} download {'successful' if download_path else 'failed'}"
        )
        
    except Exception as e:
        if not cancel_flag.is_set():
            logger.error_trace(f"Error in download processing: {e}")
            book_queue.update_status(book_id, QueueStatus.ERROR)
        else:
            logger.info(f"Download cancelled: {book_id}")
            book_queue.update_status(book_id, QueueStatus.CANCELLED)
        
        # Broadcast error/cancelled status
        if ws_manager and ws_manager.is_enabled():
            ws_manager.broadcast_status_update(queue_status())

def concurrent_download_loop() -> None:
    """Main download coordinator using ThreadPoolExecutor for concurrent downloads."""
    logger.info(f"Starting concurrent download loop with {MAX_CONCURRENT_DOWNLOADS} workers")
    
    with ThreadPoolExecutor(max_workers=MAX_CONCURRENT_DOWNLOADS, thread_name_prefix="BookDownload") as executor:
        active_futures: Dict[Future, str] = {}  # Track active download futures
        
        while True:
            # Clean up completed futures
            completed_futures = [f for f in active_futures if f.done()]
            for future in completed_futures:
                book_id = active_futures.pop(future)
                try:
                    future.result()  # This will raise any exceptions from the worker
                except Exception as e:
                    logger.error_trace(f"Future exception for {book_id}: {e}")
            
            # Start new downloads if we have capacity
            while len(active_futures) < MAX_CONCURRENT_DOWNLOADS:
                next_download = book_queue.get_next()
                if not next_download:
                    break
                    
                book_id, cancel_flag = next_download
                logger.info(f"Starting concurrent download: {book_id}")
                
                # Submit download job to thread pool
                future = executor.submit(_process_single_download, book_id, cancel_flag)
                active_futures[future] = book_id
            
            # Brief sleep to prevent busy waiting
            time.sleep(MAIN_LOOP_SLEEP_TIME)

# Start concurrent download coordinator
download_coordinator_thread = threading.Thread(
    target=concurrent_download_loop,
    daemon=True,
    name="DownloadCoordinator"
)
download_coordinator_thread.start()

logger.info(f"Download system initialized with {MAX_CONCURRENT_DOWNLOADS} concurrent workers")