TwitchDownloader/modules/utils.py

"""
Utility functions and helpers for Twitch Archive.
"""

import os
import sys
import shutil
import pathlib
import subprocess
from typing import Optional
from colorama import Fore, Style


def detect_operating_system() -> str:
    """
    Detect the current operating system.
    
    Returns:
        str: 'windows' or 'linux'
        
    Raises:
        SystemExit: If OS is not supported
    """
    if sys.platform.startswith('win32'):
        return 'windows'
    elif sys.platform.startswith('linux'):
        return 'linux'
    else:
        print(f'{Fore.RED}✗ ERROR: Unsupported operating system: {sys.platform}{Style.RESET_ALL}')
        print(f'{Fore.YELLOW}  This script only supports Windows and Linux{Style.RESET_ALL}')
        sys.exit(1)


def get_bin_path() -> str:
    """Get the path to the bin directory containing external tools."""
    return str(pathlib.Path(__file__).parent.parent.resolve() / "bin")


def get_env_value(*names: str, default: Optional[str] = None) -> Optional[str]:
    """Return the first non-empty environment variable from the provided names."""
    for name in names:
        value = os.getenv(name)
        if value not in (None, ""):
            return value
    return default


def get_ffmpeg_executable(os_type: str) -> str:
    """
    Get the platform-specific ffmpeg executable path.
    
    Args:
        os_type: Operating system type ('windows' or 'linux')
        
    Returns:
        str: Path to ffmpeg executable
    """
    bin_path = get_bin_path()
    if os_type == 'windows':
        return os.path.join(bin_path, 'ffmpeg.exe')

    system_ffmpeg = shutil.which('ffmpeg')
    if system_ffmpeg:
        return system_ffmpeg

    return os.path.join(bin_path, 'ffmpeg')


def get_twitch_downloader_executable(os_type: str) -> str:
    """
    Get the platform-specific TwitchDownloaderCLI executable path.
    
    Args:
        os_type: Operating system type ('windows' or 'linux')
        
    Returns:
        str: Path to TwitchDownloaderCLI executable
    """
    bin_path = get_bin_path()
    if os_type == 'windows':
        return os.path.join(bin_path, 'TwitchDownloaderCLI.exe')

    system_twitch_downloader = shutil.which('TwitchDownloaderCLI')
    if system_twitch_downloader:
        return system_twitch_downloader

    return os.path.join(bin_path, 'TwitchDownloaderCLI')


def get_unique_filename(filepath: str) -> str:
    """
    Generate a unique filename by appending a counter if file already exists.
    
    Args:
        filepath: The desired file path
        
    Returns:
        str: A unique file path (original or with _N suffix)
        
    Example:
        If 'video.mp4' exists, returns 'video_1.mp4'
        If 'video_1.mp4' also exists, returns 'video_2.mp4'
    """
    if not os.path.exists(filepath):
        return filepath
    
    # Split into components
    directory = os.path.dirname(filepath)
    filename = os.path.basename(filepath)
    name, ext = os.path.splitext(filename)
    
    # Find next available counter
    counter = 1
    while True:
        new_filepath = os.path.join(directory, f"{name}_{counter}{ext}")
        if not os.path.exists(new_filepath):
            return new_filepath
        counter += 1


def get_video_duration(video_path: str, ffmpeg_path: str) -> Optional[float]:
    """
    Get the duration of a video file in seconds using FFmpeg.
    
    Args:
        video_path: Path to the video file
        ffmpeg_path: Path to FFmpeg executable
        
    Returns:
        float: Duration in seconds, or None if failed
    """
    if not os.path.exists(video_path):
        return None
    
    try:
        # Use ffprobe (comes with ffmpeg) to get duration
        ffprobe_path = ffmpeg_path.replace('ffmpeg', 'ffprobe')
        
        cmd = [
            ffprobe_path,
            '-v', 'error',
            '-show_entries', 'format=duration',
            '-of', 'default=noprint_wrappers=1:nokey=1',
            video_path
        ]
        
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=30
        )
        
        if result.returncode == 0 and result.stdout.strip():
            return float(result.stdout.strip())
    except Exception:
        pass
    
    return None


def verify_streamlink() -> bool:
    """
    Verify that streamlink is available.
    
    Returns:
        bool: True if streamlink is available, False otherwise
    """
    try:
        result = subprocess.run(['streamlink', '--version'], 
                              capture_output=True, 
                              text=True, 
                              timeout=5)
        if result.returncode == 0:
            version = result.stdout.strip().split()[1] if len(result.stdout.split()) > 1 else 'unknown'
            print(f'{Fore.GREEN}✓ Streamlink v{version} found{Style.RESET_ALL}')
            return True
        else:
            raise FileNotFoundError()
    except (FileNotFoundError, subprocess.TimeoutExpired, IndexError):
        print(f'{Fore.RED}✗ ERROR: Streamlink not found{Style.RESET_ALL}')
        print(f'{Fore.CYAN}  → Install streamlink: pip install streamlink{Style.RESET_ALL}')
        print(f'{Fore.CYAN}  → Or download from: https://streamlink.github.io/{Style.RESET_ALL}')
        return False


def verify_rclone() -> bool:
    """Verify that rclone is available on PATH."""
    try:
        result = subprocess.run(['rclone', 'version'],
                              capture_output=True,
                              text=True,
                              timeout=5)
        if result.returncode == 0:
            version_line = result.stdout.strip().splitlines()[0] if result.stdout.strip() else 'unknown'
            print(f'{Fore.GREEN}✓ Rclone found ({version_line}){Style.RESET_ALL}')
            return True
        raise FileNotFoundError()
    except (FileNotFoundError, subprocess.TimeoutExpired, IndexError):
        print(f'{Fore.RED}✗ ERROR: rclone not found{Style.RESET_ALL}')
        print(f'{Fore.CYAN}  → Install rclone and ensure it is on PATH{Style.RESET_ALL}')
        return False


def verify_ffmpeg(os_type: str) -> bool:
    """
    Verify that ffmpeg is available.
    
    Args:
        os_type: Operating system type ('windows' or 'linux')
        
    Returns:
        bool: True if ffmpeg is available, False otherwise
    """
    try:
        ffmpeg_path = get_ffmpeg_executable(os_type)
        if os.path.exists(ffmpeg_path):
            print(f'{Fore.GREEN}✓ FFmpeg found at {ffmpeg_path}{Style.RESET_ALL}')
            return True
        else:
            print(f'{Fore.YELLOW}⚠ Warning: FFmpeg not found at {ffmpeg_path}{Style.RESET_ALL}')
            print(f'{Fore.YELLOW}  → Download FFmpeg and place it in the bin/ folder{Style.RESET_ALL}')
            return False
    except Exception as e:
        print(f'{Fore.YELLOW}⚠ Warning: Could not verify FFmpeg: {e}{Style.RESET_ALL}')
        return False


def verify_twitch_downloader(os_type: str) -> bool:
    """
    Verify that TwitchDownloaderCLI is available.
    
    Args:
        os_type: Operating system type ('windows' or 'linux')
        
    Returns:
        bool: True if TwitchDownloaderCLI is available, False otherwise
    """
    try:
        downloader_path = get_twitch_downloader_executable(os_type)
        if os.path.exists(downloader_path):
            print(f'{Fore.GREEN}✓ TwitchDownloaderCLI found{Style.RESET_ALL}')
            return True
        else:
            print(f'{Fore.YELLOW}⚠ Warning: TwitchDownloaderCLI not found at {downloader_path}{Style.RESET_ALL}')
            print(f'{Fore.YELLOW}  → Download from: https://github.com/lay295/TwitchDownloader/releases{Style.RESET_ALL}')
            return False
    except Exception as e:
        print(f'{Fore.YELLOW}⚠ Warning: Could not verify TwitchDownloaderCLI: {e}{Style.RESET_ALL}')
        return False


def detect_hardware_acceleration(hwaccel_config: str, os_type: str) -> Optional[str]:
    """
    Detect available hardware acceleration based on config and system.
    
    Args:
        hwaccel_config: Hardware acceleration configuration ('auto', 'nvenc', 'qsv', 'amf', 'vaapi', 'none')
        os_type: Operating system type ('windows' or 'linux')
        
    Returns:
        str: Hardware acceleration type or None
    """
    # If user explicitly set to 'none', disable hardware acceleration
    if hwaccel_config == 'none':
        return 'none'
    
    # If user specified a particular type, use it
    if hwaccel_config in ['nvenc', 'qsv', 'amf', 'vaapi']:
        return hwaccel_config
    
    # Auto-detect: choose only hardware we can reasonably prove is present.
    if hwaccel_config == 'auto':
        if is_nvidia_runtime_available():
            return 'nvenc'
        if is_vaapi_runtime_available():
            return 'vaapi'
        return 'none'
    
    return None


def is_nvidia_runtime_available() -> bool:
    """Return True when the current runtime appears to expose an NVIDIA GPU."""
    visible_devices = os.getenv('NVIDIA_VISIBLE_DEVICES', '').strip().lower()
    if visible_devices in {'void', 'none'}:
        return False
    if visible_devices and visible_devices != 'all':
        return True

    if shutil.which('nvidia-smi'):
        return True

    return any(
        os.path.exists(device_path)
        for device_path in ('/dev/nvidiactl', '/dev/nvidia0', '/dev/nvidia-modeset')
    )


def is_vaapi_runtime_available() -> bool:
    """Return True when Linux VAAPI render nodes are present."""
    return any(
        os.path.exists(device_path)
        for device_path in ('/dev/dri/renderD128', '/dev/dri/card0')
    )


def resolve_hwaccel_type(hwaccel_type: Optional[str], os_type: str) -> Optional[str]:
    """Return a safe hardware acceleration choice for the current runtime."""
    if hwaccel_type in (None, 'none'):
        return 'none'

    if hwaccel_type == 'nvenc':
        return 'nvenc' if is_nvidia_runtime_available() else 'none'

    if hwaccel_type == 'vaapi':
        return 'vaapi' if is_vaapi_runtime_available() else 'none'

    # Leave explicit QSV/AMF unchanged for non-container users; container auto-detect no longer picks them blindly.
    return hwaccel_type


def get_hwaccel_encoder(hwaccel_type: str) -> str:
    """
    Get the appropriate hardware-accelerated encoder for the given acceleration type.
    
    Args:
        hwaccel_type: Type of hardware acceleration ('nvenc', 'qsv', 'amf', 'vaapi', 'auto', 'none')
        
    Returns:
        str: FFmpeg encoder name (e.g., 'h264_nvenc', 'libx264')
    """
    encoder_map = {
        'nvenc': 'h264_nvenc',  # NVIDIA
        'qsv': 'h264_qsv',      # Intel Quick Sync
        'amf': 'h264_amf',      # AMD
        'vaapi': 'h264_vaapi',  # Linux VA-API
    }
    
    if hwaccel_type in encoder_map:
        return encoder_map[hwaccel_type]
    elif hwaccel_type == 'auto':
        # Try NVENC first (most common), fall back to libx264
        # In real usage, auto will attempt to use what's available
        return 'h264_nvenc'
    else:
        return 'libx264'  # Software encoding fallback