app/services/ffprobe_service.py

from __future__ import annotations

import asyncio
import json
from fractions import Fraction
from pathlib import Path
from typing import Any

from app.core.config import Settings
from app.core.exceptions import ProcessingError
from app.core.logger import get_logger

logger = get_logger(__name__)


class FFprobeService:
    def __init__(self, settings: Settings) -> None:
        self.settings = settings
        self._semaphore = asyncio.Semaphore(settings.max_workers)

    async def probe(self, path: Path) -> dict[str, Any]:
        command = [
            self.settings.ffprobe_binary,
            "-v",
            "error",
            "-show_format",
            "-show_streams",
            "-print_format",
            "json",
            str(path),
        ]
        try:
            async with self._semaphore:
                process = await asyncio.create_subprocess_exec(
                    *command,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=120)
        except FileNotFoundError as exc:
            raise ProcessingError("FFprobe is not installed or not available") from exc
        except asyncio.TimeoutError as exc:
            if "process" in locals():
                process.kill()
                await process.wait()
            raise ProcessingError("FFprobe timed out") from exc
        if process.returncode != 0:
            message = stderr.decode("utf-8", errors="replace")[-2_000:]
            logger.error("ffprobe failed", extra={"command": command, "stderr": message})
            raise ProcessingError("Unable to read media metadata")
        try:
            raw = json.loads(stdout)
        except json.JSONDecodeError as exc:
            raise ProcessingError("FFprobe returned invalid metadata") from exc
        return self._normalize(raw)

    @classmethod
    def _normalize(cls, raw: dict[str, Any]) -> dict[str, Any]:
        streams = raw.get("streams", [])
        fmt = raw.get("format", {})
        videos = [stream for stream in streams if stream.get("codec_type") == "video"]
        audios = [stream for stream in streams if stream.get("codec_type") == "audio"]
        subtitles = [stream for stream in streams if stream.get("codec_type") == "subtitle"]
        video = videos[0] if videos else {}
        tags = fmt.get("tags", {})
        video_tags = video.get("tags", {})
        rotation = video_tags.get("rotate")
        for item in video.get("side_data_list", []):
            if "rotation" in item:
                rotation = item["rotation"]
        duration = cls._float(fmt.get("duration"))
        if duration is None:
            duration = cls._float(video.get("duration"))
        return {
            "duration": duration,
            "resolution": (
                {
                    "width": video.get("width"),
                    "height": video.get("height"),
                }
                if video
                else None
            ),
            "fps": cls._fps(video.get("avg_frame_rate") or video.get("r_frame_rate")),
            "bitrate": cls._int(fmt.get("bit_rate")),
            "codec": video.get("codec_name") or (audios[0].get("codec_name") if audios else None),
            "video_streams": [cls._stream_summary(stream) for stream in videos],
            "audio_streams": [cls._stream_summary(stream) for stream in audios],
            "subtitle_streams": [cls._stream_summary(stream) for stream in subtitles],
            "rotation": cls._int(rotation),
            "container": fmt.get("format_name"),
            "creation_date": tags.get("creation_time") or video_tags.get("creation_time"),
            "size": cls._int(fmt.get("size")),
            "tags": tags,
        }

    @classmethod
    def _stream_summary(cls, stream: dict[str, Any]) -> dict[str, Any]:
        return {
            "index": stream.get("index"),
            "codec": stream.get("codec_name"),
            "profile": stream.get("profile"),
            "pixel_format": stream.get("pix_fmt"),
            "field_order": stream.get("field_order"),
            "sample_aspect_ratio": stream.get("sample_aspect_ratio"),
            "frame_count": cls._int(stream.get("nb_frames")),
            "bitrate": cls._int(stream.get("bit_rate")),
            "sample_rate": cls._int(stream.get("sample_rate")),
            "channels": stream.get("channels"),
            "width": stream.get("width"),
            "height": stream.get("height"),
            "language": stream.get("tags", {}).get("language"),
        }

    @staticmethod
    def _fps(value: str | None) -> float | None:
        if not value or value == "0/0":
            return None
        try:
            return round(float(Fraction(value)), 4)
        except (ValueError, ZeroDivisionError):
            return None

    @staticmethod
    def _int(value: Any) -> int | None:
        try:
            return int(value)
        except (TypeError, ValueError):
            return None

    @staticmethod
    def _float(value: Any) -> float | None:
        try:
            return round(float(value), 6)
        except (TypeError, ValueError):
            return None