VidEnc/src/models/streams.py

from pydantic import BaseModel, Field
from enum import Enum
from datetime import datetime
from pathlib import Path
from functools import cached_property

from lib.fs import sanitize_filename

class OriginalMediaType(str, Enum):
    DVD = "DVD"
    BluRay = "BluRay"
    TVRip = "TVRip"
    WebRip = "WebRip"
    Other = "Other"


class Tags(BaseModel):
    title: str | None = Field(None)
    language: str | None = Field(None)


class Disposition(BaseModel):
    default: bool = Field(False)
    comment: bool = Field(False)


class VideoStream(BaseModel):
    index: int | None = Field(None)
    width: int | None = Field(None)
    height: int | None = Field(None)
    avg_frame_rate: str | None = Field(None, alias="avg_frame_rate")
    codec_name: str | None = Field(None)
    disposition: Disposition
    tags: Tags | None = Field(None)
    preset: int = Field(2)
    crf: int = Field(30)
    original_media_type: str | None = Field(None)
    film_grain: int = Field(20)

    @cached_property
    def fps(self) -> float:
        """Calculate frames per second from avg_frame_rate string."""
        try:
            if not self.avg_frame_rate:
                return 0.0
            num, denom = (int(x) for x in self.avg_frame_rate.split("/"))
            return num / denom if denom else 0.0
        except Exception:
            return 0.0

    def cmd_flags(self) -> list[str]:
        """Generate ffmpeg command flags for this video stream."""
        return [
            "-map",
            f"0:{self.index}",
            "-c:v",
            "libsvtav1",
            "-crf",
            str(self.crf),
            "-preset",
            str(self.preset),
            "-svtav1-params",
            f"tune=0:film-grain={self.film_grain}:scd=1",
            "-vf",
            "'scale=iw*sar:ih,setsar=1,scale=-2:ih:lanczos'",
            "-g",
            str(round(self.fps) * 5),
        ]


class AudioStream(BaseModel):
    index: int
    codec_name: str
    sample_fmt: str
    channel_layout: str
    disposition: Disposition
    tags: Tags | None = Field(None)

    def cmd_flags(self, index: int) -> list[str]:
        """Generate ffmpeg command flags for this audio stream."""
        side_loaded = []
        if "side" in self.channel_layout:
            side_loaded = [f"-filter:a:{index}", "'channelmap=channel_layout=5.1'"]
        default_disposition = (
            [f"-disposition:a:{index}", "default"] if self.disposition.default else []
        )
        language = getattr(self.tags, 'language', None) if self.tags else None
        title = getattr(self.tags, 'title', None) if self.tags else None
        meta_flags = []
        if language:
            meta_flags += [f"-metadata:s:a:{index}", f"language={language}"]
        if title:
            meta_flags += [f"-metadata:s:a:{index}", f"title='{title}'"]
        return (
            ["-map", f"0:{self.index}", f"-c:a:{index}", "libopus"]
            + side_loaded
            + meta_flags
            + default_disposition
        )


class SubtitleStream(BaseModel):
    index: int
    codec_name: str
    disposition: Disposition
    tags: Tags | None = Field(None)

    def cmd_flags(self, index: int) -> list[str]:
        """Generate ffmpeg command flags for this subtitle stream."""
        language = getattr(self.tags, 'language', None) if self.tags else None
        meta_flags = []
        if language:
            meta_flags += [f"-metadata:s:s:{index}", f"language={language}"]
        return ["-map", f"0:s:{index}", f"-c:s:{index}", "copy"] + meta_flags


class VideoFile(BaseModel):
    video_path: Path
    title: str | None = Field(None)
    date_released: datetime | None = Field(None)
    imdb: str | None = Field(None)
    video_streams: list[VideoStream] = Field(default_factory=list)
    audio_streams: list[AudioStream] = Field(default_factory=list)
    subtitle_streams: list[SubtitleStream] = Field(default_factory=list)
    allow_subtitle: bool = Field(True)
    original_media_type: OriginalMediaType = Field(OriginalMediaType.DVD)

    def cmd(self) -> list[str]:
        """Generate the full ffmpeg command for this video file."""
        cmd = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel",
            "error",
            "-i",
            f"'{str(self.video_path.absolute())}'",
            "-map_metadata",
            "-1",
        ]
        if self.title:
            cmd += ["-metadata", f"title='{self.title}'"]
        if self.date_released:
            cmd += ["-metadata", f"date_released={self.date_released.strftime('%Y-%m-%d')}"]
        if self.imdb:
            cmd += ["-metadata", f"imdb={self.imdb}"]
        cmd += ["-metadata", f"original_media_type={self.original_media_type.value}"]
        cmd += ["-pix_fmt", "yuv420p10le"]
        for stream in self.video_streams:
            cmd.extend(stream.cmd_flags())
        for i, stream in enumerate(self.audio_streams):
            cmd.extend(stream.cmd_flags(i))
        if self.allow_subtitle:
            for i, stream in enumerate(self.subtitle_streams):
                cmd.extend(stream.cmd_flags(i))
        year = self.date_released.year if self.date_released else "unknown"
        output_file = (
            self.video_path.absolute().parent
            / f"{sanitize_filename(self.title or 'untitled')}.{year}.mkv"
        )
        cmd.append(f"'{str(output_file)}'")
        return cmd