# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import YarnClientConfig, kill_application
from spark_executor.tools.connections import store as conn_store
from spark_executor.tools.logs import _unknown_job_error

store = JobStore()


def kill_job(job_id: str) -> dict[str, str]:
    """Kill a running job.

    `job_id` accepts either the local job_id (returned by
    confirm_submit_job) or the YARN application_id.
    """
    logger.debug(f"kill_job enter job_id={job_id}")
    job = store.get_either(job_id)
    if job is None:
        raise _unknown_job_error(job_id)
    conn = conn_store.get(job.connection)
    if conn is None:
        raise KeyError(f"Connection not found: {job.connection}")
    config = YarnClientConfig.from_connection(conn)
    kill_application(job.application_id, config)
    logger.info(f"kill_job ok job_id={job.job_id} application_id={job.application_id}")
    return {
        "job_id": job.job_id,
        "application_id": job.application_id,
        "status": "KILLED",
    }