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

from common.config import settings
from common.logging import logger


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,
    extra_args: dict[str, str] | None = None,
) -> list[str]:
    # cmd[0] is the Spark CLI binary name, configurable via
    # SPARK_EXECUTOR_SPARK_SUBMIT_BIN. Defaults to 'spark-submit' but
    # can be 'spark2-submit' (mixed-version hosts), 'pyspark' (PySpark
    # entrypoint), or a path to a wrapper script.
    cmd = [
        settings.spark_submit_bin,
        "--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}"])
    for key, value in (extra_args or {}).items():
        cmd.extend([f"--{key}", str(value)])
    cmd.append(script_path)
    logger.debug(f"build_spark_submit_command -> {cmd}")
    return cmd


def run_spark_submit(cmd: list[str]) -> "subprocess.CompletedProcess[str]":
    """Invoke `spark-submit` and return the CompletedProcess.

    On non-zero exit, raises SparkSubmitError with the failed result
    attached as `exc.result` so callers that want to salvage information
    from stderr (notably the YARN application_id, which is often present
    even when spark-submit itself failed) can do so.
    """
    logger.debug(f"run_spark_submit exec: {cmd}")
    result = subprocess.run(cmd, capture_output=True, text=True, errors="replace")
    logger.debug(
        f"run_spark_submit done rc={result.returncode} "
        f"stdout_len={len(result.stdout)} stderr_len={len(result.stderr)}"
    )
    if result.returncode != 0:
        logger.error(f"spark-submit failed (rc={result.returncode}): {result.stderr[:500]}")
        err = SparkSubmitError(
            f"spark-submit failed (rc={result.returncode}): {result.stderr}"
        )
        # Attach the failed result so callers can still parse stderr for
        # the YARN application_id and surface it in the pending record —
        # see confirm_submit_job's failure branch.
        err.result = result
        raise err
    return result