moma/renamer/extractors/mediainfo_extractor.py

from pathlib import Path
from pymediainfo import MediaInfo
from collections import Counter
from ..constants import FRAME_CLASSES
from ..formatters.color_formatter import ColorFormatter


class MediaInfoExtractor:
    """Class to extract information from MediaInfo"""

    def __init__(self, file_path: Path):
        self.file_path = file_path
        try:
            self.media_info = MediaInfo.parse(file_path)
            self.video_tracks = [t for t in self.media_info.tracks if t.track_type == 'Video']
            self.audio_tracks = [t for t in self.media_info.tracks if t.track_type == 'Audio']
            self.sub_tracks = [t for t in self.media_info.tracks if t.track_type == 'Text']
        except Exception:
            self.media_info = None
            self.video_tracks = []
            self.audio_tracks = []
            self.sub_tracks = []

    def _get_frame_class_from_height(self, height: int) -> str:
        """Get frame class from video height using FRAME_CLASSES constant"""
        for frame_class, info in FRAME_CLASSES.items():
            if height == info['nominal_height']:
                return frame_class
        return 'Unclassified'

    def extract_frame_class(self) -> str | None:
        """Extract frame class from media info (480p, 720p, 1080p, etc.)"""
        if not self.video_tracks:
            return 'Unclassified'
        height = getattr(self.video_tracks[0], 'height', None)
        if height:
            return self._get_frame_class_from_height(height)
        return 'Unclassified'

    def extract_resolution(self) -> str | None:
        """Extract actual video resolution (WIDTHxHEIGHT) from media info"""
        if not self.video_tracks:
            return None
        width = getattr(self.video_tracks[0], 'width', None)
        height = getattr(self.video_tracks[0], 'height', None)
        if width and height:
            return f"{width}x{height}"
        return None

    def extract_aspect_ratio(self) -> str | None:
        """Extract video aspect ratio from media info"""
        if not self.video_tracks:
            return None
        aspect_ratio = getattr(self.video_tracks[0], 'display_aspect_ratio', None)
        if aspect_ratio:
            return str(aspect_ratio)
        return None

    def extract_hdr(self) -> str | None:
        """Extract HDR info from media info"""
        if not self.video_tracks:
            return None
        profile = getattr(self.video_tracks[0], 'format_profile', '')
        if 'HDR' in profile.upper():
            return 'HDR'
        return None

    def extract_audio_langs(self) -> str:
        """Extract audio languages from media info"""
        if not self.audio_tracks:
            return ''
        lang_map = {
            'en': 'eng', 'fr': 'fre', 'de': 'ger', 'uk': 'ukr', 'ru': 'rus',
            'es': 'spa', 'it': 'ita', 'pt': 'por', 'ja': 'jpn', 'ko': 'kor',
            'zh': 'chi', 'und': 'und'
        }
        langs = [getattr(a, 'language', 'und').lower()[:3] for a in self.audio_tracks]
        langs = [lang_map.get(lang, lang) for lang in langs]
        lang_counts = Counter(langs)
        audio_langs = [f"{count}{lang}" if count > 1 else lang for lang, count in lang_counts.items()]
        return ','.join(audio_langs)

    def extract_video_dimensions(self) -> tuple[int, int] | None:
        """Extract video width and height"""
        if not self.video_tracks:
            return None
        width = getattr(self.video_tracks[0], 'width', None)
        height = getattr(self.video_tracks[0], 'height', None)
        if width and height:
            return width, height
        return None

    def extract_tracks(self) -> str:
        """Extract compact media track information"""
        tracks_info = []
        try:
            # Video tracks
            for i, v in enumerate(self.video_tracks[:2]):  # Up to 2 videos
                codec = getattr(v, 'format', None) or getattr(v, 'codec', None) or 'unknown'
                width = getattr(v, 'width', None) or '?'
                height = getattr(v, 'height', None) or '?'
                bitrate = getattr(v, 'bit_rate', None)
                fps = getattr(v, 'frame_rate', None)
                profile = getattr(v, 'format_profile', None)
                
                video_str = f"{codec} {width}x{height}"
                if bitrate:
                    video_str += f" {bitrate}bps"
                if fps:
                    video_str += f" {fps}fps"
                if profile:
                    video_str += f" ({profile})"
                
                tracks_info.append(ColorFormatter.green(f"Video {i+1}: {video_str}"))
            
            # Audio tracks
            for i, a in enumerate(self.audio_tracks[:3]):  # Up to 3 audios
                codec = getattr(a, 'format', None) or getattr(a, 'codec', None) or 'unknown'
                channels = getattr(a, 'channel_s', None) or '?'
                lang = getattr(a, 'language', None) or 'und'
                bitrate = getattr(a, 'bit_rate', None)
                
                audio_str = f"{codec} {channels}ch {lang}"
                if bitrate:
                    audio_str += f" {bitrate}bps"
                
                tracks_info.append(ColorFormatter.yellow(f"Audio {i+1}: {audio_str}"))
            
            # Subtitle tracks
            for i, s in enumerate(self.sub_tracks[:3]):  # Up to 3 subs
                lang = getattr(s, 'language', None) or 'und'
                format = getattr(s, 'format', None) or getattr(s, 'codec', None) or 'unknown'
                
                sub_str = f"{lang} ({format})"
                tracks_info.append(ColorFormatter.magenta(f"Sub {i+1}: {sub_str}"))
                
        except Exception as e:
            tracks_info.append(ColorFormatter.red(f"Track info error: {str(e)}"))
        
        return "\n".join(tracks_info) if tracks_info else ""