Zordon/src/api/preview_manager.py

import logging
import os
import subprocess
import threading
from pathlib import Path
from typing import Dict, Optional

from src.utilities.ffmpeg_helper import generate_thumbnail, save_first_frame

logger = logging.getLogger()
supported_video_formats = ['.mp4', '.mov', '.avi', '.mpg', '.mpeg', '.mxf', '.m4v', '.mkv', '.webm']
supported_image_formats = ['.jpg', '.png', '.exr', '.tif', '.tga', '.bmp', '.webp']


class PreviewManager:
    _default_instance: Optional['PreviewManager'] = None

    storage_path: Optional[str] = None

    def __init__(self) -> None:
        self.storage_path = None
        self._running_jobs: Dict = {}

    @classmethod
    def _sync_class(cls) -> None:
        if cls._default_instance is not None:
            cls.storage_path = cls._default_instance.storage_path

    def _generate_job_preview_worker(self, job, replace_existing=False, max_width=480):

        # Determine best source file to use for thumbs
        job_file_list = job.file_list()
        source_files = job_file_list if job_file_list else [job.input_path]
        preview_label = "output" if job_file_list else "input"

        # filter by type
        found_image_files = [f for f in source_files if os.path.splitext(f)[-1].lower() in supported_image_formats]
        found_video_files = [f for f in source_files if os.path.splitext(f)[-1].lower() in supported_video_formats]

        # check if we even have any valid files to work from
        if source_files and not found_video_files and not found_image_files:
            logger.warning(f"No valid image or video files found in files from job: {job}")
            return

        os.makedirs(self.storage_path, exist_ok=True)
        base_path = os.path.join(self.storage_path, f"{job.id}-{preview_label}-{max_width}")
        preview_video_path = base_path + '.mp4'
        preview_image_path = base_path + '.jpg'

        if replace_existing:
            for x in [preview_image_path, preview_video_path]:
                try:
                    os.remove(x)
                except OSError:
                    pass

        # Generate image previews
        if (found_video_files or found_image_files) and not os.path.exists(preview_image_path):
            try:
                path_of_source = found_image_files[-1] if found_image_files else found_video_files[-1]
                logger.debug(f"Generating image preview for {path_of_source}")
                save_first_frame(source_path=path_of_source, dest_path=preview_image_path, max_width=max_width)
                logger.debug(f"Successfully created image preview for {path_of_source}")
            except Exception as e:
                logger.error(f"Error generating image preview for {job}: {e}")

        # Generate video previews
        if found_video_files and not os.path.exists(preview_video_path):
            try:
                path_of_source = found_video_files[0]
                logger.debug(f"Generating video preview for {path_of_source}")
                generate_thumbnail(source_path=path_of_source, dest_path=preview_video_path, max_width=max_width)
                logger.debug(f"Successfully created video preview for {path_of_source}")
            except subprocess.CalledProcessError as e:
                logger.error(f"Error generating video preview for {job}: {e}")

    def _update_previews_for_job(self, job, replace_existing=False, wait_until_completion=False, timeout=None):
        job_thread = self._running_jobs.get(job.id)
        if job_thread and job_thread.is_alive():
            logger.debug(f'Preview generation job already running for {job}')
        else:
            job_thread = threading.Thread(target=self._generate_job_preview_worker, args=(job, replace_existing,))
            job_thread.start()
            self._running_jobs[job.id] = job_thread

        if wait_until_completion:
            job_thread.join(timeout=timeout)

    def _get_previews_for_job(self, job):

        results = {}
        try:
            directory_path = Path(self.storage_path)
            preview_files_for_job = [f for f in directory_path.iterdir() if f.is_file() and f.name.startswith(job.id)]

            for preview_filename in preview_files_for_job:
                try:
                    pixel_width = str(preview_filename).split('-')[-1]
                    preview_label = str(os.path.basename(preview_filename)).split('-')[1]
                    extension = os.path.splitext(preview_filename)[-1].lower()
                    kind = 'video' if extension in supported_video_formats else \
                        'image' if extension in supported_image_formats else 'unknown'
                    results[preview_label] = results.get(preview_label, [])
                    results[preview_label].append({'filename': str(preview_filename), 'width': pixel_width, 'kind': kind})
                except IndexError:  # ignore invalid filenames
                    pass
        except FileNotFoundError:
            pass
        return results

    def _delete_previews_for_job(self, job):
        all_previews = self.get_previews_for_job(job)
        flattened_list = [item for sublist in all_previews.values() for item in sublist]
        for preview in flattened_list:
            try:
                logger.debug(f"Removing preview: {preview['filename']}")
                os.remove(preview['filename'])
            except OSError as e:
                logger.error(f"Error removing preview '{preview.get('filename')}': {e}")

    # --- Forwarders for backward compatibility ---

    @classmethod
    def update_previews_for_job(cls, job, replace_existing=False, wait_until_completion=False, timeout=None):
        if cls._default_instance is not None:
            cls._default_instance._update_previews_for_job(job, replace_existing, wait_until_completion, timeout)

    @classmethod
    def get_previews_for_job(cls, job):
        if cls._default_instance is not None:
            return cls._default_instance._get_previews_for_job(job)
        return {}

    @classmethod
    def delete_previews_for_job(cls, job):
        if cls._default_instance is not None:
            cls._default_instance._delete_previews_for_job(job)