Zordon/src/engines/engine_manager.py

import logging
import os
import shutil
import threading
import concurrent.futures
from pathlib import Path
from typing import Type, List, Dict, Any, Optional

from src.engines.core.base_engine import BaseRenderEngine
from src.engines.blender.blender_engine import Blender
from src.engines.ffmpeg.ffmpeg_engine import FFMPEG
from src.utilities.misc_helper import current_system_os, current_system_cpu

logger = logging.getLogger()

ENGINE_CLASSES = [Blender, FFMPEG]

class EngineManager:
    """Class that manages different versions of installed render engines and handles fetching and downloading new versions,
    if possible.
    """

    _default_instance: Optional['EngineManager'] = None

    engines_path: Optional[str] = None
    download_tasks: List[Any] = []

    def __init__(self) -> None:
        self.engines_path: Optional[str] = None
        self.download_tasks: List[Any] = []

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

    @staticmethod
    def supported_engines() -> list[type[BaseRenderEngine]]:
        return ENGINE_CLASSES

    # --- Installed Engines ---

    def _engine_class_for_project_path(self, path: str) -> Type[BaseRenderEngine]:
        _, extension = os.path.splitext(path)
        extension = extension.lower().strip('.')
        for engine_class in self.supported_engines():
            engine = self.get_latest_engine_instance(engine_class)
            if extension in engine.supported_extensions():
                return engine_class
        undefined_renderer_support = [x for x in self.supported_engines() if not self.get_latest_engine_instance(x).supported_extensions()]
        return undefined_renderer_support[0]

    def _engine_class_with_name(self, engine_name: str) -> Optional[Type[BaseRenderEngine]]:
        for obj in self.supported_engines():
            if obj.name().lower() == engine_name.lower():
                return obj
        return None

    def _get_latest_engine_instance(self, engine_class: Type[BaseRenderEngine]) -> BaseRenderEngine:
        newest = self.newest_installed_engine_data(engine_class.name())
        engine = engine_class(newest["path"])
        return engine

    def _get_installed_engine_data(self, filter_name: Optional[str] = None, include_corrupt: bool = False,
                                  ignore_system: bool = False) -> List[Dict[str, Any]]:
        if not self.engines_path:
            raise FileNotFoundError("Engine path is not set")

        results = []
        try:
            all_items = os.listdir(self.engines_path)
            all_directories = [item for item in all_items if os.path.isdir(os.path.join(self.engines_path, item))]
            keys = ["engine", "version", "system_os", "cpu"]

            for directory in all_directories:
                segments = directory.split('-')
                result_dict = {keys[i]: segments[i] for i in range(min(len(keys), len(segments)))}
                result_dict['type'] = 'managed'

                binary_name = result_dict['engine'].lower()
                eng = self.engine_class_with_name(result_dict['engine'])
                binary_name = eng.binary_names.get(result_dict['system_os'], binary_name)

                search_root = self.engines_path / directory
                match = next((p for p in search_root.rglob(binary_name) if p.is_file()), None)
                path = str(match) if match else None
                result_dict['path'] = path

                if not filter_name or filter_name == result_dict['engine']:
                    results.append(result_dict)
        except FileNotFoundError as e:
            logger.warning(f"Cannot find local engines download directory: {e}")

        def fetch_engine_details(eng, include_corrupt=False):
            version = eng().version()
            if not version and not include_corrupt:
                return

            return {
                'engine': eng.name(),
                'version': version or 'error',
                'system_os': current_system_os(),
                'cpu': current_system_cpu(),
                'path': eng.default_engine_path(),
                'type': 'system'
            }

        if not ignore_system:
            with concurrent.futures.ThreadPoolExecutor() as executor:
                futures = {
                    executor.submit(fetch_engine_details, eng, include_corrupt): eng.name()
                    for eng in self.supported_engines()
                    if eng.default_engine_path() and (not filter_name or filter_name == eng.name())
                }

                for future in concurrent.futures.as_completed(futures):
                    result = future.result()
                    if result:
                        results.append(result)

        return results

    # --- Check for Updates ---

    def _update_all_engines(self) -> None:
        for engine in self.downloadable_engines():
            update_available = self.is_engine_update_available(engine)
            if update_available:
                update_available['name'] = engine.name()
                self.download_engine(engine.name(), update_available['version'], background=True)

    def _all_version_data_for_engine(self, engine_name: str, include_corrupt=False, ignore_system=False) -> list:
        versions = self.get_installed_engine_data(filter_name=engine_name, include_corrupt=include_corrupt, ignore_system=ignore_system)
        sorted_versions = sorted(versions, key=lambda x: x['version'], reverse=True)
        return sorted_versions

    def _newest_installed_engine_data(self, engine_name: str, system_os=None, cpu=None, ignore_system=None) -> list:
        system_os = system_os or current_system_os()
        cpu = cpu or current_system_cpu()

        try:
            filtered = [x for x in self.all_version_data_for_engine(engine_name, ignore_system=ignore_system)
                        if x['system_os'] == system_os and x['cpu'] == cpu]
            return filtered[0]
        except IndexError:
            logger.error(f"Cannot find newest engine version for {engine_name}-{system_os}-{cpu}")
        return []

    def _is_version_installed(self, engine_name: str, version: str, system_os=None, cpu=None, ignore_system=False):
        system_os = system_os or current_system_os()
        cpu = cpu or current_system_cpu()

        filtered = [x for x in self.get_installed_engine_data(filter_name=engine_name, ignore_system=ignore_system) if
                    x['system_os'] == system_os and x['cpu'] == cpu and x['version'] == version]
        return filtered[0] if filtered else False

    def _version_is_available_to_download(self, engine_name: str, version, system_os=None, cpu=None):
        try:
            downloader = self.engine_class_with_name(engine_name).downloader()
            return downloader.version_is_available_to_download(version=version, system_os=system_os, cpu=cpu)
        except Exception as e:
            logger.debug(f"Exception in version_is_available_to_download: {e}")
            return None

    def _find_most_recent_version(self, engine_name: str, system_os=None, cpu=None, lts_only=False) -> dict:
        try:
            downloader = self.engine_class_with_name(engine_name).downloader()
            return downloader.find_most_recent_version(system_os=system_os, cpu=cpu)
        except Exception as e:
            logger.debug(f"Exception in find_most_recent_version: {e}")
            return {}

    def _is_engine_update_available(self, engine_class: Type[BaseRenderEngine], ignore_system_installs=False):
        logger.debug(f"Checking for updates to {engine_class.name()}")
        latest_version = engine_class.downloader().find_most_recent_version()

        if not latest_version:
            logger.warning(f"Could not find most recent version of {engine_class.name()} to download")
            return None

        version_num = latest_version.get('version')
        if self.is_version_installed(engine_class.name(), version_num, ignore_system=ignore_system_installs):
            logger.debug(f"Latest version of {engine_class.name()} ({version_num}) already downloaded")
            return None

        return latest_version

    # --- Downloads ---

    def _downloadable_engines(self):
        return [engine for engine in self.supported_engines() if hasattr(engine, "downloader") and engine.downloader()]

    def _get_existing_download_task(self, engine_name, version, system_os=None, cpu=None):
        for task in self.download_tasks:
            task_parts = task.name.split('-')
            task_engine, task_version, task_system_os, task_cpu = task_parts[:4]

            if engine_name == task_engine and version == task_version:
                if system_os in (task_system_os, None) and cpu in (task_cpu, None):
                    return task
        return None

    def _download_engine(self, engine_name, version, system_os=None, cpu=None, background=False, ignore_system=False):
        engine_to_download = self.engine_class_with_name(engine_name)
        existing_task = self.get_existing_download_task(engine_name, version, system_os, cpu)
        if existing_task:
            logger.debug(f"Already downloading {engine_name} {version}")
            if not background:
                existing_task.join()
            return None
        elif not engine_to_download.downloader():
            logger.warning("No valid downloader for this engine. Please update this software manually.")
            return None
        elif not self.engines_path:
            raise FileNotFoundError("Engines path must be set before requesting downloads")

        thread = EngineDownloadWorker(engine_name, version, system_os, cpu)
        self.download_tasks.append(thread)
        thread.start()

        if background:
            return thread

        thread.join()
        found_engine = self.is_version_installed(engine_name, version, system_os, cpu, ignore_system)
        if not found_engine:
            logger.error(f"Error downloading {engine_name}")
        return found_engine

    def _delete_engine_download(self, engine_name, version, system_os=None, cpu=None):
        logger.info(f"Requested deletion of engine: {engine_name}-{version}")

        found = self.is_version_installed(engine_name, version, system_os, cpu)
        if found and found['type'] == 'managed':
            root_dir_name = '-'.join([engine_name, version, found['system_os'], found['cpu']])
            remove_path = os.path.join(found['path'].split(root_dir_name)[0], root_dir_name)
            logger.info(f"Deleting engine at path: {remove_path}")
            shutil.rmtree(remove_path, ignore_errors=False)
            logger.info(f"Engine {engine_name}-{version}-{found['system_os']}-{found['cpu']} successfully deleted")
            return True
        elif found:
            logger.error(f'Cannot delete requested {engine_name} {version}. Managed externally.')
        else:
            logger.error(f"Cannot find engine: {engine_name}-{version}")
        return False

    # --- Background Tasks ---

    def _active_downloads(self) -> list:
        return [x for x in self.download_tasks if x.is_alive()]

    def _create_worker(self, engine_name: str, input_path: Path, output_path: Path, engine_version=None, args=None, parent=None, name=None):
        worker_class = self.engine_class_with_name(engine_name).worker_class()

        all_versions = self.all_version_data_for_engine(engine_name)
        if not all_versions:
            raise FileNotFoundError(f"Cannot find any installed '{engine_name}' engines")

        engine_path = None
        if engine_version and engine_version != 'latest':
            for ver in all_versions:
                if ver['version'] == engine_version:
                    engine_path = ver['path']
                    break

            if not engine_path:
                download_result = self.download_engine(engine_name, engine_version)
                if not download_result:
                    raise FileNotFoundError(f"Cannot download requested version: {engine_name} {engine_version}")
                engine_path = download_result['path']
                logger.info("Engine downloaded. Creating worker.")
        else:
            logger.debug(f"Using latest engine version ({all_versions[0]['version']})")
            engine_path = all_versions[0]['path']

        if not engine_path:
            raise FileNotFoundError(f"Cannot find requested engine version {engine_version}")

        return worker_class(input_path=str(input_path), output_path=str(output_path), engine_path=engine_path, args=args,
                            parent=parent, name=name)

    # --- Forwarders for backward compatibility ---

    @classmethod
    def engine_class_for_project_path(cls, path):
        if cls._default_instance is not None:
            return cls._default_instance._engine_class_for_project_path(path)

    @classmethod
    def engine_class_with_name(cls, engine_name):
        if cls._default_instance is not None:
            return cls._default_instance._engine_class_with_name(engine_name)

    @classmethod
    def get_latest_engine_instance(cls, engine_class):
        if cls._default_instance is not None:
            return cls._default_instance._get_latest_engine_instance(engine_class)

    @classmethod
    def get_installed_engine_data(cls, filter_name=None, include_corrupt=False, ignore_system=False):
        if cls._default_instance is not None:
            return cls._default_instance._get_installed_engine_data(filter_name, include_corrupt, ignore_system)
        return []

    @classmethod
    def update_all_engines(cls):
        if cls._default_instance is not None:
            cls._default_instance._update_all_engines()

    @classmethod
    def all_version_data_for_engine(cls, engine_name, include_corrupt=False, ignore_system=False):
        if cls._default_instance is not None:
            return cls._default_instance._all_version_data_for_engine(engine_name, include_corrupt, ignore_system)
        return []

    @classmethod
    def newest_installed_engine_data(cls, engine_name, system_os=None, cpu=None, ignore_system=None):
        if cls._default_instance is not None:
            return cls._default_instance._newest_installed_engine_data(engine_name, system_os, cpu, ignore_system)
        return []

    @classmethod
    def is_version_installed(cls, engine_name, version, system_os=None, cpu=None, ignore_system=False):
        if cls._default_instance is not None:
            return cls._default_instance._is_version_installed(engine_name, version, system_os, cpu, ignore_system)
        return False

    @classmethod
    def version_is_available_to_download(cls, engine_name, version, system_os=None, cpu=None):
        if cls._default_instance is not None:
            return cls._default_instance._version_is_available_to_download(engine_name, version, system_os, cpu)
        return None

    @classmethod
    def find_most_recent_version(cls, engine_name, system_os=None, cpu=None, lts_only=False):
        if cls._default_instance is not None:
            return cls._default_instance._find_most_recent_version(engine_name, system_os, cpu, lts_only)
        return {}

    @classmethod
    def is_engine_update_available(cls, engine_class, ignore_system_installs=False):
        if cls._default_instance is not None:
            return cls._default_instance._is_engine_update_available(engine_class, ignore_system_installs)
        return None

    @classmethod
    def downloadable_engines(cls):
        if cls._default_instance is not None:
            return cls._default_instance._downloadable_engines()
        return []

    @classmethod
    def get_existing_download_task(cls, engine_name, version, system_os=None, cpu=None):
        if cls._default_instance is not None:
            return cls._default_instance._get_existing_download_task(engine_name, version, system_os, cpu)
        return None

    @classmethod
    def download_engine(cls, engine_name, version, system_os=None, cpu=None, background=False, ignore_system=False):
        if cls._default_instance is not None:
            return cls._default_instance._download_engine(engine_name, version, system_os, cpu, background, ignore_system)
        return None

    @classmethod
    def delete_engine_download(cls, engine_name, version, system_os=None, cpu=None):
        if cls._default_instance is not None:
            return cls._default_instance._delete_engine_download(engine_name, version, system_os, cpu)
        return False

    @classmethod
    def active_downloads(cls):
        if cls._default_instance is not None:
            return cls._default_instance._active_downloads()
        return []

    @classmethod
    def create_worker(cls, engine_name, input_path, output_path, engine_version=None, args=None, parent=None, name=None):
        if cls._default_instance is not None:
            return cls._default_instance._create_worker(engine_name, input_path, output_path, engine_version, args, parent, name)
        raise RuntimeError("EngineManager is not initialized")


class EngineDownloadWorker(threading.Thread):
    def __init__(self, engine, version, system_os=None, cpu=None):
        super().__init__()
        self.engine = engine
        self.version = version
        self.system_os = system_os
        self.cpu = cpu
        self.percent_complete = 0

    def _update_progress(self, current_progress):
        self.percent_complete = current_progress

    def run(self):
        try:
            existing_download = EngineManager.is_version_installed(self.engine, self.version, self.system_os, self.cpu,
                                                                    ignore_system=True)
            if existing_download:
                logger.info(f"Requested download of {self.engine} {self.version}, but local copy already exists")
                return existing_download

            downloader = EngineManager.engine_class_with_name(self.engine).downloader()
            downloader.download_engine(self.version, download_location=EngineManager.engines_path,
                                        system_os=self.system_os, cpu=self.cpu, timeout=300, progress_callback=self._update_progress)
        except Exception as e:
            logger.error(f"Error in download worker: {e}")
        finally:
            try:
                if EngineManager._default_instance is not None:
                    EngineManager._default_instance.download_tasks.remove(self)
            except ValueError:
                pass


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    # print(EngineManager.newest_engine_version('blender', 'macos', 'arm64'))
    # EngineManager.delete_engine_download('blender', '3.2.1', 'macos', 'a')
    EngineManager.engines_path = "/Users/brettwilliams/zordon-uploads/engines"
    # print(EngineManager.is_version_downloaded("ffmpeg", "6.0"))
    print(EngineManager.get_installed_engine_data())