Zordon/lib/client/dashboard_window.py

import datetime
import logging
import tkinter as tk
import threading
import time
import socket
import os
from tkinter import ttk, messagebox, simpledialog
from PIL import Image, ImageTk
from lib.client.new_job_window import NewJobWindow
from lib.server.server_proxy import RenderServerProxy
from lib.server.zeroconf_server import ZeroconfServer
from lib.workers.base_worker import RenderStatus
from lib.utilities.misc_helper import launch_url, file_exists_in_mounts, get_time_elapsed

logger = logging.getLogger()


def sort_column(tree, col, reverse=False):
    data = [(tree.set(child, col), child) for child in tree.get_children('')]
    data.sort(reverse=reverse)
    for index, (_, child) in enumerate(data):
        tree.move(child, '', index)


def make_sortable(tree):
    for col in tree["columns"]:
        tree.heading(col, text=col, command=lambda c=col: sort_column(tree, c))


class DashboardWindow:

    lib_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    image_path = os.path.join(lib_path, 'server', 'static', 'images')
    default_image = Image.open(os.path.join(image_path, 'desktop.png'))

    def __init__(self):

        # Create a Treeview widget
        self.root = tk.Tk()
        self.root.title("Zordon Dashboard")
        self.current_hostname = None
        self.server_proxies = {}
        self.added_hostnames = []

        # Setup zeroconf
        ZeroconfServer.configure("_zordon._tcp.local.", socket.gethostname(), 8080)
        ZeroconfServer.start(listen_only=True)

        # Setup photo preview
        photo_pad = tk.Frame(self.root, background="gray")
        photo_pad.pack(fill=tk.BOTH, pady=5, padx=5)
        self.photo_label = tk.Label(photo_pad, height=500)
        self.photo_label.pack(fill=tk.BOTH, expand=True)
        self.set_image(self.default_image)

        server_frame = tk.LabelFrame(self.root, text="Server")
        server_frame.pack(fill=tk.BOTH, pady=5, padx=5, expand=True)

        # Create server tree
        left_frame = tk.Frame(server_frame)
        left_frame.pack(side=tk.LEFT, expand=True, fill=tk.BOTH)
        self.server_tree = ttk.Treeview(left_frame, show="headings")
        self.server_tree.pack(expand=True, fill=tk.BOTH)
        self.server_tree["columns"] = ("Server", "Status")
        self.server_tree.bind("<<TreeviewSelect>>", self.server_picked)
        self.server_tree.column("Server", width=200)
        self.server_tree.column("Status", width=80)

        left_button_frame = tk.Frame(left_frame)
        left_button_frame.pack(side=tk.BOTTOM, fill=tk.X, padx=5, pady=5, expand=False)

        # Create buttons
        self.remove_server_button = tk.Button(left_button_frame, text="-", command=self.remove_server_button)
        self.remove_server_button.pack(side=tk.RIGHT)
        self.remove_server_button.config(state='disabled')
        add_server_button = tk.Button(left_button_frame, text="+", command=self.add_server_button)
        add_server_button.pack(side=tk.RIGHT)

        # Create separator
        separator = ttk.Separator(server_frame, orient=tk.VERTICAL)
        separator.pack(side=tk.LEFT, padx=5, pady=5, fill=tk.Y)

        # Setup the Tree
        self.job_tree = ttk.Treeview(server_frame, show="headings")
        self.job_tree.tag_configure(RenderStatus.RUNNING.value, background='lawn green', font=('', 0, 'bold'))
        self.job_tree.bind("<<TreeviewSelect>>", self.job_picked)
        self.job_tree["columns"] = ("id", "Name", "Renderer", "Priority", "Status", "Time Elapsed", "Frames",
                                    "Date Added", "Parent", "")

        # Format the columns
        self.job_tree.column("id", width=0, stretch=False)
        self.job_tree.column("Name", width=300)
        self.job_tree.column("Renderer", width=100, stretch=False)
        self.job_tree.column("Priority", width=50, stretch=False)
        self.job_tree.column("Status", width=100, stretch=False)
        self.job_tree.column("Time Elapsed", width=100, stretch=False)
        self.job_tree.column("Frames", width=50, stretch=False)
        self.job_tree.column("Date Added", width=150, stretch=True)
        self.job_tree.column("Parent", width=250, stretch=True)

        # Create the column headings
        for name in self.job_tree['columns']:
            self.job_tree.heading(name, text=name)

        # Pack the Treeview widget
        self.job_tree.pack(fill=tk.BOTH, expand=True)

        button_frame = tk.Frame(server_frame)
        button_frame.pack(pady=5, fill=tk.X, expand=False)

        # Create buttons
        self.logs_button = tk.Button(button_frame, text="Logs", command=self.open_logs)
        self.show_files_button = tk.Button(button_frame, text="Show Files", command=self.show_files)
        self.stop_button = tk.Button(button_frame, text="Stop", command=self.stop_job)
        self.delete_button = tk.Button(button_frame, text="Delete", command=self.delete_job)
        add_job_button = tk.Button(button_frame, text="Add Job", command=self.show_new_job_window)

        # Pack the buttons in the frame
        self.stop_button.pack(side=tk.LEFT)
        self.stop_button.config(state='disabled')
        self.delete_button.pack(side=tk.LEFT)
        self.delete_button.config(state='disabled')
        self.show_files_button.pack(side=tk.LEFT)
        self.show_files_button.config(state='disabled')
        self.logs_button.pack(side=tk.LEFT)
        self.logs_button.config(state='disabled')
        add_job_button.pack(side=tk.RIGHT)

        # Start the Tkinter event loop
        self.root.geometry("500x600+300+300")
        self.root.maxsize(width=2000, height=1200)
        self.root.minsize(width=900, height=800)
        make_sortable(self.job_tree)
        make_sortable(self.server_tree)

        # update servers
        self.update_servers()
        try:
            selected_server = self.server_tree.get_children()[0]
            self.server_tree.selection_set(selected_server)
            self.server_picked()
        except IndexError:
            pass

        # update jobs
        self.update_jobs()
        try:
            selected_job = self.job_tree.get_children()[0]
            self.job_tree.selection_set(selected_job)
            self.job_picked()
        except IndexError:
            pass

        # start background update
        x = threading.Thread(target=self.__background_update)
        x.daemon = True
        x.start()

    @property
    def current_server_proxy(self):
        return self.server_proxies.get(self.current_hostname, None)

    def remove_server_button(self):
        new_hostname = self.server_tree.selection()[0]
        if new_hostname in self.added_hostnames:
            self.added_hostnames.remove(new_hostname)
            self.update_servers()
            if self.server_tree.get_children():
                self.server_tree.selection_set(self.server_tree.get_children()[0])
                self.server_picked(event=None)

    def add_server_button(self):
        hostname = simpledialog.askstring("Server Hostname", "Enter the server hostname to add:")
        if hostname:
            hostname = hostname.strip()
            if hostname not in self.added_hostnames:
                if RenderServerProxy(hostname=hostname).connect():
                    self.added_hostnames.append(hostname)
                    self.update_servers()
                else:
                    messagebox.showerror("Cannot Connect", f"Cannot connect to server at hostname: '{hostname}'")

    def server_picked(self, event=None):
        try:
            new_hostname = self.server_tree.selection()[0]
            self.remove_server_button.config(state="normal" if new_hostname in self.added_hostnames else "disabled")
            if self.current_hostname == new_hostname:
                return
            self.current_hostname = new_hostname
            self.update_jobs(clear_table=True)
        except IndexError:
            pass

    def selected_job_ids(self):
        selected_items = self.job_tree.selection()  # Get the selected item
        row_data = [self.job_tree.item(item) for item in selected_items]  # Get the text of the selected item
        job_ids = [row['values'][0] for row in row_data]
        return job_ids

    def stop_job(self):
        job_ids = self.selected_job_ids()
        for job_id in job_ids:
            self.current_server_proxy.cancel_job(job_id, confirm=True)
        self.update_jobs(clear_table=True)

    def delete_job(self):
        job_ids = self.selected_job_ids()
        if len(job_ids) == 1:
            job = next((d for d in self.current_server_proxy.get_all_jobs() if d.get('id') == job_ids[0]), None)
            display_name = job['name'] or os.path.basename(job['input_path'])
            message = f"Are you sure you want to delete the job:\n{display_name}?"
        else:
            message = f"Are you sure you want to delete these {len(job_ids)} jobs?"

        result = messagebox.askyesno("Confirmation", message)
        if result:
            for job_id in job_ids:
                self.current_server_proxy.request_data(f'job/{job_id}/delete?confirm=true')
            self.update_jobs(clear_table=True)

    def set_image(self, image):
        thumb_image = ImageTk.PhotoImage(image)
        if thumb_image:
            self.photo_label.configure(image=thumb_image)
            self.photo_label.image = thumb_image

    def job_picked(self, event=None):
        job_id = self.selected_job_ids()[0] if self.selected_job_ids() else None
        if job_id:
            # update thumb
            def fetch_preview():
                try:
                    before_fetch_hostname = self.current_server_proxy.hostname
                    response = self.current_server_proxy.request(f'job/{job_id}/thumbnail?size=big')
                    if response.ok:
                        import io
                        image_data = response.content
                        image = Image.open(io.BytesIO(image_data))
                        if self.current_server_proxy.hostname == before_fetch_hostname and job_id == self.selected_job_ids()[0]:
                            self.set_image(image)
                except ConnectionError as e:
                    logger.error(f"Connection error fetching image: {e}")
                except Exception as e:
                    logger.error(f"Error fetching image: {e}")

            fetch_thread = threading.Thread(target=fetch_preview)
            fetch_thread.daemon = True
            fetch_thread.start()
        else:
            self.set_image(self.default_image)

        # update button status
        current_jobs = self.current_server_proxy.get_all_jobs() or []
        job = next((d for d in current_jobs if d.get('id') == job_id), None)
        stop_button_state = 'normal' if job and job['status'] == RenderStatus.RUNNING.value else 'disabled'
        self.stop_button.config(state=stop_button_state)
        
        generic_button_state = 'normal' if job else 'disabled'
        self.show_files_button.config(state=generic_button_state)
        self.delete_button.config(state=generic_button_state)
        self.logs_button.config(state=generic_button_state)

    def show_files(self):
        if not self.selected_job_ids():
            return

        job = next((d for d in self.current_server_proxy.get_all_jobs() if d.get('id') == self.selected_job_ids()[0]), None)
        output_path = os.path.dirname(job['output_path'])  # check local filesystem
        if not os.path.exists(output_path):
            output_path = file_exists_in_mounts(output_path)  # check any attached network shares
        if not output_path:
            return messagebox.showerror("File Not Found", "The file could not be found. Check your network mounts.")
        launch_url(output_path)

    def open_logs(self):
        if self.selected_job_ids():
            url = f'http://{self.current_server_proxy.hostname}:{self.current_server_proxy.port}/api/job/{self.selected_job_ids()[0]}/logs'
            launch_url(url)

    def mainloop(self):
        self.root.mainloop()

    def __background_update(self):
        while True:
            self.update_servers()
            self.update_jobs()
            time.sleep(1)

    def update_servers(self):

        def update_row(tree, id, new_values, tags=None):
            for item in tree.get_children():
                values = tree.item(item, "values")
                if values[0] == id:
                    if tags:
                        tree.item(item, values=new_values, tags=tags)
                    else:
                        tree.item(item, values=new_values)
                    break

        current_servers = list(set(ZeroconfServer.found_clients() + self.added_hostnames))
        for hostname in current_servers:
            if not self.server_proxies.get(hostname, None):
                new_proxy = RenderServerProxy(hostname=hostname)
                new_proxy.start_background_update()
                self.server_proxies[hostname] = new_proxy

        try:
            for hostname, proxy in self.server_proxies.items():
                if hostname not in self.server_tree.get_children():
                    self.server_tree.insert("", tk.END, iid=hostname, values=(hostname, proxy.status(), ))
                else:
                    update_row(self.server_tree, hostname, new_values=(hostname, proxy.status()))
        except RuntimeError:
            pass

        # remove any servers that don't belong
        for row in self.server_tree.get_children():
            if row not in current_servers:
                self.server_tree.delete(row)
                proxy = self.server_proxies.get(row, None)
                if proxy:
                    proxy.stop_background_update()
                self.server_proxies.pop(row)

    def update_jobs(self, clear_table=False):

        if not self.current_server_proxy:
            return

        def update_row(tree, id, new_values, tags=None):
            for item in tree.get_children():
                values = tree.item(item, "values")
                if values[0] == id:
                    tree.item(item, values=new_values, tags=tags)
                    break

        if clear_table:
            self.job_tree.delete(*self.job_tree.get_children())

        job_fetch = self.current_server_proxy.get_all_jobs(ignore_token=clear_table)
        if job_fetch:
            for job in job_fetch:
                display_status = job['status'] if job['status'] != RenderStatus.RUNNING.value else \
                    ('%.0f%%' % (job['percent_complete'] * 100))  # if running, show percent, otherwise just show status
                tags = (job['status'],)
                start_time = datetime.datetime.fromisoformat(job['start_time']) if job['start_time'] else None
                end_time = datetime.datetime.fromisoformat(job['end_time']) if job['end_time'] else None

                time_elapsed = "" if (job['status'] != RenderStatus.RUNNING.value and not end_time) else \
                    get_time_elapsed(start_time, end_time)

                values = (job['id'],
                          job['name'] or os.path.basename(job['input_path']),
                          job['renderer'] + "-" + job['renderer_version'],
                          job['priority'],
                          display_status,
                          time_elapsed,
                          job['total_frames'],
                          job['date_created'],
                          job['parent'])
                try:
                    if self.job_tree.exists(job['id']):
                        update_row(self.job_tree, job['id'], new_values=values, tags=tags)
                    else:
                        self.job_tree.insert("", tk.END, iid=job['id'], values=values, tags=tags)
                except tk.TclError:
                    pass

            # remove any jobs that don't belong
            all_job_ids = [job['id'] for job in job_fetch]
            for row in self.job_tree.get_children():
                if row not in all_job_ids:
                    self.job_tree.delete(row)

    def show_new_job_window(self):
        new_window = tk.Toplevel(self.root)
        new_window.title("New Window")
        new_window.geometry("500x600+300+300")
        new_window.resizable(False, height=True)
        x = NewJobWindow(parent=new_window, clients=list(self.server_tree.get_children()))
        x.pack()


def start_client():

    logging.basicConfig(format='%(asctime)s: %(levelname)s: %(module)s: %(message)s', datefmt='%d-%b-%y %H:%M:%S',
                        level='INFO'.upper())

    x = DashboardWindow()
    x.mainloop()


if __name__ == '__main__':
    start_client()