from typing import List
from typing import Optional
from typing import Dict
from typing import Any


def create_job_override_command(
    command: List[str],
    command_worker: Optional[List[str]] = None,
    num_nodes: Optional[int] = None,
) -> Dict[str, Any]:
    """Format the command into the form needed for the AWS Batch `submit_job` API.

    command: A list of tokens comprising the command
    command_worker: List of command tokens for worker nodes, if present.
    num_nodes: If `None`, this is a single-node job. If not `None`, a multi-node job.

    Returns a dictionary to be passed to `submit_job` as kwargs."""
    batch_job_args: Dict[str, Any] = {}

    if num_nodes is None:
        # Single-node jobs override the container.
        batch_job_args["containerOverrides"] = {"command": command}
    else:

        # Make a shallow copy so we can suppress output on these nodes
        node_overrides = [
            {
                "targetNodes": "0",
                "containerOverrides": {"command": command},
            },
            {"targetNodes": "1:", "containerOverrides": {"command": command_worker}},
        ]

        batch_job_args["nodeOverrides"] = {"nodePropertyOverrides": node_overrides}

    return batch_job_args