"""
ComfyUI Deadline submission nodes.

This module owns the ComfyUI-side submission flow. It packages the current API
prompt, stages referenced input assets, and submits a render job to Deadline.
"""

import copy
import filecmp
import hashlib
import json
import ntpath
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
import uuid
from typing import Any, Dict, List, Optional, Tuple


DEADLINE_COMMAND_PATHS = {
    "windows": "C:\\Program Files\\Thinkbox\\Deadline10\\bin\\deadlinecommand.exe",
    "linux": "/opt/Thinkbox/Deadline10/bin/deadlinecommand",
}

DEADLINE_SUBMIT_NODE_TYPES = {"DeadlineSubmit", "SaveAndSubmitNode"}
OUTPUT_NODE_TYPES = {"SaveImage", "PreviewImage", "SaveVideo", "VHS_VideoCombine"}
INPUT_LOADER_FIELDS = {
    "LoadImage": ("image",),
    "LoadImageMask": ("image",),
    "LoadAudio": ("audio",),
    "LoadVideo": ("file",),
    # VideoHelperSuite file-based loaders.  Directory loaders such as
    # VHS_LoadImagesPath intentionally remain unsupported here.
    "VHS_LoadVideo": ("video",),
    "VHS_LoadVideoFFmpeg": ("video",),
    "VHS_LoadAudioUpload": ("audio",),
}
ABSOLUTE_FILE_LOADER_FIELDS = {
    "VHS_LoadVideoPath": ("video",),
    "VHS_LoadVideoFFmpegPath": ("video",),
    "VHS_LoadImagePath": ("image",),
    "VHS_LoadAudio": ("audio_file",),
}
MEDIA_EXTENSIONS = {
    ".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif", ".tif", ".tiff", ".exr",
    ".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v",
    ".wav", ".mp3", ".flac", ".ogg", ".m4a", ".aac",
}

# A worker that is allowed to reuse a long-running GUI endpoint must prove that
# the endpoint belongs to this installation.  A random value created at module
# import also distinguishes a restarted ComfyUI session from the old process.
DEADLINE_ENDPOINT_SESSION_ID = uuid.uuid4().hex


def _validate_path_text(path: str) -> None:
    if not isinstance(path, str) or not path or any(c in path for c in '\x00\r\n"'):
        raise ValueError("Path must be nonempty and contain no quotes or control characters.")
    normalized = path.replace("\\", "/")
    if normalized.startswith(("//?/", "//./")):
        raise ValueError("Windows device paths are not allowed.")
    _, tail = ntpath.splitdrive(path)
    for part in tail.replace("\\", "/").split("/"):
        if not part or part == ".":
            continue
        if part == ".." or part.rstrip(" .") != part or ":" in part or re.fullmatch(r"(?i:CON|PRN|AUX|NUL|COM[1-9]|LPT[1-9])(?:\..*)?", part):
            raise ValueError("Path contains traversal, a device name, or an alternate data stream.")


def _approved_path(path: str, roots: List[str]) -> str:
    _validate_path_text(path)
    if not os.path.isabs(path):
        raise ValueError("An absolute path inside an approved storage root is required.")
    # Preserve the rooted UNC share anchor, including when realpath drops its slash.
    candidate = os.path.normpath(os.path.abspath(path) + os.sep)
    for root in roots:
        _validate_path_text(root)
        if not os.path.isabs(root):
            raise ValueError("Approved storage roots must be absolute paths.")
        # Check the spelling before resolving links: an unapproved UNC path must
        # never cause a connection (and Windows authentication) to another host.
        root = os.path.normpath(os.path.abspath(root) + os.sep)
        try:
            if os.path.commonpath([os.path.normcase(root), os.path.normcase(candidate)]) != os.path.normcase(root):
                continue
            real_root = os.path.normcase(os.path.normpath(os.path.realpath(root) + os.sep))
            real_candidate = os.path.normcase(os.path.normpath(os.path.realpath(candidate) + os.sep))
            if os.path.commonpath([real_root, real_candidate]) == real_root:
                return candidate
        except ValueError:
            continue
    raise ValueError("Path is outside the approved storage roots.")


def _deadline_value(value: Any) -> str:
    value = str(value)
    if any(c in value for c in "\r\n\x00"):
        raise ValueError("Deadline job values cannot contain newlines or NUL characters.")
    return value


def _farm_path(path: str, settings: Dict[str, Any]) -> str:
    drive, tail = ntpath.splitdrive(path)
    network_root = settings.get("drive_mappings", {}).get(drive.upper())
    if network_root:
        return ntpath.join(network_root, tail.lstrip("\\/"))
    return path


def _coerce_bool(value: Any) -> bool:
    if isinstance(value, str):
        return value.strip().lower() in {"1", "true", "yes", "on"}
    return bool(value)


class NodeDefaults:
    JOB_NAME = "ComfyUI via Deadline"
    PRIORITY = 50
    POOL = "none"
    GROUP = "none"
    BATCH_COUNT = 1
    CHUNK_SIZE = 1
    MAX_BATCH_COUNT = 10000
    MAX_CHUNK_SIZE = 256
    MAX_PRIORITY = 100


class DeadlineCommandHelper:
    @staticmethod
    def get_deadline_command() -> str:
        deadline_bin = os.environ.get("DEADLINE_PATH", "")

        if not deadline_bin and os.path.exists("/Users/Shared/Thinkbox/DEADLINE_PATH"):
            try:
                with open("/Users/Shared/Thinkbox/DEADLINE_PATH", "r", encoding="utf-8") as handle:
                    deadline_bin = handle.read().strip()
            except Exception:
                deadline_bin = ""

        candidates = []
        if deadline_bin:
            candidates.append(os.path.join(deadline_bin, "deadlinecommand.exe" if os.name == "nt" else "deadlinecommand"))
        candidates.append(DEADLINE_COMMAND_PATHS["windows"] if sys.platform.startswith("win") else DEADLINE_COMMAND_PATHS["linux"])

        for candidate in candidates:
            if os.path.isabs(candidate) and os.path.isfile(candidate):
                return candidate
        return ""

    @staticmethod
    def call_deadline_command(arguments: List[str], hide_window: bool = True) -> str:
        deadline_command = DeadlineCommandHelper.get_deadline_command()
        if not deadline_command:
            raise RuntimeError("Deadline command not found. Set DEADLINE_PATH or install Deadline Client.")

        startupinfo = None
        creationflags = 0
        if os.name == "nt" and hide_window:
            try:
                startupinfo = subprocess.STARTUPINFO()
                startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            except Exception:
                startupinfo = None
        elif os.name == "nt":
            creationflags = 0x08000000

        process = subprocess.Popen(
            [deadline_command] + arguments,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            startupinfo=startupinfo,
            creationflags=creationflags,
            shell=False,
        )
        stdout, stderr = process.communicate()
        output = stdout.decode(errors="replace") if isinstance(stdout, bytes) else str(stdout)
        errors = stderr.decode(errors="replace") if isinstance(stderr, bytes) else str(stderr)

        if process.returncode != 0:
            raise RuntimeError(f"deadlinecommand failed with code {process.returncode}: {errors or output}")
        return output

    @staticmethod
    def get_job_id_from_submission(submission_results: str) -> str:
        for token in submission_results.replace("\r", "\n").split():
            if token.startswith("JobID="):
                return token.split("=", 1)[1].strip()
        return ""

    @staticmethod
    def get_farm_settings() -> Dict[str, Any]:
        script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "scripts", "read_farm_settings.py")
        output = DeadlineCommandHelper.call_deadline_command(["-ExecuteScriptNoGui", script])
        marker = "COMFY_DEADLINE_FARM_SETTINGS="
        for line in output.splitlines():
            if not line.startswith(marker):
                continue
            settings = json.loads(line[len(marker):])
            if not isinstance(settings, dict):
                raise ValueError("Deadline returned invalid farm storage settings.")
            for key in ("output_roots", "input_roots"):
                roots = settings.get(key)
                if not isinstance(roots, list):
                    raise ValueError("Deadline returned invalid farm storage roots.")
                for root in roots:
                    _validate_path_text(root)
                    if not os.path.isabs(root):
                        raise ValueError("Farm storage roots must be absolute paths.")
            mappings = settings.get("drive_mappings", {})
            if not isinstance(mappings, dict):
                raise ValueError("Deadline returned invalid drive mappings.")
            for drive, network_root in mappings.items():
                _validate_path_text(network_root)
                if not re.fullmatch(r"[A-Z]:", drive) or not network_root.startswith("\\\\"):
                    raise ValueError("Deadline returned invalid drive mappings.")
            return settings
        raise RuntimeError("Could not read ComfyUI farm settings from Deadline. Check the Deadline Client connection.")


class WorkflowProcessor:
    @staticmethod
    def normalize_prompt(prompt: Dict[str, Any]) -> Dict[str, Any]:
        if not isinstance(prompt, dict) or not prompt:
            raise ValueError("ComfyUI did not provide a valid API prompt.")
        return copy.deepcopy(prompt)

    @staticmethod
    def prepare_for_worker(prompt: Dict[str, Any]) -> Dict[str, Any]:
        prepared = WorkflowProcessor.normalize_prompt(prompt)
        removed = []
        for node_id, node in list(prepared.items()):
            if isinstance(node, dict) and node.get("class_type") in DEADLINE_SUBMIT_NODE_TYPES:
                removed.append(node_id)
                del prepared[node_id]

        if removed:
            print(f"Deadline Submission: Removed submit node(s) from worker prompt: {', '.join(map(str, removed))}")

        WorkflowProcessor.validate_worker_prompt(prepared)
        return prepared

    @staticmethod
    def validate_worker_prompt(prompt: Dict[str, Any]) -> None:
        if not prompt:
            raise ValueError("Worker prompt is empty after removing Deadline submit nodes.")

        has_output = any(
            isinstance(node, dict) and node.get("class_type") in OUTPUT_NODE_TYPES
            for node in prompt.values()
        )
        if not has_output:
            print("Deadline Submission: Warning - worker prompt has no known output node.")


class InputAssetStager:
    def __init__(self, output_directory: str, job_name: str, submission_id: str):
        self.output_directory = os.path.abspath(output_directory)
        self.job_name = job_name
        self.submission_id = submission_id
        self.farm_settings = None

    def stage_referenced_assets(self, prompt: Dict[str, Any]) -> Tuple[str, str, List[Dict[str, Any]]]:
        input_dir = self._get_local_input_directory()
        references = self._collect_references(prompt, input_dir)
        staging_dir = self._staging_directory()
        manifest_path = os.path.join(staging_dir, f"deadline_input_manifest_{self.submission_id}.json")

        if not references:
            os.makedirs(staging_dir, exist_ok=True)
            manifest = {
                "submission_id": self.submission_id,
                "input_directory": staging_dir,
                "assets": [],
            }
            self._write_manifest(manifest_path, manifest)
            return staging_dir, manifest_path, []

        os.makedirs(staging_dir, exist_ok=True)
        assets = []
        for original_reference, source_path in sorted(references.items()):
            # Absolute-path loader nodes are converted to a farm-visible staged
            # absolute path later.  Use only the basename inside the staging
            # directory; never reproduce the submitter's local directory tree.
            destination_rel_path = (
                os.path.basename(source_path)
                if os.path.isabs(original_reference)
                else original_reference
            )
            staged_rel_path, destination_path = self._resolve_destination(staging_dir, destination_rel_path, source_path)
            os.makedirs(os.path.dirname(destination_path), exist_ok=True)
            if not os.path.exists(destination_path):
                shutil.copy2(source_path, destination_path)
            assets.append({
                "original_relative_path": original_reference.replace("\\", "/"),
                "staged_relative_path": staged_rel_path.replace("\\", "/"),
                "relative_path": staged_rel_path.replace("\\", "/"),
                "source": source_path,
                "destination": destination_path,
                "size": os.path.getsize(destination_path),
            })

        manifest = {
            "submission_id": self.submission_id,
            "input_directory": staging_dir,
            "assets": assets,
        }
        self._write_manifest(manifest_path, manifest)
        print(f"Deadline Submission: Staged {len(assets)} input asset(s) to {staging_dir}")
        return staging_dir, manifest_path, assets

    def _get_local_input_directory(self) -> str:
        try:
            import folder_paths
            return os.path.abspath(folder_paths.get_input_directory())
        except Exception as exc:
            raise RuntimeError(f"Could not resolve ComfyUI input directory: {exc}")

    def _collect_references(self, prompt: Dict[str, Any], input_dir: str) -> Dict[str, str]:
        references: Dict[str, str] = {}
        for node in prompt.values():
            if not isinstance(node, dict):
                continue

            class_type = node.get("class_type", "")
            inputs = node.get("inputs", {})
            if not isinstance(inputs, dict):
                continue

            candidate_values: List[Tuple[Any, bool, bool]] = []
            for field_name in INPUT_LOADER_FIELDS.get(class_type, ()):
                if field_name in inputs:
                    candidate_values.append((inputs[field_name], True, False))

            for field_name in ABSOLUTE_FILE_LOADER_FIELDS.get(class_type, ()):
                if field_name in inputs:
                    candidate_values.append((inputs[field_name], True, True))

            for value in inputs.values():
                if isinstance(value, str):
                    candidate_values.append((value, False, False))

            for value, strict, allow_absolute in candidate_values:
                if not isinstance(value, str):
                    continue
                resolved = self._resolve_input_file(value, input_dir, strict, allow_absolute)
                if not resolved:
                    continue
                rel_path, source_path = resolved
                references[rel_path] = source_path
        return references

    def _resolve_destination(self, staging_dir: str, rel_path: str, source_path: str) -> Tuple[str, str]:
        destination_path = _approved_path(os.path.join(staging_dir, rel_path), [staging_dir])
        if not os.path.exists(destination_path):
            return rel_path, destination_path

        if os.path.isfile(destination_path) and filecmp.cmp(source_path, destination_path, shallow=False):
            return rel_path, destination_path

        stem, extension = os.path.splitext(rel_path)
        staged_rel_path = f"{stem}_{self.submission_id}{extension}"
        destination_path = _approved_path(os.path.join(staging_dir, staged_rel_path), [staging_dir])
        if os.path.lexists(destination_path):
            raise ValueError("Staged input collision would overwrite an existing file.")
        return staged_rel_path, destination_path

    def _resolve_input_file(
        self,
        value: str,
        input_dir: str,
        strict: bool,
        allow_absolute: bool = False,
    ) -> Optional[Tuple[str, str]]:
        clean_value, annotation = self._strip_annotation(value)
        if annotation in {"output", "temp"}:
            return None
        if not clean_value:
            return None
        if os.path.splitext(clean_value)[1].lower() not in MEDIA_EXTENSIONS:
            return None

        if os.path.isabs(clean_value):
            if not allow_absolute:
                return None
            try:
                candidate = _approved_path(clean_value, [input_dir])
            except ValueError:
                if self.farm_settings is None:
                    self.farm_settings = DeadlineCommandHelper.get_farm_settings()
                candidate = _approved_path(_farm_path(clean_value, self.farm_settings), [input_dir] + self.farm_settings["input_roots"])
            if not os.path.isfile(candidate):
                if not strict:
                    return None
                raise FileNotFoundError(f"Referenced input asset was not found: {value} ({candidate})")
            return clean_value, candidate

        candidate = _approved_path(os.path.join(input_dir, clean_value), [input_dir])
        if not os.path.isfile(candidate):
            if not strict:
                return None
            raise FileNotFoundError(f"Referenced input asset was not found: {value} ({candidate})")

        rel_path = os.path.relpath(candidate, input_dir)
        return rel_path, candidate

    def _strip_annotation(self, value: str) -> Tuple[str, Optional[str]]:
        match = re.match(r"^(.*)\s+\[(input|output|temp)\]\s*$", value)
        if not match:
            return value.strip().replace("/", os.sep), None
        return match.group(1).strip().replace("/", os.sep), match.group(2)

    def _staging_directory(self) -> str:
        if not re.fullmatch(r"[A-Za-z0-9_-]+", self.submission_id):
            raise ValueError("Submission ID must contain only letters, digits, underscores, or hyphens.")
        return _approved_path(os.path.join(self.output_directory, ".deadline-inputs", self.submission_id), [self.output_directory])

    def _write_manifest(self, manifest_path: str, manifest: Dict[str, Any]) -> None:
        _approved_path(manifest_path, [self._staging_directory()])
        with open(manifest_path, "w", encoding="utf-8") as handle:
            json.dump(manifest, handle, indent=2)


class DeadlineJobSubmitter:
    def __init__(self, workflow_data: Dict[str, Any], job_config: Dict[str, Any]):
        self.workflow_data = workflow_data
        self.job_config = job_config

    def submit_job(self) -> Tuple[bool, str]:
        try:
            submission_dir = tempfile.mkdtemp(prefix="comfy_deadline_job_")
            job_info_file, plugin_info_file, auxiliary_files = self._create_submission_files(submission_dir)
            result = DeadlineCommandHelper.call_deadline_command([job_info_file, plugin_info_file] + auxiliary_files)
            job_id = DeadlineCommandHelper.get_job_id_from_submission(result)
            if not job_id:
                return False, f"Deadline submission did not return a JobID. Output: {result}"
            return True, job_id
        except Exception as exc:
            return False, str(exc)

    def _create_submission_files(self, submission_dir: str) -> Tuple[str, str, List[str]]:
        job_info_file = os.path.join(submission_dir, "job_info.txt")
        plugin_info_file = os.path.join(submission_dir, "plugin_info.txt")
        prompt_file = os.path.join(submission_dir, "prompt_to_execute.json")
        standard_workflow_file = os.path.join(submission_dir, "workflow.json")

        with open(prompt_file, "w", encoding="utf-8") as handle:
            json.dump(self.workflow_data, handle, indent=2)

        auxiliary_files = [prompt_file]
        if config_package := self.job_config.get("spool_package"):
            package_file = os.path.join(submission_dir, "spool_request.json")
            with open(package_file, "w", encoding="utf-8") as handle:
                json.dump(config_package, handle, ensure_ascii=False, indent=2)
            auxiliary_files.append(package_file)
        standard_workflow = self.job_config.get("standard_workflow")
        if standard_workflow:
            with open(standard_workflow_file, "w", encoding="utf-8") as handle:
                json.dump(standard_workflow, handle, indent=2)
            auxiliary_files.append(standard_workflow_file)

        self._write_job_info(job_info_file)
        self._write_plugin_info(plugin_info_file)
        return job_info_file, plugin_info_file, auxiliary_files

    def _write_job_info(self, path: str) -> None:
        config = self.job_config
        batch_count = int(config["batch_count"])
        chunk_size = max(1, int(config["chunk_size"]))
        output_dir = os.path.abspath(config["output_directory"])

        for key in ("job_name", "comment", "department", "pool", "group", "output_directory"):
            _deadline_value(config.get(key, ""))
        if config.get("spool_package"):
            for key in ("submissionId", "attemptId"):
                _deadline_value(config["spool_package"][key])
        with open(path, "w", encoding="utf-8") as handle:
            handle.write("Plugin=ComfyUI\n")
            handle.write(f"Name={config['job_name']}\n")
            handle.write(f"Comment={config.get('comment', '')}\n")
            handle.write(f"Department={config.get('department', '')}\n")
            handle.write(f"Pool={'' if config['pool'] == 'none' else config['pool']}\n")
            handle.write(f"Group={'' if config['group'] == 'none' else config['group']}\n")
            handle.write(f"Priority={int(config['priority'])}\n")
            handle.write(f"Frames=0-{batch_count - 1}\n")
            handle.write(f"ChunkSize={chunk_size}\n")
            handle.write(f"OutputDirectory0={output_dir}\n")
            if config.get("spool_package"):
                package = config["spool_package"]
                handle.write("ExtraInfo0=spool:" + package["submissionId"] + "\n")
                handle.write("ExtraInfoKeyValue0=SpoolSubmissionId=" + package["submissionId"] + "\n")
                handle.write("ExtraInfoKeyValue1=SpoolAttemptId=" + package["attemptId"] + "\n")
            machine_list = config.get("machine_list", "")
            if machine_list:
                if not re.fullmatch(r"[A-Za-z0-9_,.\-]+", machine_list):
                    raise ValueError("Machine list must contain comma-separated Worker names.")
                handle.write(f"Whitelist={machine_list}\n")

    def _write_plugin_info(self, path: str) -> None:
        config = self.job_config
        entries = {
            "SpoolRequestFile": "spool_request.json" if config.get("spool_package") else "",
            "StandardWorkflowFile": "workflow.json" if config.get("standard_workflow") else "",
            "JobOutputDirectory": os.path.abspath(config["output_directory"]),
            "JobInputDirectory": os.path.abspath(config["input_directory"]),
            "InputManifestFile": os.path.abspath(config["input_manifest"]),
            "SubmissionId": config["submission_id"],
            "BatchCount": str(int(config["batch_count"])),
            "BatchMode": "True",
            "DefaultCudaDeviceZero": "True",
            "SeedMode": "fixed",
            "WorkerMode": "False",
            "DistributedMode": "False",
            "ForceNewInstance": "True",
            "DisableDynamicVRAM": str(bool(config.get("disable_dynamic_vram", False))),
            "ReserveVRAM": str(max(0.0, min(64.0, float(config.get("reserve_vram", 0.0))))),
        }

        with open(path, "w", encoding="utf-8") as handle:
            for key, value in entries.items():
                if value != "":
                    handle.write(f"{key}={_deadline_value(value)}\n")


class DeadlineSeed:
    """
    Deadline-compatible seed node.
    Batch mode varies by Deadline task ID; distributed mode varies by worker ID.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "seed": ("INT", {
                    "default": 1125899906842,
                    "min": 0,
                    "max": 1125899906842624,
                    "forceInput": False,
                }),
            },
            "hidden": {
                "task_id": ("INT", {"default": 0}),
                "batch_mode": ("BOOLEAN", {"default": False}),
                "is_worker": ("BOOLEAN", {"default": False}),
                "worker_id": ("STRING", {"default": ""}),
            },
        }

    RETURN_TYPES = ("INT",)
    RETURN_NAMES = ("seed",)
    FUNCTION = "distribute"
    CATEGORY = "deadline"

    def distribute(self, seed, task_id=0, batch_mode=False, is_worker=False, worker_id=""):
        seed = int(seed)

        if _coerce_bool(batch_mode):
            try:
                task_id = int(task_id)
            except (TypeError, ValueError):
                task_id = 0
            return (seed + task_id,)

        if _coerce_bool(is_worker):
            try:
                worker_id = str(worker_id)
                if worker_id.startswith("worker_"):
                    worker_index = int(worker_id.split("_")[1])
                else:
                    worker_index = int(worker_id)
                return (seed + worker_index + 1,)
            except (TypeError, ValueError, IndexError):
                return (seed,)

        return (seed,)


class LegacyDeadlineSeedAlias(DeadlineSeed):
    DEPRECATED = True


class DeadlineSubmitNode:
    @classmethod
    def INPUT_TYPES(cls):
        pools = cls._get_deadline_pools()
        groups = cls._get_deadline_groups()
        return {
            "required": {
                "output_directory": ("STRING", {
                    "default": "",
                    "multiline": False,
                    "placeholder": "Full farm-visible output path",
                    "tooltip": "Enter a full output path, e.g. X:\\outputs\\shot010. It must be inside the admin's Allowed Storage Locations in Deadline.",
                }),
                "batch_count": ("INT", {
                    "default": NodeDefaults.BATCH_COUNT,
                    "min": 1,
                    "max": NodeDefaults.MAX_BATCH_COUNT,
                    "step": 1,
                }),
                "chunk_size": ("INT", {
                    "default": NodeDefaults.CHUNK_SIZE,
                    "min": 1,
                    "max": NodeDefaults.MAX_CHUNK_SIZE,
                    "step": 1,
                }),
                "priority": ("INT", {
                    "default": NodeDefaults.PRIORITY,
                    "min": 0,
                    "max": NodeDefaults.MAX_PRIORITY,
                }),
                "pool": (pools, {"default": NodeDefaults.POOL}),
                "group": (groups, {"default": NodeDefaults.GROUP}),
                "job_name": ("STRING", {"default": NodeDefaults.JOB_NAME}),
            },
            "optional": {
                "comment": ("STRING", {"default": ""}),
                "department": ("STRING", {"default": ""}),
                "machine_list": ("STRING", {"default": "", "tooltip": "Optional comma-separated Deadline Worker allowlist."}),
                "disable_dynamic_vram": ("BOOLEAN", {"default": False}),
                "reserve_vram": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 64.0, "step": 0.5, "tooltip": "GB reserved for activations and other GPU users."}),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("job_id",)
    FUNCTION = "submit_to_deadline"
    CATEGORY = "deadline"
    OUTPUT_NODE = True

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return f"deadline_submit_{time.time()}_{uuid.uuid4()}"

    @classmethod
    def _get_deadline_pools(cls) -> List[str]:
        try:
            output = DeadlineCommandHelper.call_deadline_command(["-pools"])
            pools = [line.strip() for line in output.splitlines() if line.strip()]
            return pools or [NodeDefaults.POOL]
        except Exception as exc:
            print(f"Deadline Submission: Could not query Deadline pools: {exc}")
            return [NodeDefaults.POOL]

    @classmethod
    def _get_deadline_groups(cls) -> List[str]:
        try:
            output = DeadlineCommandHelper.call_deadline_command(["-groups"])
            groups = [line.strip() for line in output.splitlines() if line.strip()]
            return groups or [NodeDefaults.GROUP]
        except Exception as exc:
            print(f"Deadline Submission: Could not query Deadline groups: {exc}")
            return [NodeDefaults.GROUP]

    def submit_to_deadline(
        self,
        output_directory,
        batch_count,
        chunk_size,
        priority,
        pool,
        group,
        job_name,
        comment="",
        department="",
        machine_list="",
        disable_dynamic_vram=False,
        reserve_vram=0.0,
        prompt=None,
        extra_pnginfo=None,
        **_legacy_inputs,
    ):
        try:
            output_directory = self._prepare_output_directory(output_directory)
            batch_count = max(1, int(batch_count))
            chunk_size = max(1, min(int(chunk_size), batch_count))
            spool_package = copy.deepcopy((extra_pnginfo or {}).get("spoolSubmission"))
            if spool_package and (batch_count != 1 or chunk_size != 1):
                raise ValueError("Spool jobs require one assigned take per job")
            submission_id = spool_package["submissionId"] if spool_package else uuid.uuid4().hex[:12]

            if prompt is None:
                raise ValueError("ComfyUI did not inject the current API prompt.")

            worker_prompt = WorkflowProcessor.prepare_for_worker(prompt)
            stager = InputAssetStager(output_directory, job_name, submission_id)
            input_dir, manifest_file, assets = stager.stage_referenced_assets(worker_prompt)
            self._rewrite_prompt_asset_references(worker_prompt, assets)
            standard_workflow = self._extract_standard_workflow(extra_pnginfo)
            self._rewrite_standard_workflow_assets(standard_workflow, assets)

            if spool_package:
                spool_package["stagedPrompt"] = copy.deepcopy(worker_prompt)
                spool_package["resultRoot"] = output_directory
                for asset in spool_package["assets"]:
                    value = worker_prompt[asset["nodeId"]]["inputs"][asset["input"]]
                    staged = _approved_path(os.path.join(input_dir, value), [input_dir])
                    with open(staged, "rb") as handle:
                        digest = hashlib.sha256(handle.read()).hexdigest()
                    if digest != asset["sha256"] or os.path.getsize(staged) != asset["byteSize"]:
                        raise ValueError("Stager changed a frozen Spool dependency")
                    asset["staged"] = value
                from .spool_protocol import atomic_json
                package_dir = _approved_path(os.path.join(output_directory, "spool-submissions", submission_id), [output_directory])
                _approved_path(os.path.join(package_dir, "request.json"), [output_directory])
                atomic_json(os.path.join(package_dir, "request.json"), spool_package)
            job_config = {
                "spool_package": spool_package,
                "submission_id": submission_id,
                "job_name": job_name.strip() or NodeDefaults.JOB_NAME,
                "priority": int(priority),
                "pool": pool,
                "group": group,
                "batch_count": batch_count,
                "chunk_size": chunk_size,
                "output_directory": output_directory,
                "input_directory": input_dir,
                "input_manifest": manifest_file,
                "standard_workflow": standard_workflow,
                "comment": comment,
                "department": department,
                "machine_list": ",".join(part.strip() for part in machine_list.split(",") if part.strip()),
                "disable_dynamic_vram": disable_dynamic_vram,
                "reserve_vram": reserve_vram,
            }

            submitter = DeadlineJobSubmitter(worker_prompt, job_config)
            success, result = submitter.submit_job()
            if not success:
                raise RuntimeError(result)

            if spool_package:
                _approved_path(os.path.join(package_dir, "ack.json"), [output_directory])
                atomic_json(os.path.join(package_dir, "ack.json"), {"schemaVersion": 2, "submissionId": submission_id, "attemptId": spool_package["attemptId"], "fingerprint": spool_package["fingerprint"], "jobId": result})
            print(f"Deadline Submission: Submitted job {result} with {batch_count} variation(s), chunk size {chunk_size}, {len(assets)} staged asset(s).")
            return (result,)
        except Exception as exc:
            raise

    def _rewrite_prompt_asset_references(self, prompt: Dict[str, Any], assets: List[Dict[str, Any]]) -> None:
        staged_by_original = self._asset_map(assets, "staged_relative_path", allow_absolute=True)
        staged_absolute_by_original = self._asset_map(assets, "destination", allow_absolute=True)
        if not staged_by_original:
            return

        for node in prompt.values():
            if not isinstance(node, dict):
                continue
            inputs = node.get("inputs")
            if not isinstance(inputs, dict):
                continue

            class_type = node.get("class_type", "")
            for field_name in ABSOLUTE_FILE_LOADER_FIELDS.get(class_type, ()):
                value = inputs.get(field_name)
                normalized = self._normalize_asset_reference(value, allow_absolute=True)
                if normalized in staged_absolute_by_original:
                    inputs[field_name] = staged_absolute_by_original[normalized]

            for field_name in INPUT_LOADER_FIELDS.get(class_type, ()):
                value = inputs.get(field_name)
                normalized = self._normalize_asset_reference(value)
                if normalized in staged_by_original:
                    inputs[field_name] = staged_by_original[normalized]

            for field_name, value in list(inputs.items()):
                normalized = self._normalize_asset_reference(value)
                if normalized in staged_by_original:
                    inputs[field_name] = staged_by_original[normalized]

    def _rewrite_standard_workflow_assets(self, workflow: Optional[Dict[str, Any]], assets: List[Dict[str, Any]]) -> None:
        staged_absolute_by_original = self._asset_map(assets, "destination", allow_absolute=True)
        if not workflow or not staged_absolute_by_original:
            return

        nodes = workflow.get("nodes", [])
        if not isinstance(nodes, list):
            return

        for node in nodes:
            if not isinstance(node, dict):
                continue
            widgets = node.get("widgets_values")
            if not isinstance(widgets, list):
                continue

            for index, value in enumerate(widgets):
                normalized = self._normalize_asset_reference(value)
                if normalized in staged_absolute_by_original:
                    widgets[index] = staged_absolute_by_original[normalized]

    def _asset_map(
        self,
        assets: List[Dict[str, Any]],
        target_key: str,
        allow_absolute: bool = False,
    ) -> Dict[str, str]:
        mapping = {}
        for asset in assets:
            original = self._normalize_asset_reference(
                asset.get("original_relative_path"), allow_absolute=allow_absolute
            )
            target = asset.get(target_key)
            if original and target:
                mapping[original] = target
        return mapping

    def _normalize_asset_reference(self, value: Any, allow_absolute: bool = False) -> Optional[str]:
        if not isinstance(value, str) or not value:
            return None
        cleaned = re.sub(r"\s+\[(input|output|temp)\]\s*$", "", value.strip())
        if os.path.isabs(cleaned) and not allow_absolute:
            return None
        return os.path.normpath(cleaned.replace("/", os.sep)).replace("\\", "/")

    def _extract_standard_workflow(self, extra_pnginfo: Any) -> Optional[Dict[str, Any]]:
        if not isinstance(extra_pnginfo, dict):
            return None

        workflow = extra_pnginfo.get("workflow")
        if isinstance(workflow, dict):
            return copy.deepcopy(workflow)
        if isinstance(workflow, str):
            try:
                parsed = json.loads(workflow)
                return parsed if isinstance(parsed, dict) else None
            except json.JSONDecodeError:
                return None
        return None

    def _prepare_output_directory(self, output_directory: str) -> str:
        output_directory = (output_directory or "").strip()
        if not output_directory or not os.path.isabs(output_directory):
            raise ValueError("Set output_directory to a full farm-visible path, e.g. X:\\outputs\\shot010.")
        _validate_path_text(output_directory)
        import folder_paths
        settings = DeadlineCommandHelper.get_farm_settings()
        local_output = _farm_path(folder_paths.get_output_directory(), settings)
        output_directory = _approved_path(_farm_path(output_directory, settings), [local_output] + settings["output_roots"])
        os.makedirs(output_directory, exist_ok=True)
        if not os.path.isdir(output_directory):
            raise ValueError(f"Output path is not a directory: {output_directory}")
        return output_directory


def on_prompt(json_data: Dict[str, Any]) -> Dict[str, Any]:
    prompt = json_data.get("prompt")
    if not isinstance(prompt, dict):
        return json_data

    submit_ids = [
        str(node_id)
        for node_id, node in prompt.items()
        if isinstance(node, dict) and node.get("class_type") in DEADLINE_SUBMIT_NODE_TYPES
    ]

    if submit_ids:
        json_data["partial_execution_targets"] = submit_ids[:1]
        if len(submit_ids) > 1:
            print(f"Deadline Submission: Multiple submit nodes detected; only node {submit_ids[0]} will execute locally.")
    return json_data


def register_on_prompt_handler() -> None:
    try:
        import server
        instance = getattr(getattr(server, "PromptServer", None), "instance", None)
        if instance and hasattr(instance, "add_on_prompt_handler"):
            instance.add_on_prompt_handler(on_prompt)
            print("Deadline Submission: Registered submit-only prompt handler.")
    except Exception as exc:
        print(f"Deadline Submission: Could not register prompt handler: {exc}")

    try:
        from aiohttp import web
        import folder_paths
        import server

        instance = getattr(getattr(server, "PromptServer", None), "instance", None)
        if instance and hasattr(instance, "routes"):
            @instance.routes.get("/deadline/session")
            async def deadline_session(_request):
                return web.json_response({
                    "product": "ComfyUI-Deadline-Plugin",
                    "protocol": 1,
                    "session_id": DEADLINE_ENDPOINT_SESSION_ID,
                    "pid": os.getpid(),
                    "comfyui_root": os.path.abspath(folder_paths.base_path),
                })
            print(f"Deadline Submission: Registered endpoint identity {DEADLINE_ENDPOINT_SESSION_ID}.")
    except Exception as exc:
        print(f"Deadline Submission: Could not register endpoint identity route: {exc}")


NODE_CLASS_MAPPINGS = {
    "DeadlineSubmit": DeadlineSubmitNode,
    "DeadlineSeed": DeadlineSeed,
    "DeadlineDistributedSeed": LegacyDeadlineSeedAlias,
    "DistributedSeed": LegacyDeadlineSeedAlias,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DeadlineSubmit": "Submit to Deadline",
    "DeadlineSeed": "Deadline Seed",
    "DeadlineDistributedSeed": "Deadline Seed",
    "DistributedSeed": "Deadline Seed",
}