Zordon/src/engines/engine_manager.py

import logging
import os
import shutil
import threading

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

logger = logging.getLogger()


class EngineManager:

    engines_path = None
    download_tasks = []

    @staticmethod
    def supported_engines():
        return [Blender, FFMPEG]

    @classmethod
    def engine_with_name(cls, engine_name):
        for obj in cls.supported_engines():
            if obj.name().lower() == engine_name.lower():
                return obj

    @classmethod
    def all_engines(cls):

        if not cls.engines_path:
            raise FileNotFoundError("Engines path must be set before requesting downloads")

        # Parse downloaded engine directory
        results = []
        try:
            all_items = os.listdir(cls.engines_path)
            all_directories = [item for item in all_items if os.path.isdir(os.path.join(cls.engines_path, item))]

            for directory in all_directories:
                # Split the input string by dashes to get segments
                segments = directory.split('-')

                # Create a dictionary with named keys
                keys = ["engine", "version", "system_os", "cpu"]
                result_dict = {keys[i]: segments[i] for i in range(min(len(keys), len(segments)))}
                result_dict['type'] = 'managed'

                # Figure out the binary name for the path
                binary_name = result_dict['engine'].lower()
                for eng in cls.supported_engines():
                    if eng.name().lower() == result_dict['engine']:
                        binary_name = eng.binary_names.get(result_dict['system_os'], binary_name)

                # Find path to binary
                path = None
                for root, _, files in os.walk(system_safe_path(os.path.join(cls.engines_path, directory))):
                    if binary_name in files:
                        path = os.path.join(root, binary_name)
                        break

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

        # add system installs to this list
        for eng in cls.supported_engines():
            if eng.default_renderer_path():
                results.append({'engine': eng.name(), 'version': eng().version(),
                                'system_os': current_system_os(),
                                'cpu': current_system_cpu(),
                                'path': eng.default_renderer_path(), 'type': 'system'})

        return results

    @classmethod
    def all_versions_for_engine(cls, engine):
        versions = [x for x in cls.all_engines() if x['engine'] == engine]
        sorted_versions = sorted(versions, key=lambda x: x['version'], reverse=True)
        return sorted_versions

    @classmethod
    def newest_engine_version(cls, engine, system_os=None, cpu=None):
        system_os = system_os or current_system_os()
        cpu = cpu or current_system_cpu()

        try:
            filtered = [x for x in cls.all_versions_for_engine(engine) 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}-{system_os}-{cpu}")
            return None

    @classmethod
    def is_version_downloaded(cls, engine, version, system_os=None, cpu=None):
        system_os = system_os or current_system_os()
        cpu = cpu or current_system_cpu()

        filtered = [x for x in cls.all_engines() if
                    x['engine'] == engine and x['system_os'] == system_os and x['cpu'] == cpu and
                    x['version'] == version]
        return filtered[0] if filtered else False

    @classmethod
    def version_is_available_to_download(cls, engine, version, system_os=None, cpu=None):
        try:
            downloader = cls.engine_with_name(engine).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

    @classmethod
    def find_most_recent_version(cls, engine=None, system_os=None, cpu=None, lts_only=False):
        try:
            downloader = cls.engine_with_name(engine).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 None

    @classmethod
    def get_existing_download_task(cls, engine, version, system_os=None, cpu=None):
        for task in cls.download_tasks:
            task_parts = task.name.split('-')
            task_engine, task_version, task_system_os, task_cpu = task_parts[:4]

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

    @classmethod
    def download_engine(cls, engine, version, system_os=None, cpu=None, background=False):

        engine_to_download = cls.engine_with_name(engine)
        existing_task = cls.get_existing_download_task(engine, version, system_os, cpu)
        if existing_task:
            logger.debug(f"Already downloading {engine} {version}")
            if not background:
                existing_task.join()  # If download task exists, wait until it's done downloading
            return
        elif not engine_to_download.downloader():
            logger.warning("No valid downloader for this engine. Please update this software manually.")
            return
        elif not cls.engines_path:
            raise FileNotFoundError("Engines path must be set before requesting downloads")

        thread = EngineDownloadWorker(engine, version, system_os, cpu)
        cls.download_tasks.append(thread)
        thread.start()

        if background:
            return thread
        else:
            thread.join()
            found_engine = cls.is_version_downloaded(engine, version, system_os, cpu)  # Check that engine downloaded
            if not found_engine:
                logger.error(f"Error downloading {engine}")
            return found_engine

    @classmethod
    def delete_engine_download(cls, engine, version, system_os=None, cpu=None):
        logger.info(f"Requested deletion of engine: {engine}-{version}")

        found = cls.is_version_downloaded(engine, version, system_os, cpu)
        if found and found['type'] == 'managed':  # don't delete system installs
            # find the root directory of the engine executable
            root_dir_name = '-'.join([engine, version, found['system_os'], found['cpu']])
            remove_path = os.path.join(found['path'].split(root_dir_name)[0], root_dir_name)
            # delete the file path
            logger.info(f"Deleting engine at path: {remove_path}")
            shutil.rmtree(remove_path, ignore_errors=False)
            logger.info(f"Engine {engine}-{version}-{found['system_os']}-{found['cpu']} successfully deleted")
            return True
        elif found:  # these are managed by the system / user. Don't delete these.
            logger.error(f'Cannot delete requested {engine} {version}. Managed externally.')
        else:
            logger.error(f"Cannot find engine: {engine}-{version}")
        return False

    @classmethod
    def update_all_engines(cls):
        def engine_update_task(engine_class):
            logger.debug(f"Checking for updates to {engine_class.name()}")
            latest_version = engine_class.downloader().find_most_recent_version()
            if latest_version:
                logger.debug(f"Latest version of {engine_class.name()} available: {latest_version.get('version')}")
                if not cls.is_version_downloaded(engine_class.name(), latest_version.get('version')):
                    logger.info(f"Downloading latest version of {engine_class.name()}...")
                    cls.download_engine(engine=engine_class.name(), version=latest_version['version'], background=True)
            else:
                logger.warning(f"Unable to get check for updates for {engine.name()}")

        logger.info(f"Checking for updates for render engines...")
        threads = []
        for engine in cls.supported_engines():
            if engine.downloader():
                thread = threading.Thread(target=engine_update_task, args=(engine,))
                threads.append(thread)
                thread.start()

    @classmethod
    def create_worker(cls, renderer, input_path, output_path, engine_version=None, args=None, parent=None, name=None):

        worker_class = cls.engine_with_name(renderer).worker_class()

        # check to make sure we have versions installed
        all_versions = cls.all_versions_for_engine(renderer)
        if not all_versions:
            raise FileNotFoundError(f"Cannot find any installed {renderer} engines")

        # Find the path to the requested engine version or use default
        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

            # Download the required engine if not found locally
            if not engine_path:
                download_result = cls.download_engine(renderer, engine_version)
                if not download_result:
                    raise FileNotFoundError(f"Cannot download requested version: {renderer} {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=input_path, output_path=output_path, engine_path=engine_path, args=args,
                            parent=parent, name=name)

    @classmethod
    def engine_for_project_path(cls, path):
        _, extension = os.path.splitext(path)
        extension = extension.lower().strip('.')
        for engine in cls.supported_engines():
            if extension in engine.supported_extensions():
                return engine
        undefined_renderer_support = [x for x in cls.supported_engines() if not x.supported_extensions()]
        return undefined_renderer_support[0]


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

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

        # Get the appropriate downloader class based on the engine type
        EngineManager.engine_with_name(self.engine).downloader().download_engine(
            self.version, download_location=EngineManager.engines_path, system_os=self.system_os, cpu=self.cpu,
            timeout=300)

        # remove itself from the downloader list
        EngineManager.download_tasks.remove(self)


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.all_engines())