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

from common.logging import logger
from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import YarnClientConfig, get_application_status
from spark_executor.tools.connections import store as conn_store
from spark_executor.models import JobResult
from spark_executor.tools.logs import _unknown_job_error

store = JobStore()


def get_job_result(job_id: str) -> JobResult:
    """Query YARN for a job's terminal result view.

    `job_id` accepts either the local job_id (returned by
    confirm_submit_job) or the YARN application_id.
    """
    logger.debug(f"get_job_result enter job_id={job_id}")
    job = store.get_either(job_id)
    if job is None:
        raise _unknown_job_error(
            job_id,
            external_tool_hint="get_external_job_result(application_id, connection_name)",
        )
    conn = conn_store.get(job.connection)
    if conn is None:
        raise KeyError(f"Connection not found: {job.connection}")
    config = YarnClientConfig.from_connection(conn)
    state, raw = get_application_status(job.application_id, config)
    app = json.loads(raw).get("app", {})
    result = JobResult(
        application_id=job.application_id,
        state=state,
        final_status=app.get("finalStatus"),
        diagnostics=app.get("diagnostics"),
        tracking_url=app.get("trackingUrl"),
        started_time=app.get("startedTime"),
        finished_time=app.get("finishedTime"),
    )
    logger.info(
        f"get_job_result ok job_id={job.job_id} application_id={job.application_id} "
        f"state={state} final_status={result.final_status}"
    )
    return result