# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
import subprocess


class SparkSubmitError(Exception):
    """Raised when spark-submit exits with a non-zero return code."""


def build_spark_submit_command(
    *,
    master: str,
    deploy_mode: str,
    script_path: str,
    queue: str,
    executor_memory: str,
    executor_cores: int,
    num_executors: int,
    spark_conf: dict[str, str] | None = None,
) -> list[str]:
    cmd = [
        "spark-submit",
        "--master", master,
        "--deploy-mode", deploy_mode,
        "--queue", queue,
        "--executor-memory", executor_memory,
        "--executor-cores", str(executor_cores),
        "--num-executors", str(num_executors),
    ]
    for key, value in (spark_conf or {}).items():
        cmd.extend(["--conf", f"{key}={value}"])
    cmd.append(script_path)
    return cmd


def run_spark_submit(cmd: list[str]) -> "subprocess.CompletedProcess[str]":
    result = subprocess.run(cmd, capture_output=True, text=True, errors="replace")
    if result.returncode != 0:
        raise SparkSubmitError(
            f"spark-submit failed (rc={result.returncode}): {result.stderr}"
        )
    return result