mule-image/backend/app/routers/download.py

"""
Download router — streams a .zip of every photo in a folder (recursively)
or a heap back to the browser.

Auth: both endpoints accept the regular Authorization header *or* a
``?token=JWT`` query string, mirroring the media endpoints. That lets the
frontend trigger a download with a plain ``<a href>`` (which can't set a
header), keeping the client side a one-liner.

Implementation: we build the zip into a ``NamedTemporaryFile`` and then
stream its bytes back, deleting the temp file on the way out. Stored
(uncompressed) mode because photos and videos are already compressed —
deflating them again just burns CPU for a fraction of a percent. For
very large libraries the temp-file route is mildly wasteful vs. a true
streaming zip (zipstream-ng etc), but it avoids a new dependency and
handles arbitrary folder sizes without blowing out RAM.
"""
import logging
import os
import re
import tempfile
import zipfile
from typing import List

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.database import get_db
from app.dependencies import get_current_user_media
from app.models import Folder, Heap, Photo, SourceRoot
from app.models.heaps import heap_photos
from app.models.user import User

logger = logging.getLogger(__name__)

router = APIRouter()


def _safe_filename(name: str) -> str:
    """Strip characters that Content-Disposition or Windows filesystems
    would choke on. Keeps the download's filename readable without
    needing any escaping on the client side."""
    cleaned = re.sub(r'[\\/:*?"<>|\r\n\t]', '_', name).strip().strip('.')
    return cleaned or 'download'


async def _collect_folder_photos(
    folder_id: str,
    user: User,
    db: AsyncSession,
) -> tuple[str, str, List[Photo]]:
    """Resolve a folder or source-root id → (base_path, display_name,
    photos). ``base_path`` is the prefix we strip off each photo's
    filepath when naming zip entries, so the archive mirrors the user's
    on-disk structure under that folder.
    """
    folder = (await db.execute(
        select(Folder).where(Folder.id == folder_id, Folder.user_id == user.id)
    )).scalar_one_or_none()

    base_path: str
    display_name: str
    if folder is not None:
        base_path = os.path.normpath(folder.path)
        display_name = folder.name or os.path.basename(base_path)
    else:
        sr = (await db.execute(
            select(SourceRoot).where(
                SourceRoot.id == folder_id,
                SourceRoot.user_id == user.id,
            )
        )).scalar_one_or_none()
        if sr is None:
            raise HTTPException(status_code=404, detail="Folder not found")
        base_path = os.path.normpath(sr.path)
        display_name = sr.name or os.path.basename(base_path)

    # Every photo whose filepath is at or below the base path — matches
    # the same prefix convention folders.py uses for recursive deletes.
    descendant_prefix = base_path.rstrip(os.sep) + os.sep
    result = await db.execute(
        select(Photo).where(
            Photo.user_id == user.id,
            Photo.is_discarded == False,  # noqa: E712
            (Photo.filepath == base_path) | (Photo.filepath.like(descendant_prefix + '%')),
        )
    )
    photos = list(result.scalars().all())
    return base_path, display_name, photos


def _build_zip(
    photos: List[Photo],
    arcname_fn,
) -> tempfile.NamedTemporaryFile:
    """Write ``photos`` into a fresh ZIP_STORED temp file.

    ``arcname_fn(photo, used_names)`` returns the entry name to use for
    the given photo; the caller supplies it because folder downloads
    want path-preserving names while heap downloads flatten to bare
    filenames (with a collision suffix).
    """
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.zip')
    try:
        used: set[str] = set()
        with zipfile.ZipFile(tmp, 'w', zipfile.ZIP_STORED, allowZip64=True) as zf:
            for p in photos:
                if not p.filepath or not os.path.exists(p.filepath):
                    # Silent skip: the scanner may have indexed files
                    # that have since been moved / unlinked by a shell.
                    continue
                name = arcname_fn(p, used)
                used.add(name)
                try:
                    zf.write(p.filepath, name)
                except OSError as e:
                    logger.warning(f"Skipping {p.filepath} in zip: {e}")
        tmp.close()
        return tmp
    except Exception:
        tmp.close()
        try:
            os.unlink(tmp.name)
        except OSError:
            pass
        raise


def _stream_and_cleanup(path: str):
    """Yield the temp zip in 1 MiB chunks and unlink it when the
    iterator is exhausted (or GC'd, if the client disconnects early)."""
    try:
        with open(path, 'rb') as f:
            while True:
                chunk = f.read(1024 * 1024)
                if not chunk:
                    break
                yield chunk
    finally:
        try:
            os.unlink(path)
        except OSError as e:
            logger.debug(f"Temp zip cleanup failed for {path}: {e}")


def _dedupe(name: str, used: set[str]) -> str:
    """Return ``name`` (or ``name (2)``, ``name (3)`` ...) such that the
    result doesn't collide with anything in ``used``. Needed for heap
    downloads where two members can have identical filenames from
    different folders."""
    if name not in used:
        return name
    stem, ext = os.path.splitext(name)
    n = 2
    while True:
        cand = f"{stem} ({n}){ext}"
        if cand not in used:
            return cand
        n += 1


@router.get("/folders/{folder_id}")
async def download_folder(
    folder_id: str,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user_media),
):
    """Zip every (non-discarded) photo under a folder/source-root and
    stream it back. Entries preserve the folder structure relative to
    the downloaded root so the resulting archive is a faithful snapshot.
    """
    base_path, display_name, photos = await _collect_folder_photos(
        folder_id, current_user, db
    )
    if not photos:
        raise HTTPException(status_code=404, detail="No photos to download")

    def arcname(p: Photo, _used: set[str]) -> str:
        # Relative path from the download root, falling back to the
        # bare filename if the photo somehow lives outside base_path.
        abs_path = os.path.normpath(p.filepath)
        if abs_path.startswith(base_path + os.sep):
            rel = abs_path[len(base_path) + 1:]
        elif abs_path == base_path:
            rel = os.path.basename(abs_path)
        else:
            rel = p.filename or os.path.basename(abs_path)
        # Nest everything under display_name so users see one top-level
        # folder inside the zip rather than loose files.
        return os.path.join(_safe_filename(display_name), rel)

    tmp = _build_zip(photos, arcname)
    filename = _safe_filename(display_name) + '.zip'
    return StreamingResponse(
        _stream_and_cleanup(tmp.name),
        media_type='application/zip',
        headers={
            'Content-Disposition': f'attachment; filename="{filename}"',
            'Content-Length': str(os.path.getsize(tmp.name)),
        },
    )


@router.get("/heaps/{heap_id}")
async def download_heap(
    heap_id: str,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user_media),
):
    """Zip every photo in a heap. Heaps are flat collections, so entries
    use the original filename (with a ``(2)`` collision suffix when
    two members share a name)."""
    heap = (await db.execute(
        select(Heap).where(Heap.id == heap_id, Heap.user_id == current_user.id)
    )).scalar_one_or_none()
    if heap is None:
        raise HTTPException(status_code=404, detail="Heap not found")

    result = await db.execute(
        select(Photo)
        .join(heap_photos, heap_photos.c.photo_id == Photo.id)
        .where(
            heap_photos.c.heap_id == heap_id,
            Photo.is_discarded == False,  # noqa: E712
        )
    )
    photos = list(result.scalars().all())
    if not photos:
        raise HTTPException(status_code=404, detail="Heap is empty")

    def arcname(p: Photo, used: set[str]) -> str:
        bare = p.filename or os.path.basename(p.filepath or 'photo')
        entry = os.path.join(_safe_filename(heap.name), _dedupe(bare, used))
        return entry

    tmp = _build_zip(photos, arcname)
    filename = _safe_filename(heap.name) + '.zip'
    return StreamingResponse(
        _stream_and_cleanup(tmp.name),
        media_type='application/zip',
        headers={
            'Content-Disposition': f'attachment; filename="{filename}"',
            'Content-Length': str(os.path.getsize(tmp.name)),
        },
    )