freepik/freepik_cli/utils/files.py

"""File I/O utilities: base64 encoding, output path generation, download."""

from __future__ import annotations

import base64
from datetime import datetime
from pathlib import Path
from typing import Optional

import httpx
from rich.console import Console
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    SpinnerColumn,
    TimeElapsedColumn,
    TransferSpeedColumn,
)


def image_to_base64(path: Path) -> str:
    """Read an image file and return a base64-encoded string.

    Uses Pillow to detect the actual format rather than trusting the file
    extension — mismatched extensions (e.g. a JPEG saved as .png) would
    produce an incorrect MIME type that causes silent failures with some models.
    """
    from PIL import Image

    _pillow_to_mime = {
        "JPEG": "image/jpeg",
        "PNG":  "image/png",
        "GIF":  "image/gif",
        "WEBP": "image/webp",
    }
    with Image.open(path) as img:
        fmt = img.format or "JPEG"
        mime = _pillow_to_mime.get(fmt, "image/jpeg")

    with open(path, "rb") as f:
        encoded = base64.b64encode(f.read()).decode()
    return f"data:{mime};base64,{encoded}"


def video_to_base64(path: Path) -> str:
    """Read a video file and return a base64-encoded string."""
    suffix = path.suffix.lower().lstrip(".")
    mime_map = {
        "mp4": "video/mp4",
        "mov": "video/quicktime",
        "avi": "video/x-msvideo",
        "webm": "video/webm",
    }
    mime = mime_map.get(suffix, "video/mp4")
    with open(path, "rb") as f:
        encoded = base64.b64encode(f.read()).decode()
    return f"data:{mime};base64,{encoded}"


def auto_output_path(task_type: str, model: str, ext: str = "jpg", output_dir: str = ".") -> Path:
    """Generate a timestamped output filename."""
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe_model = model.replace("/", "-").replace(" ", "-")
    filename = f"freepik_{task_type}_{safe_model}_{ts}.{ext}"
    return Path(output_dir) / filename


def save_from_url(url: str, dest: Path, console: Console) -> None:
    """Stream-download a URL to a file, displaying a Rich progress bar."""
    dest.parent.mkdir(parents=True, exist_ok=True)

    with Progress(
        SpinnerColumn(style="bold magenta"),
        "[progress.description]{task.description}",
        BarColumn(bar_width=30, style="magenta", complete_style="green"),
        "[progress.percentage]{task.percentage:>3.0f}%",
        DownloadColumn(),
        TransferSpeedColumn(),
        TimeElapsedColumn(),
        console=console,
        transient=True,
    ) as progress:
        with httpx.stream("GET", url, follow_redirects=True, timeout=120) as r:
            r.raise_for_status()
            total = int(r.headers.get("content-length", 0)) or None
            task = progress.add_task(
                f"[dim]Saving[/dim] [bold]{dest.name}[/bold]",
                total=total,
            )
            with open(dest, "wb") as f:
                for chunk in r.iter_bytes(chunk_size=65536):
                    f.write(chunk)
                    progress.advance(task, len(chunk))


def get_image_dimensions(path: Path) -> tuple[Optional[int], Optional[int]]:
    """Return (width, height) of an image using Pillow."""
    try:
        from PIL import Image

        with Image.open(path) as img:
            return img.width, img.height
    except Exception:
        return None, None