Zordon/src/engines/engine_manager.py

import logging
import os
import platform
import shutil

from src.engines.blender.blender_downloader import BlenderDownloader
from src.engines.blender.blender_engine import Blender
from src.engines.ffmpeg.ffmpeg_downloader import FFMPEGDownloader
from src.engines.ffmpeg.ffmpeg_engine import FFMPEG
from src.utilities.misc_helper import system_safe_path

logger = logging.getLogger()


class EngineManager:

    engines_path = "~/zordon-uploads/engines"
    downloader_classes = {
        "blender": BlenderDownloader,
        "ffmpeg": FFMPEGDownloader,
        # Add more engine types and corresponding downloader classes as needed
    }

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

    @classmethod
    def all_engines(cls):
        results = []
        # Parse downloaded engine directory
        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:
            logger.warning("Cannot find local engines download directory")

        # 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': cls.system_os(),
                                'cpu': cls.system_cpu(),
                                'path': eng.default_renderer_path(), 'type': 'system'})

        return results

    @classmethod
    def all_versions_for_engine(cls, engine):
        return [x for x in cls.all_engines() if x['engine'] == engine]

    @classmethod
    def newest_engine_version(cls, engine, system_os=None, cpu=None):
        system_os = system_os or cls.system_os()
        cpu = cpu or cls.system_cpu()

        try:
            filtered = [x for x in cls.all_engines() if x['engine'] == engine and x['system_os'] == system_os and x['cpu'] == cpu]
            versions = sorted(filtered, key=lambda x: x['version'], reverse=True)
            return versions[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 cls.system_os()
        cpu = cpu or cls.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

    @staticmethod
    def system_os():
        return platform.system().lower().replace('darwin', 'macos')

    @staticmethod
    def system_cpu():
        return platform.machine().lower().replace('amd64', 'x64')

    @classmethod
    def version_is_available_to_download(cls, engine, version, system_os=None, cpu=None):
        try:
            return cls.downloader_classes[engine].version_is_available_to_download(version=version, system_os=system_os,
                                                                                   cpu=cpu)
        except Exception as e:
            return None

    @classmethod
    def find_most_recent_version(cls, engine=None, system_os=None, cpu=None, lts_only=False):
        try:
            return cls.downloader_classes[engine].find_most_recent_version(system_os=system_os, cpu=cpu)
        except Exception as e:
            return None

    @classmethod
    def download_engine(cls, engine, version, system_os=None, cpu=None):
        existing_download = cls.is_version_downloaded(engine, version, system_os, cpu)
        if existing_download:
            logger.info(f"Requested download of {engine} {version}, but local copy already exists")
            return existing_download

        # Check if the provided engine type is valid
        if engine not in cls.downloader_classes:
            logger.error("No valid engine found")
            return

        # Get the appropriate downloader class based on the engine type
        cls.downloader_classes[engine].download_engine(version, download_location=cls.engines_path,
                                                       system_os=system_os, cpu=cpu, timeout=300)

        # Check that engine was properly downloaded
        found_engine = cls.is_version_downloaded(engine, version, system_os, cpu)
        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:
            dir_path = os.path.dirname(found['path'])
            shutil.rmtree(dir_path, ignore_errors=True)
            logger.info(f"Engine {engine}-{version}-{found['system_os']}-{found['cpu']} successfully deleted")
            return True
        else:
            logger.error(f"Cannot find engine: {engine}-{version}")

    @classmethod
    def update_all_engines(cls):
        logger.info(f"Checking for updates for render engines...")
        for engine, engine_downloader in cls.downloader_classes.items():
            latest_version = engine_downloader.find_most_recent_version().get('version')
            if latest_version and not cls.is_version_downloaded(engine, latest_version):
                logger.info(f"Downloading newest version of {engine} ({latest_version})")
                cls.download_engine(engine, latest_version)


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')