utils.py

from __future__ import annotations

import io
import logging
import os
import tempfile
from collections.abc import Callable
from pathlib import Path
from typing import Any, TypeVar

import numpy as np
from PIL import Image

LOGGER = logging.getLogger(__name__)
T = TypeVar("T")


def image_to_bytes(image: Image.Image | np.ndarray | str | Path) -> bytes:
    if isinstance(image, (str, Path)):
        image = Image.open(image)
    elif isinstance(image, np.ndarray):
        image = Image.fromarray(image)

    if not isinstance(image, Image.Image):
        raise TypeError("Expected a PIL image, NumPy array, or image path.")

    buffer = io.BytesIO()
    image.convert("RGB").save(buffer, format="JPEG", quality=92)
    return buffer.getvalue()


WAV_PREFIX = "thirdeye_"


def bytes_to_wav(audio_bytes: bytes) -> str:
    if not audio_bytes:
        raise ValueError("Audio data is empty.")

    output = tempfile.NamedTemporaryFile(delete=False, prefix=WAV_PREFIX, suffix=".wav")
    output.write(audio_bytes)
    output.close()
    return output.name


def cleanup_wav(path: str | None) -> None:
    if path and os.path.exists(path):
        try:
            os.unlink(path)
        except OSError:
            pass


def prune_old_wavs(max_age_seconds: float = 600.0) -> int:
    """Best-effort cleanup of stale Third Eye temp WAVs.

    Only removes files older than ``max_age_seconds`` so a clip currently being
    served to the browser is never deleted. Returns the number of files removed.
    """
    import time

    removed = 0
    temp_dir = Path(tempfile.gettempdir())
    cutoff = time.time() - max_age_seconds
    try:
        candidates = temp_dir.glob(f"{WAV_PREFIX}*.wav")
    except OSError:
        return 0
    for candidate in candidates:
        try:
            if candidate.stat().st_mtime < cutoff:
                candidate.unlink()
                removed += 1
        except OSError:
            continue
    return removed


def read_audio_bytes(audio_path: str | Path) -> bytes:
    return Path(audio_path).read_bytes()


def safe_call(
    fn: Callable[..., T],
    *args: Any,
    fallback: T | None = None,
    warn: str = "Something went wrong.",
    **kwargs: Any,
) -> T | None:
    try:
        return fn(*args, **kwargs)
    except Exception as exc:
        LOGGER.exception("Third Eye stage failed: %s", getattr(fn, "__name__", fn))
        detail = f"{warn} ({type(exc).__name__}: {exc})"
        try:
            import gradio as gr

            gr.Warning(detail if len(detail) < 200 else warn)
        except Exception:
            LOGGER.warning(detail)
        return fallback