Zordon/src/api/job_import_handler.py

#!/usr/bin/env python3
import logging
import os
import shutil
import tempfile
import zipfile
from datetime import datetime
from pathlib import Path

import requests
from tqdm import tqdm
from werkzeug.utils import secure_filename

from src.distributed_job_manager import DistributedJobManager

logger = logging.getLogger()


class JobImportHandler:
    """Handles job import operations for rendering projects.

    This class provides functionality to validate, download, and process
    job data and project files for the rendering queue system.
    """

    @classmethod
    def create_jobs_from_processed_data(cls, processed_job_data: dict) -> list[dict]:
        """ Takes processed job data and creates new jobs

        Args: processed_job_data: Dictionary containing job information"""
        loaded_project_local_path = processed_job_data['__loaded_project_local_path']

        # prepare child job data
        job_data_to_create = []
        if processed_job_data.get("child_jobs"):
            for child_job_diffs in processed_job_data["child_jobs"]:
                processed_child_job_data = processed_job_data.copy()
                processed_child_job_data.pop("child_jobs")
                processed_child_job_data.update(child_job_diffs)
                processed_child_job_data['__use_output_subdir'] = True
                job_data_to_create.append(processed_child_job_data)
        else:
            job_data_to_create.append(processed_job_data)

        # create the jobs
        created_jobs = []
        for job_data in job_data_to_create:
            new_job = DistributedJobManager.create_render_job(job_data, loaded_project_local_path)
            created_jobs.append(new_job)

        # Save notes to .txt
        if processed_job_data.get("notes"):
            parent_dir = Path(loaded_project_local_path).parent.parent
            notes_name = processed_job_data['name'] + "-notes.txt"
            with (Path(parent_dir) / notes_name).open("w") as f:
                f.write(processed_job_data["notes"])
        return [x.json() for x in created_jobs]

    @classmethod
    def validate_job_data(cls, new_job_data: dict, upload_directory: Path, uploaded_file=None) -> dict:
        """Validates and prepares job data for import.

        This method validates the job data dictionary, handles project file
        acquisition (upload, download, or local copy), and prepares the job
        directory structure.

        Args:
            new_job_data: Dictionary containing job information including
                'name', 'engine_name', and optionally 'url' or 'local_path'.
            upload_directory: Base directory for storing uploaded jobs.
            uploaded_file: Optional uploaded file object from the request.

        Returns:
            The validated job data dictionary with additional metadata.

        Raises:
            KeyError: If required fields 'name' or 'engine_name' are missing.
            FileNotFoundError: If no valid project file can be found.
        """
        loaded_project_local_path = None

        # check for required keys
        job_name = new_job_data.get('name')
        engine_name = new_job_data.get('engine_name')
        if not job_name:
            raise KeyError("Missing job name")
        if not engine_name:
            raise KeyError("Missing engine name")

        project_url = new_job_data.get('url', None)
        local_path = new_job_data.get('local_path', None)
        downloaded_file_url = None

        if uploaded_file and uploaded_file.filename:
            referred_name = os.path.basename(uploaded_file.filename)
        elif project_url:
            referred_name, downloaded_file_url = cls.download_project_from_url(project_url)
            if not referred_name:
                raise FileNotFoundError(f"Error downloading file from URL: {project_url}")
        elif local_path and os.path.exists(local_path):
            referred_name = os.path.basename(local_path)
        else:
            raise FileNotFoundError("Cannot find any valid project paths")

        # Prepare the local filepath
        cleaned_path_name = job_name.replace(' ', '-')
        folder_name = f"{cleaned_path_name}-{engine_name}-{datetime.now().strftime('%Y.%m.%d_%H.%M.%S')}"
        job_dir = Path(upload_directory) / folder_name
        os.makedirs(job_dir, exist_ok=True)
        project_source_dir = Path(job_dir) / 'source'
        os.makedirs(project_source_dir, exist_ok=True)

        # Move projects to their work directories
        if uploaded_file and uploaded_file.filename:
            # Handle file uploading
            filename = secure_filename(uploaded_file.filename)
            loaded_project_local_path = Path(project_source_dir) / filename
            uploaded_file.save(str(loaded_project_local_path))
            logger.info(f"Transfer complete for {loaded_project_local_path.relative_to(upload_directory)}")

        elif project_url:
            # Handle downloading project from a URL
            loaded_project_local_path = Path(project_source_dir) / referred_name
            shutil.move(str(downloaded_file_url), str(loaded_project_local_path))
            logger.info(f"Download complete for {loaded_project_local_path.relative_to(upload_directory)}")

        elif local_path:
            # Handle local files
            loaded_project_local_path = Path(project_source_dir) / referred_name
            shutil.copy(str(local_path), str(loaded_project_local_path))
            logger.info(f"Import complete for {loaded_project_local_path.relative_to(upload_directory)}")

        if loaded_project_local_path.suffix == ".zip":
            loaded_project_local_path = cls.process_zipped_project(loaded_project_local_path)

        new_job_data["__loaded_project_local_path"] = loaded_project_local_path

        return new_job_data

    @staticmethod
    def download_project_from_url(project_url: str):
        """Downloads a project file from the given URL.

        Downloads the file from the specified URL to a temporary directory
        with progress tracking. Returns the filename and temporary path.

        Args:
            project_url: The URL to download the project file from.

        Returns:
            A tuple of (filename, temp_file_path) if successful,
            (None, None) if download fails.
        """
        # This nested function is to handle downloading from a URL
        logger.info(f"Downloading project from url: {project_url}")
        referred_name = os.path.basename(project_url)

        try:
            response = requests.get(project_url, stream=True, timeout=300)
            if response.status_code == 200:
                # Get the total file size from the "Content-Length" header
                file_size = int(response.headers.get("Content-Length", 0))
                # Create a progress bar using tqdm
                progress_bar = tqdm(total=file_size, unit="B", unit_scale=True)
                # Open a file for writing in binary mode
                downloaded_file_url = os.path.join(tempfile.gettempdir(), referred_name)
                with open(downloaded_file_url, "wb") as file:
                    for chunk in response.iter_content(chunk_size=1024):
                        if chunk:
                            # Write the chunk to the file
                            file.write(chunk)
                            # Update the progress bar
                            progress_bar.update(len(chunk))
                # Close the progress bar
                progress_bar.close()
                return referred_name, downloaded_file_url
        except Exception as e:
            logger.error(f"Error downloading file: {e}")
        return None, None

    @staticmethod
    def process_zipped_project(zip_path: Path) -> Path:
        """
        Processes a zipped project.

        This method takes a path to a zip file, extracts its contents, and returns the path to the extracted project
        file. If the zip file contains more than one project file or none, an error is raised.

        Args:
            zip_path (Path): The path to the zip file.

        Raises:
            ValueError: If there's more than 1 project file or none in the zip file.

        Returns:
            Path: The path to the main project file.
        """
        work_path = zip_path.parent

        try:
            with zipfile.ZipFile(zip_path, 'r') as myzip:
                myzip.extractall(str(work_path))

            project_files = [p for p in work_path.iterdir() if p.is_file() and p.suffix.lower() != ".zip"]

            logger.debug(f"Zip files: {project_files}")

            # supported_exts = RenderWorkerFactory.class_for_name(engine).engine.supported_extensions
            # if supported_exts:
            #     project_files = [file for file in project_files if any(file.endswith(ext) for ext in supported_exts)]

            # If there's more than 1 project file or none, raise an error
            if len(project_files) != 1:
                raise ValueError(f'Cannot find a valid project file in {zip_path.name}')

            extracted_project_path = work_path / project_files[0]
            logger.info(f"Extracted zip file to {extracted_project_path}")

        except (zipfile.BadZipFile, zipfile.LargeZipFile) as e:
            logger.error(f"Error processing zip file: {e}")
            raise ValueError(f'Error processing zip file: {e}') from e
        return extracted_project_path