# 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, get_application_logs
from spark_executor.tools.connections import store as conn_store

store = JobStore()


def _unknown_job_error(uid: str) -> KeyError:
    """Standard "we tried both IDs and found nothing" message.

    The agent gets this from confirm_submit_job's response:
        {"job_id": "a1b2c3d4e5f6", "application_id": "application_...", ...}
    and routinely confuses which to pass here. Spelling out that BOTH
    IDs were tried (and what they look like) saves a round trip.
    """
    return KeyError(
        f"No Job found for id={uid!r} (neither as job_id nor as "
        f"application_id). Pass the job_id from confirm_submit_job's "
        f"response — it is a 12-char hex like 'a1b2c3d4e5f6'. The "
        f"application_id is the YARN ID, e.g. 'application_17400000001_0001'."
    )


def get_job_logs(job_id: str, tail_chars: int = 5000) -> str:
    """Fetch aggregated container logs for a Spark job.

    `job_id` accepts either the local job_id (returned by
    confirm_submit_job) or the YARN application_id — both are looked up
    against the same Job record.
    """
    logger.debug(f"get_job_logs enter job_id={job_id} tail_chars={tail_chars}")
    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)
    full = get_application_logs(job.application_id, config)
    tailed = full[-tail_chars:] if len(full) > tail_chars else full
    logger.info(
        f"get_job_logs ok job_id={job.job_id} application_id={job.application_id} "
        f"full_chars={len(full)} returned_chars={len(tailed)}"
    )
    return tailed