Zordon/addons/blender/zordon_blender/__init__.py

bl_info = {
    'name': 'Zordon Render Submitter',
    'author': 'Zordon',
    'version': (0, 1, 0),
    'blender': (3, 6, 0),
    'location': 'Properties > Render > Zordon',
    'description': 'Submit the current Blender file to a Zordon render server.',
    'category': 'Render',
}

import getpass
import http.client
import json
import os
import re
import socket
import subprocess
import time
from urllib.parse import urlparse
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

import bpy
from bpy.props import BoolProperty, EnumProperty, IntProperty, StringProperty
from bpy.types import AddonPreferences, Operator, Panel, PropertyGroup


DEFAULT_SERVER = 'localhost:8080'
DEFAULT_TIMEOUT = 10
UPLOAD_TIMEOUT = 300
ZORDON_SERVICE_TYPE = '_zordon._tcp.local.'


def _addon_preferences(context):
    return context.preferences.addons[__name__].preferences


def _configured_server_items(self, context):
    preferences = _addon_preferences(context)
    raw_servers = [server.strip() for server in preferences.servers.split(',')]
    servers = [server for server in raw_servers if server]
    if not servers:
        servers = [DEFAULT_SERVER]
    return [(server, server, '', index) for index, server in enumerate(servers)]


def _normalize_server_url(server, fallback_port):
    if not server:
        server = DEFAULT_SERVER

    server = server.strip()
    if not server.startswith(('http://', 'https://')):
        server = f'http://{server}'

    parsed = urlparse(server)
    hostname = parsed.hostname or 'localhost'
    scheme = parsed.scheme or 'http'
    port = parsed.port or fallback_port or 8080
    return scheme, hostname, int(port)


def _api_url(server, port, path):
    scheme, hostname, parsed_port = _normalize_server_url(server, port)
    return f'{scheme}://{hostname}:{parsed_port}{path}'


def _get_text(server, port, path, timeout=DEFAULT_TIMEOUT):
    request = Request(_api_url(server, port, path), method='GET')
    with urlopen(request, timeout=timeout) as response:
        return response.status, response.read().decode('utf-8')


def _server_list_from_preferences(preferences):
    return [server.strip() for server in preferences.servers.split(',') if server.strip()]


def _save_server_list(preferences, servers):
    unique_servers = []
    for server in servers:
        if server and server not in unique_servers:
            unique_servers.append(server)
    preferences.servers = ', '.join(unique_servers)


def _discover_servers_with_zeroconf(timeout=3.0):
    try:
        from zeroconf import ServiceBrowser, Zeroconf
    except ImportError:
        return []

    discovered = []
    zeroconf = Zeroconf()

    class Listener:
        def add_service(self, zc, service_type, name):
            info = zc.get_service_info(service_type, name, timeout=1000)
            if info and info.server and info.port:
                discovered.append(f'{info.server.rstrip(".")}:{info.port}')

        def update_service(self, zc, service_type, name):
            self.add_service(zc, service_type, name)

        def remove_service(self, zc, service_type, name):
            pass

    try:
        ServiceBrowser(zeroconf, ZORDON_SERVICE_TYPE, Listener())
        time.sleep(timeout)
    finally:
        zeroconf.close()

    return discovered


def _discover_service_names_with_dns_sd(timeout=2.0):
    try:
        process = subprocess.Popen(
            ['dns-sd', '-B', '_zordon', '_tcp', 'local'],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
        )
    except (FileNotFoundError, OSError):
        return []

    try:
        stdout, _ = process.communicate(timeout=timeout)
    except subprocess.TimeoutExpired:
        process.terminate()
        stdout, _ = process.communicate(timeout=1)

    service_names = []
    for line in stdout.splitlines():
        if ' Add ' not in line:
            continue
        parts = line.split()
        if len(parts) >= 7:
            service_name = ' '.join(parts[6:])
            if service_name and service_name not in service_names:
                service_names.append(service_name)
    return service_names


def _resolve_service_with_dns_sd(service_name, timeout=2.0):
    try:
        process = subprocess.Popen(
            ['dns-sd', '-L', service_name, '_zordon', '_tcp', 'local'],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
        )
    except (FileNotFoundError, OSError):
        return None

    try:
        stdout, _ = process.communicate(timeout=timeout)
    except subprocess.TimeoutExpired:
        process.terminate()
        stdout, _ = process.communicate(timeout=1)

    for line in stdout.splitlines():
        if ' can be reached at ' not in line:
            continue
        match = re.search(r'can be reached at\s+([^:\s]+):(\d+)', line)
        if match:
            hostname, port = match.groups()
            return f'{hostname.rstrip(".")}:{port}'
    return None


def _discover_servers_with_dns_sd():
    servers = []
    for service_name in _discover_service_names_with_dns_sd():
        server = _resolve_service_with_dns_sd(service_name)
        if server and server not in servers:
            servers.append(server)
    return servers


def _discover_zordon_servers():
    discovered = _discover_servers_with_zeroconf()
    if discovered:
        return discovered
    return _discover_servers_with_dns_sd()


def _send_multipart_job(server, port, job_data, file_path):
    scheme, hostname, parsed_port = _normalize_server_url(server, port)
    if scheme != 'http':
        raise ValueError('Only http:// Zordon servers are currently supported.')

    boundary = f'----ZordonBlender{int(time.time() * 1000)}'
    json_payload = json.dumps(job_data).encode('utf-8')
    file_name = os.path.basename(file_path)
    file_size = os.path.getsize(file_path)

    parts_before_file = [
        f'--{boundary}\r\n',
        'Content-Disposition: form-data; name="json"\r\n',
        'Content-Type: application/json\r\n\r\n',
        json_payload.decode('utf-8'),
        '\r\n',
        f'--{boundary}\r\n',
        f'Content-Disposition: form-data; name="file"; filename="{file_name}"\r\n',
        'Content-Type: application/octet-stream\r\n\r\n',
    ]
    preamble = ''.join(parts_before_file).encode('utf-8')
    closing = f'\r\n--{boundary}--\r\n'.encode('utf-8')
    content_length = len(preamble) + file_size + len(closing)

    connection = http.client.HTTPConnection(hostname, parsed_port, timeout=UPLOAD_TIMEOUT)
    try:
        connection.putrequest('POST', '/api/jobs')
        connection.putheader('Content-Type', f'multipart/form-data; boundary={boundary}')
        connection.putheader('Content-Length', str(content_length))
        connection.endheaders()
        connection.send(preamble)

        with open(file_path, 'rb') as upload_file:
            while True:
                chunk = upload_file.read(1024 * 1024)
                if not chunk:
                    break
                connection.send(chunk)

        connection.send(closing)
        response = connection.getresponse()
        response_body = response.read().decode('utf-8', errors='replace')
        return response.status, response_body
    finally:
        connection.close()


def _current_file_path():
    return bpy.data.filepath


def _default_job_name():
    file_path = _current_file_path()
    if file_path:
        return os.path.splitext(os.path.basename(file_path))[0]
    return 'Blender Render'


def _camera_objects():
    return [obj for obj in bpy.data.objects if obj.type == 'CAMERA']


def _default_selected_camera_names(scene):
    cameras = _camera_objects()
    if not cameras:
        return []
    if scene.camera:
        return [scene.camera.name]
    if len(cameras) == 1:
        return [cameras[0].name]
    return []


def _selected_camera_names(props, scene=None):
    if props.selected_cameras:
        try:
            names = json.loads(props.selected_cameras)
        except json.JSONDecodeError:
            names = []
        existing_names = {camera.name for camera in _camera_objects()}
        selected_names = [name for name in names if name in existing_names]
        if selected_names:
            return selected_names
    if scene:
        return _default_selected_camera_names(scene)
    return []


def _set_selected_camera_names(props, camera_names):
    props.selected_cameras = json.dumps(list(camera_names))


def _active_camera_names(scene):
    if scene.camera:
        return [scene.camera.name]
    return _default_selected_camera_names(scene)


def _scene_export_format(scene):
    file_format = scene.render.image_settings.file_format
    return file_format or 'PNG'


class ZORDON_AddonPreferences(AddonPreferences):
    bl_idname = __name__

    servers: StringProperty(
        name='Servers',
        description='Comma-separated Zordon servers. Example: studio-render.local:8080, localhost:8080',
        default=DEFAULT_SERVER,
    )

    def draw(self, context):
        layout = self.layout
        layout.prop(self, 'servers')
        layout.operator('zordon.discover_servers', icon='VIEWZOOM')


class ZORDON_SubmissionProperties(PropertyGroup):
    server: EnumProperty(
        name='Server',
        description='Zordon server to submit this file to',
        items=_configured_server_items,
    )
    port: IntProperty(
        name='Fallback Port',
        description='Port used when the selected server does not include one',
        default=8080,
        min=1,
        max=65535,
    )
    job_name: StringProperty(
        name='Job Name',
        description='Name shown in Zordon',
        default='',
    )
    output_path: StringProperty(
        name='Output Name',
        description='Output name or path used by the Zordon job',
        default='',
    )
    notes: StringProperty(
        name='Notes',
        description='Optional notes for the Zordon job',
        default='',
    )
    save_before_submit: BoolProperty(
        name='Save Before Submit',
        description='Save the current .blend before uploading it',
        default=True,
    )
    use_scene_frame_range: BoolProperty(
        name='Use Scene Frame Range',
        description='Submit the current scene frame start and end',
        default=True,
    )
    selected_cameras: StringProperty(
        name='Selected Cameras',
        description='JSON encoded list of selected camera names',
        default='',
        options={'HIDDEN'},
    )


class ZORDON_OT_DiscoverServers(Operator):
    bl_idname = 'zordon.discover_servers'
    bl_label = 'Discover Servers'
    bl_description = 'Find Zordon servers advertised with Zeroconf'

    def execute(self, context):
        preferences = _addon_preferences(context)
        existing_servers = _server_list_from_preferences(preferences)
        discovered_servers = _discover_zordon_servers()
        if not discovered_servers:
            self.report({'WARNING'}, 'No Zordon servers found.')
            return {'CANCELLED'}

        merged_servers = existing_servers + discovered_servers
        _save_server_list(preferences, merged_servers)

        props = getattr(context.scene, 'zordon_submission', None)
        if props:
            try:
                props.server = discovered_servers[0]
            except (TypeError, ValueError):
                pass

        self.report({'INFO'}, f'Found {len(discovered_servers)} Zordon server(s).')
        return {'FINISHED'}


class ZORDON_OT_TestConnection(Operator):
    bl_idname = 'zordon.test_connection'
    bl_label = 'Test Connection'
    bl_description = 'Check whether the selected Zordon server is reachable'

    def execute(self, context):
        props = context.scene.zordon_submission
        try:
            status, body = _get_text(props.server, props.port, '/api/heartbeat')
        except (HTTPError, URLError, TimeoutError, OSError) as error:
            self.report({'ERROR'}, f'Could not reach Zordon: {error}')
            return {'CANCELLED'}

        if status == 200:
            self.report({'INFO'}, f'Connected to Zordon: {body}')
            return {'FINISHED'}

        self.report({'ERROR'}, f'Unexpected response from Zordon: HTTP {status}')
        return {'CANCELLED'}


class ZORDON_OT_RefreshCameras(Operator):
    bl_idname = 'zordon.refresh_cameras'
    bl_label = 'Refresh Cameras'
    bl_description = 'Refresh the camera checklist from the current Blender scene'

    def execute(self, context):
        _set_selected_camera_names(
            context.scene.zordon_submission,
            _default_selected_camera_names(context.scene),
        )
        return {'FINISHED'}


class ZORDON_OT_ToggleCamera(Operator):
    bl_idname = 'zordon.toggle_camera'
    bl_label = 'Toggle Camera'
    bl_description = 'Toggle whether this camera is submitted to Zordon'

    camera_name: StringProperty(name='Camera Name')

    def execute(self, context):
        props = context.scene.zordon_submission
        selected = set(_selected_camera_names(props, context.scene))
        if self.camera_name in selected:
            if len(selected) == 1:
                self.report({'WARNING'}, 'At least one camera must be selected.')
                return {'CANCELLED'}
            selected.remove(self.camera_name)
        else:
            selected.add(self.camera_name)
        _set_selected_camera_names(props, sorted(selected))
        return {'FINISHED'}


class ZORDON_OT_SelectAllCameras(Operator):
    bl_idname = 'zordon.select_all_cameras'
    bl_label = 'All'
    bl_description = 'Select all cameras for submission'

    def execute(self, context):
        props = context.scene.zordon_submission
        _set_selected_camera_names(props, [camera.name for camera in _camera_objects()])
        return {'FINISHED'}


class ZORDON_OT_SelectActiveCamera(Operator):
    bl_idname = 'zordon.select_active_camera'
    bl_label = 'Active Only'
    bl_description = 'Select only the active scene camera'

    def execute(self, context):
        active_cameras = _active_camera_names(context.scene)
        if not active_cameras:
            self.report({'WARNING'}, 'No cameras found.')
            return {'CANCELLED'}
        _set_selected_camera_names(context.scene.zordon_submission, active_cameras)
        return {'FINISHED'}


class ZORDON_OT_SubmitCurrentFile(Operator):
    bl_idname = 'zordon.submit_current_file'
    bl_label = 'Submit Current File'
    bl_description = 'Upload the current Blender file to the selected Zordon server'

    def execute(self, context):
        props = context.scene.zordon_submission
        scene = context.scene
        file_path = _current_file_path()

        if not file_path:
            self.report({'ERROR'}, 'Save this Blender file before submitting it to Zordon.')
            return {'CANCELLED'}

        if props.save_before_submit:
            bpy.ops.wm.save_as_mainfile(filepath=file_path)

        job_name = props.job_name.strip() or _default_job_name()
        output_path = props.output_path.strip() or job_name
        selected_cameras = _selected_camera_names(props, scene)
        if not selected_cameras:
            self.report({'ERROR'}, 'Add at least one camera before submitting to Zordon.')
            return {'CANCELLED'}
        resolution = (
            int(scene.render.resolution_x),
            int(scene.render.resolution_y),
        )

        job_data = {
            'owner': f'{getpass.getuser()}@{socket.gethostname()}',
            'engine_name': 'blender',
            'engine_version': 'latest',
            'name': job_name,
            'output_path': output_path,
            'start_frame': int(scene.frame_start) if props.use_scene_frame_range else int(scene.frame_current),
            'end_frame': int(scene.frame_end) if props.use_scene_frame_range else int(scene.frame_current),
            'priority': 1,
            'notes': props.notes,
            'enable_split_jobs': False,
            'split_jobs_same_os': False,
            'args': {
                'raw': '',
                'export_format': _scene_export_format(scene),
                'resolution': resolution,
                'fps': scene.render.fps,
            },
        }

        if len(selected_cameras) == 1:
            job_data['args']['camera'] = selected_cameras[0]
        elif len(selected_cameras) > 1:
            child_jobs = []
            for camera_name in selected_cameras:
                camera_suffix = camera_name.replace(' ', '-')
                child_args = dict(job_data['args'])
                child_args['camera'] = camera_name
                child_jobs.append({
                    'name': f'{job_name}_{camera_suffix}',
                    'output_path': f'{output_path}_{camera_suffix}',
                    'args': child_args,
                })
            job_data['child_jobs'] = child_jobs

        try:
            status, body = _send_multipart_job(props.server, props.port, job_data, file_path)
        except (ValueError, TimeoutError, OSError) as error:
            self.report({'ERROR'}, f'Error submitting to Zordon: {error}')
            return {'CANCELLED'}

        if 200 <= status < 300:
            self.report({'INFO'}, f'Submitted "{job_name}" to Zordon.')
            return {'FINISHED'}

        self.report({'ERROR'}, f'Zordon returned HTTP {status}: {body}')
        return {'CANCELLED'}


class ZORDON_PT_SubmitPanel(Panel):
    bl_label = 'Zordon'
    bl_idname = 'ZORDON_PT_submit_panel'
    bl_space_type = 'PROPERTIES'
    bl_region_type = 'WINDOW'
    bl_context = 'render'

    def draw(self, context):
        props = context.scene.zordon_submission
        layout = self.layout
        default_name = _default_job_name()
        cameras = _camera_objects()
        selected_cameras = set(_selected_camera_names(props, context.scene))

        layout.prop(props, 'server')
        layout.prop(props, 'port')

        row = layout.row(align=True)
        row.operator('zordon.discover_servers', icon='VIEWZOOM')
        row.operator('zordon.test_connection', icon='LINKED')

        layout.separator()
        if not props.job_name:
            layout.label(text=f'Job Name Default: {default_name}')
        layout.prop(props, 'job_name')
        if not props.output_path:
            layout.label(text=f'Output Name Default: {default_name}')
        layout.prop(props, 'output_path')
        layout.prop(props, 'notes')
        layout.prop(props, 'save_before_submit')
        layout.prop(props, 'use_scene_frame_range')

        if props.use_scene_frame_range:
            layout.label(text=f'Frames: {context.scene.frame_start} - {context.scene.frame_end}')
        else:
            layout.label(text=f'Frame: {context.scene.frame_current}')

        layout.label(text=f'Format: {_scene_export_format(context.scene)}')

        layout.separator()
        camera_header = layout.row(align=True)
        camera_header.label(text='Cameras')
        camera_header.operator('zordon.refresh_cameras', text='', icon='FILE_REFRESH')
        camera_header.operator('zordon.select_all_cameras', text='All')
        camera_header.operator('zordon.select_active_camera', text='Active Only')

        if cameras:
            for camera in cameras:
                icon = 'CHECKBOX_HLT' if camera.name in selected_cameras else 'CHECKBOX_DEHLT'
                label = f'{camera.name} - {camera.data.lens:g}mm'
                row = layout.row()
                op = row.operator('zordon.toggle_camera', text=label, icon=icon, depress=camera.name in selected_cameras)
                op.camera_name = camera.name
        else:
            layout.label(text='No cameras found')

        layout.operator('zordon.submit_current_file', icon='RENDER_ANIMATION')


classes = (
    ZORDON_AddonPreferences,
    ZORDON_SubmissionProperties,
    ZORDON_OT_DiscoverServers,
    ZORDON_OT_TestConnection,
    ZORDON_OT_RefreshCameras,
    ZORDON_OT_ToggleCamera,
    ZORDON_OT_SelectAllCameras,
    ZORDON_OT_SelectActiveCamera,
    ZORDON_OT_SubmitCurrentFile,
    ZORDON_PT_SubmitPanel,
)


def register():
    for cls in classes:
        bpy.utils.register_class(cls)
    bpy.types.Scene.zordon_submission = bpy.props.PointerProperty(type=ZORDON_SubmissionProperties)


def unregister():
    del bpy.types.Scene.zordon_submission
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)


if __name__ == '__main__':
    register()